From bc45f37b69253a056e60e8404b539ae3c479088c Mon Sep 17 00:00:00 2001
From: Shx Guo <int123c@gmail.com>
Date: Thu, 20 Jul 2023 15:50:27 +0800
Subject: [PATCH] Adjust suggestion accepting

---
 .../SuggestionInjector.swift                  | 53 ++++++++++++-------
 .../AcceptSuggestionTests.swift               | 37 +++++++++++++
 2 files changed, 71 insertions(+), 19 deletions(-)

diff --git a/Core/Sources/SuggestionInjector/SuggestionInjector.swift b/Core/Sources/SuggestionInjector/SuggestionInjector.swift
index 16216934..8558d16d 100644
--- a/Core/Sources/SuggestionInjector/SuggestionInjector.swift
+++ b/Core/Sources/SuggestionInjector/SuggestionInjector.swift
@@ -175,36 +175,52 @@ public struct SuggestionInjector {
         }
 
         // appending suffix text not in range if needed.
-        let skipAppendingDueToContinueTyping = {
+        let leftoverCount: Int = {
+            let maxCount = lastRemovedLine?.count ?? 0
             guard let first = toBeInserted.first?
                 .dropLast((toBeInserted.first?.hasSuffix("\n") ?? false) ? 1 : 0),
-                !first.isEmpty else { return false }
+                  !first.isEmpty else { return maxCount }
             guard let last = toBeInserted.last?
                 .dropLast((toBeInserted.last?.hasSuffix("\n") ?? false) ? 1 : 0),
-                !last.isEmpty else { return false }
+                !last.isEmpty else { return maxCount }
             let droppedLast = lastRemovedLine?
                 .dropLast((lastRemovedLine?.hasSuffix("\n") ?? false) ? 1 : 0)
-            guard let droppedLast, !droppedLast.isEmpty else { return false }
+            guard let droppedLast, !droppedLast.isEmpty else { return maxCount }
             // case 1: user keeps typing as the suggestion suggests.
             if first.hasPrefix(droppedLast) {
-                return true
+                return 0
             }
             // case 2: user also typed the suffix of the suggestion (or auto-completed by Xcode)
-            if cursorPosition.character < droppedLast.count {
-                let splitIndex = droppedLast.index(
-                    droppedLast.startIndex,
-                    offsetBy: cursorPosition.character
-                )
-                let prefix = droppedLast[..<splitIndex]
-                let suffix = droppedLast[splitIndex...]
-                if first.hasPrefix(prefix), last.hasSuffix(suffix) {
-                    return true
+            if end.character < droppedLast.count {
+                // locate the split index, the prefix of which matches the suggestion prefix.
+                var splitIndex: String.Index? = nil
+                for offset in end.character..<droppedLast.count {
+                    let proposedIndex = droppedLast.index(
+                        droppedLast.startIndex,
+                        offsetBy: offset
+                    )
+                    let prefix = droppedLast[..<proposedIndex]
+                    if first.hasPrefix(prefix) {
+                        splitIndex = proposedIndex
+                    }
+                }
+                
+                // then check how many characters are not in the suffix of the suggestion.
+                if let splitIndex {
+                    let suffix = droppedLast[splitIndex...]
+                    if last.hasSuffix(suffix) {
+                        return 0
+                    }
+                    return suffix.count
                 }
             }
-            return false
+            
+            return maxCount
         }()
+        
+        // appending suffix text not in range if needed.
         if let lastRemovedLine,
-           !skipAppendingDueToContinueTyping,
+           leftoverCount > 0,
            !lastRemovedLine.isEmptyOrNewLine,
            end.character >= 0,
            end.character - 1 < lastRemovedLine.count,
@@ -218,10 +234,9 @@ public struct SuggestionInjector {
             if toBeInserted[toBeInserted.endIndex - 1].hasSuffix("\n") {
                 toBeInserted[toBeInserted.endIndex - 1].removeLast(1)
             }
-            let leftover = lastRemovedLine[leftoverRange]
+            let leftover = lastRemovedLine[leftoverRange].suffix(leftoverCount)
 
-            toBeInserted[toBeInserted.endIndex - 1]
-                .append(contentsOf: leftover)
+            toBeInserted[toBeInserted.endIndex - 1].append(contentsOf: leftover)
         }
 
         let cursorCol = toBeInserted[toBeInserted.endIndex - 1].count - 1
diff --git a/Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift b/Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift
index c8258c66..c99ba417 100644
--- a/Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift
+++ b/Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift
@@ -213,6 +213,43 @@ final class AcceptSuggestionTests: XCTestCase {
         """)
     }
     
+    func test_accept_suggestion_overlap_continue_typing_suggestion_in_the_middle() async throws {
+        let content = """
+        print("He")
+        """
+        let text = """
+        print("Hello World!
+        """
+        let suggestion = CodeSuggestion(
+            text: text,
+            position: .init(line: 0, character: 6),
+            uuid: "",
+            range: .init(
+                start: .init(line: 0, character: 0),
+                end: .init(line: 0, character: 6)
+            ),
+            displayText: ""
+        )
+
+        var extraInfo = SuggestionInjector.ExtraInfo()
+        var lines = content.breakLines()
+        var cursor = CursorPosition(line: 0, character: 7)
+        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: 0, character: 20))
+        XCTAssertEqual(lines.joined(separator: ""), """
+        print("Hello World!")
+        """)
+    }
+    
     func test_accept_suggestion_overlap_continue_typing_has_suffix_typed_suggestion_has_multiple_lines() async throws {
         let content = """
         struct Cat {}