Update UpdatedContent to return a selection range

Author: intitni

Core/Sources/CopilotModel/ExportedFromLSP.swift

@@ -9,4 +9,7 @@ public extension CursorPosition {
 
 public extension CursorRange {
     static var outOfScope: CursorRange { .init(start: .outOfScope, end: .outOfScope) }
+    static func cursor(_ position: CursorPosition) -> CursorRange {
+        return .init(start: position, end: position)
+    }
 }

Core/Sources/Service/SuggestionCommandHandler/CommentBaseCommandHandler.swift

@@ -103,7 +103,7 @@ struct CommentBaseCommandHandler: SuggestionCommandHandler {
 
         return .init(
             content: String(lines.joined(separator: "")),
-            newCursor: cursorPosition,
+            newSelection: .cursor(cursorPosition),
             modifications: extraInfo.modifications
         )
     }

Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift

@@ -108,7 +108,8 @@ struct PseudoCommandHandler {
                         .presentErrorMessage("Fail to set editor content.")
                 }
 
-                if let cursor = result.newCursor {
+                #warning("MUSTDO: handle range")
+                if let cursor = result.newSelection?.start {
                     var range = convertCursorPositionToRange(cursor, in: result.content)
                     if let value = AXValueCreate(.cfRange, &range) {
                         AXUIElementSetAttributeValue(

Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift

@@ -136,47 +136,55 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
     func acceptSuggestion(editor: EditorContent) async throws -> UpdatedContent? {
         presenter.markAsProcessing(true)
         defer { presenter.markAsProcessing(false) }
-        
+
         let fileURL = try await Environment.fetchCurrentFileURL()
         let (workspace, _) = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
 
-        let result: (suggestion: CopilotCompletion, cleanup: () -> Void)? = {
+        let result: (
+            suggestion: CopilotCompletion,
+            cleanup: () -> Void,
+            startPosition: CursorPosition?
+        )? = {
             if let service = WidgetDataSource.shared.promptToCodes[fileURL]?.promptToCodeService {
-                return (CopilotCompletion(
-                    text: service.code,
-                    position: service.selectionRange.start,
-                    uuid: UUID().uuidString,
-                    range: service.selectionRange,
-                    displayText: service.code
-                ), {
-                    WidgetDataSource.shared.removePromptToCode(for: fileURL)
-                    presenter.closePromptToCode(fileURL: fileURL)
-                })
+                return (
+                    CopilotCompletion(
+                        text: service.code,
+                        position: service.selectionRange.start,
+                        uuid: UUID().uuidString,
+                        range: service.selectionRange,
+                        displayText: service.code
+                    ),
+                    {
+                        WidgetDataSource.shared.removePromptToCode(for: fileURL)
+                        presenter.closePromptToCode(fileURL: fileURL)
+                    },
+                    service.selectionRange.start
+                )
             }
-            
+
             if let acceptedSuggestion = workspace.acceptSuggestion(
                 forFileAt: fileURL,
                 editor: editor
             ) {
-                return (acceptedSuggestion, {
-                    presenter.discardSuggestion(fileURL: fileURL)
-                })
+                return (
+                    acceptedSuggestion,
+                    {
+                        presenter.discardSuggestion(fileURL: fileURL)
+                    },
+                    nil
+                )
             }
-            
+
             return nil
         }()
-        
+
         guard let result else { return nil }
-        
+
         let injector = SuggestionInjector()
         var lines = editor.lines
         var cursorPosition = editor.cursorPosition
         var extraInfo = SuggestionInjector.ExtraInfo()
-        injector.rejectCurrentSuggestions(
-            from: &lines,
-            cursorPosition: &cursorPosition,
-            extraInfo: &extraInfo
-        )
+
         injector.acceptSuggestion(
             intoContentWithoutSuggestion: &lines,
             cursorPosition: &cursorPosition,
@@ -185,11 +193,16 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
         )
 
         result.cleanup()
-        
+
         return .init(
             content: String(lines.joined(separator: "")),
-            newCursor: cursorPosition,
-            modifications: extraInfo.modifications
+            newSelection: {
+                if let startPosition = result.startPosition {
+                    return .init(start: startPosition, end: cursorPosition)
+                }
+                return .cursor(cursorPosition)
+            }(),
+            modifications : extraInfo.modifications
         )
     }
 

Core/Sources/Service/SuggestionPresenter/PresentInCommentSuggestionPresenter.swift

@@ -25,7 +25,7 @@ struct PresentInCommentSuggestionPresenter {
         guard let completion = await filespace.presentingSuggestion else {
             return .init(
                 content: originalContent,
-                newCursor: cursorPosition,
+                newSelection: .cursor(cursorPosition),
                 modifications: extraInfo.modifications
             )
         }
@@ -40,7 +40,7 @@ struct PresentInCommentSuggestionPresenter {
 
         return .init(
             content: String(lines.joined(separator: "")),
-            newCursor: cursorPosition,
+            newSelection: .cursor(cursorPosition),
             modifications: extraInfo.modifications
         )
     }
@@ -65,7 +65,7 @@ struct PresentInCommentSuggestionPresenter {
 
         return .init(
             content: String(lines.joined(separator: "")),
-            newCursor: cursorPosition,
+            newSelection: .cursor(cursorPosition),
             modifications: extraInfo.modifications
         )
     }

Core/Sources/XPCShared/Models.swift

@@ -47,14 +47,14 @@ public struct EditorContent: Codable {
 }
 
 public struct UpdatedContent: Codable {
-    public init(content: String, newCursor: CursorPosition? = nil, modifications: [Modification]) {
+    public init(content: String, newSelection: CursorRange? = nil, modifications: [Modification]) {
         self.content = content
-        self.newCursor = newCursor
+        self.newSelection = newSelection
         self.modifications = modifications
     }
 
     public var content: String
-    public var newCursor: CursorPosition?
+    public var newSelection: CursorRange?
     public var modifications: [Modification]
 }
 

EditorExtension/Helpers.swift

@@ -19,15 +19,15 @@ extension XCSourceEditorCommandInvocation {
     }
 
     func accept(_ updatedContent: UpdatedContent) {
-        if let newCursor = updatedContent.newCursor {
+        if let newSelection = updatedContent.newSelection {
             mutateCompleteBuffer(
                 modifications: updatedContent.modifications,
                 restoringSelections: false
             )
             buffer.selections.removeAllObjects()
             buffer.selections.add(XCSourceTextRange(
-                start: .init(line: newCursor.line, column: newCursor.character),
-                end: .init(line: newCursor.line, column: newCursor.character)
+                start: .init(line: newSelection.start.line, column: newSelection.start.character),
+                end: .init(line: newSelection.end.line, column: newSelection.end.character)
             ))
         } else {
             mutateCompleteBuffer(