Reuse suggestion if user is typing according to it

Author: intitni

Core/Sources/Service/RealtimeSuggestionController.swift

@@ -26,7 +26,7 @@ public class RealtimeSuggestionController {
     private var activeApplicationMonitorTask: Task<Void, Error>?
     private var editorObservationTask: Task<Void, Error>?
     private var focusedUIElement: AXUIElement?
-    private var sourceEditor: AXUIElement?
+    private var sourceEditor: SourceEditor?
 
     var isCommentMode: Bool {
         UserDefaults.shared.value(for: \.suggestionPresentationMode) == .comment
@@ -116,7 +116,7 @@ public class RealtimeSuggestionController {
         }
 
         guard focusElementType == "Source Editor" else { return }
-        sourceEditor = focusElement
+        sourceEditor = SourceEditor(runningApplication: activeXcode, element: focusElement)
 
         editorObservationTask?.cancel()
         editorObservationTask = nil
@@ -138,11 +138,7 @@ public class RealtimeSuggestionController {
                     self.triggerPrefetchDebounced()
                     await self.notifyEditingFileChange(editor: focusElement)
                 case kAXSelectedTextChangedNotification:
-                    guard let editor = sourceEditor else { continue }
-                    let sourceEditor = SourceEditor(
-                        runningApplication: activeXcode,
-                        element: editor
-                    )
+                    guard let sourceEditor else { continue }
                     await PseudoCommandHandler()
                         .invalidateRealtimeSuggestionsIfNeeded(sourceEditor: sourceEditor)
                 default:

Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift

@@ -39,9 +39,21 @@ struct PseudoCommandHandler {
         ))
     }
 
-    func generateRealtimeSuggestions(sourceEditor: AXUIElement?) async {
+    func generateRealtimeSuggestions(sourceEditor: SourceEditor?) async {
         // Can't use handler if content is not available.
-        guard let editor = await getEditorContent(sourceEditor: sourceEditor) else { return }
+        guard
+            let editor = await getEditorContent(sourceEditor: sourceEditor),
+            let filespace = await getFilespace()
+        else { return }
+
+        if await filespace.validateSuggestions(
+            lines: editor.lines,
+            cursorPosition: editor.cursorPosition
+        ) {
+            return
+        } else {
+            PresentInWindowSuggestionPresenter().discardSuggestion(fileURL: filespace.fileURL)
+        }
 
         // Otherwise, get it from pseudo handler directly.
         let mode = UserDefaults.shared.value(for: \.suggestionPresentationMode)
@@ -246,11 +258,9 @@ extension PseudoCommandHandler {
     }
 
     @ServiceActor
-    func getEditorContent(sourceEditor: AXUIElement?) async -> EditorContent? {
-        guard
-            let filespace = await getFilespace(),
-            let content = await getFileContent(sourceEditor: sourceEditor)
-        else { return nil }
+    func getEditorContent(sourceEditor: SourceEditor?) async -> EditorContent? {
+        guard let filespace = await getFilespace(), let sourceEditor else { return nil }
+        let content = sourceEditor.content
         let uti = filespace.uti ?? ""
         let tabSize = filespace.tabSize ?? 4
         let indentSize = filespace.indentSize ?? 4

Core/Sources/Service/Workspace.swift

@@ -79,6 +79,13 @@ final class Filespace {
             return false
         }
 
+        let editingLine = lines[cursorPosition.line].dropLast(1) // dropping \n
+        let suggestionFirstLine = presentingSuggestion?.text.split(separator: "\n").first ?? ""
+        if !suggestionFirstLine.hasPrefix(editingLine) {
+            reset()
+            return false
+        }
+        
         return true
     }
 }

Core/Sources/SuggestionInjector/SuggestionInjector.swift

@@ -154,6 +154,7 @@ public struct SuggestionInjector {
 
         var toBeInserted = suggestionContent.breakLines(appendLineBreakToLastLine: true)
 
+        // prepending prefix text not in range if needed.
         if let firstRemovedLine,
            !firstRemovedLine.isEmptyOrNewLine,
            start.character > 0,
@@ -175,8 +176,16 @@ public struct SuggestionInjector {
             )
         }
 
+        // appending suffix text not in range if needed.
         let cursorCol = toBeInserted[toBeInserted.endIndex - 1].count - 1
+        let skipAppendingDueToContinueTyping = {
+            guard let first = toBeInserted.first?.dropLast(1), !first.isEmpty else { return false }
+            let droppedLast = lastRemovedLine?.dropLast(1)
+            guard let droppedLast, !droppedLast.isEmpty else { return false }
+            return first.hasPrefix(droppedLast)
+        }()
         if let lastRemovedLine,
+           !skipAppendingDueToContinueTyping,
            !lastRemovedLine.isEmptyOrNewLine,
            end.character >= 0,
            end.character - 1 < lastRemovedLine.count,
@@ -191,6 +200,7 @@ public struct SuggestionInjector {
                 toBeInserted[toBeInserted.endIndex - 1].removeLast(1)
             }
             let leftover = lastRemovedLine[leftoverRange]
+            
             toBeInserted[toBeInserted.endIndex - 1]
                 .append(contentsOf: leftover)
         }

Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift

@@ -131,6 +131,49 @@ final class AcceptSuggestionTests: XCTestCase {
         }
         """)
     }
+    
+    func test_accept_suggestion_overlap_continue_typing() async throws {
+        let content = """
+        struct Cat {
+            var name: Str
+        }
+        """
+        let text = """
+            var name: String
+            var age: String
+        """
+        let suggestion = CodeSuggestion(
+            text: text,
+            position: .init(line: 1, character: 12),
+            uuid: "",
+            range: .init(
+                start: .init(line: 1, character: 0),
+                end: .init(line: 1, character: 12)
+            ),
+            displayText: ""
+        )
+
+        var extraInfo = SuggestionInjector.ExtraInfo()
+        var lines = content.breakLines()
+        var cursor = CursorPosition(line: 0, character: 0)
+        SuggestionInjector().acceptSuggestion(
+            intoContentWithoutSuggestion: &lines,
+            cursorPosition: &cursor,
+            completion: suggestion,
+            extraInfo: &extraInfo
+        )
+        XCTAssertTrue(extraInfo.didChangeContent)
+        XCTAssertTrue(extraInfo.didChangeCursorPosition)
+        XCTAssertNil(extraInfo.suggestionRange)
+        XCTAssertEqual(lines, content.breakLines().applying(extraInfo.modifications))
+        XCTAssertEqual(cursor, .init(line: 2, character: 19))
+        XCTAssertEqual(lines.joined(separator: ""), """
+        struct Cat {
+            var name: String
+            var age: String
+        }
+        """)
+    }
 
     func test_propose_suggestion_partial_overlap() async throws {
         let content = "func quickSort() {}}\n"