Support sending back multiple selections to editor extension

Author: intitni

Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift

@@ -459,7 +459,7 @@ extension PseudoCommandHandler {
 
         // recover selection range
 
-        if let selection = result.newSelection {
+        if let selection = result.newSelections.first {
             var range = SourceEditor.convertCursorRangeToRange(selection, in: result.content)
             if let value = AXValueCreate(.cfRange, &range) {
                 AXUIElementSetAttributeValue(

Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift

@@ -241,13 +241,7 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
 
             return .init(
                 content: String(lines.joined(separator: "")),
-                newSelection: {
-                    if ranges.isEmpty {
-                        .init(start: cursorPosition, end: cursorPosition)
-                    } else {
-                        ranges.last
-                    }
-                }(),
+                newSelections: ranges,
                 modifications: extraInfo.modifications
             )
         }

EditorExtension/Helpers.swift

@@ -1,5 +1,5 @@
-import SuggestionBasic
 import Foundation
+import SuggestionBasic
 import XcodeKit
 import XPCShared
 
@@ -19,16 +19,21 @@ extension XCSourceEditorCommandInvocation {
     }
 
     func accept(_ updatedContent: UpdatedContent) {
-        if let newSelection = updatedContent.newSelection {
+        if !updatedContent.newSelections.isEmpty {
             mutateCompleteBuffer(
                 modifications: updatedContent.modifications,
                 restoringSelections: false
             )
             buffer.selections.removeAllObjects()
-            buffer.selections.add(XCSourceTextRange(
-                start: .init(line: newSelection.start.line, column: newSelection.start.character),
-                end: .init(line: newSelection.end.line, column: newSelection.end.character)
-            ))
+            for newSelection in updatedContent.newSelections {
+                buffer.selections.add(XCSourceTextRange(
+                    start: .init(
+                        line: newSelection.start.line,
+                        column: newSelection.start.character
+                    ),
+                    end: .init(line: newSelection.end.line, column: newSelection.end.character)
+                ))
+            }
         } else {
             mutateCompleteBuffer(
                 modifications: updatedContent.modifications,
@@ -47,17 +52,17 @@ extension EditorContent {
             uti: buffer.contentUTI,
             cursorPosition: ((buffer.selections.lastObject as? XCSourceTextRange)?.end).map {
                 CursorPosition(line: $0.line, character: $0.column)
-            } ?? CursorPosition(line: 0, character: 0), 
+            } ?? CursorPosition(line: 0, character: 0),
             cursorOffset: -1,
             selections: buffer.selections.map {
                 let sl = ($0 as? XCSourceTextRange)?.start.line ?? 0
                 let sc = ($0 as? XCSourceTextRange)?.start.column ?? 0
                 let el = ($0 as? XCSourceTextRange)?.end.line ?? 0
                 let ec = ($0 as? XCSourceTextRange)?.end.column ?? 0
-                
+
                 return Selection(
-                    start: CursorPosition( line: sl, character: sc ),
-                    end: CursorPosition( line: el, character: ec )
+                    start: CursorPosition(line: sl, character: sc),
+                    end: CursorPosition(line: el, character: ec)
                 )
             },
             tabSize: buffer.tabWidth,
@@ -96,3 +101,4 @@ extension Task where Failure == Error {
 private struct TimeoutError: LocalizedError {
     var errorDescription: String? = "Task timed out before completion"
 }
+

Tool/Sources/XPCShared/Models.swift

@@ -53,12 +53,18 @@ public struct EditorContent: Codable {
 public struct UpdatedContent: Codable {
     public init(content: String, newSelection: CursorRange? = nil, modifications: [Modification]) {
         self.content = content
-        self.newSelection = newSelection
+        self.newSelections = if let newSelection { [newSelection] } else { [] }
+        self.modifications = modifications
+    }
+    
+    public init(content: String, newSelections: [CursorRange], modifications: [Modification]) {
+        self.content = content
+        self.newSelections = newSelections
         self.modifications = modifications
     }
 
     public var content: String
-    public var newSelection: CursorRange?
+    public var newSelections: [CursorRange]
     public var modifications: [Modification]
 }