Reimplement suffix recovering

Author: intitni

Core/Sources/SuggestionInjector/SuggestionInjector.swift

@@ -174,97 +174,77 @@ public struct SuggestionInjector {
             )
         }
 
-        #warning(
-            "TODO: I feel like the implementation is doing a lot of receptive work to recover suffix."
+        let recoveredSuffixLength = recoverSuffixIfNeeded(
+            endOfReplacedContent: end,
+            toBeInserted: &toBeInserted,
+            lastRemovedLine: lastRemovedLine
         )
 
-        /// The number of character that is in the last modified line but not included.
-        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 maxCount }
-            guard let last = toBeInserted.last?
-                .dropLast((toBeInserted.last?.hasSuffix("\n") ?? false) ? 1 : 0),
-                !last.isEmpty else { return maxCount }
-            let droppedLast = lastRemovedLine?
-                .dropLast((lastRemovedLine?.hasSuffix("\n") ?? false) ? 1 : 0)
-            guard let droppedLast, !droppedLast.isEmpty else { return maxCount }
-            // case 1: user keeps typing as the suggestion suggests.
-            if first.hasPrefix(droppedLast) {
-                return 0
-            }
-            // case 2: user also typed the suffix of the suggestion (or auto-completed by Xcode)
-            if end.character < droppedLast.count {
-                // locate the split index, the prefix of which matches the suggestion prefix.
-                var splitIndex: String.Index?
-                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
-                    }
-                }
+        let cursorCol = toBeInserted[toBeInserted.endIndex - 1].count - 1 - recoveredSuffixLength
+        let insertingIndex = min(start.line, content.endIndex)
+        content.insert(contentsOf: toBeInserted, at: insertingIndex)
+        extraInfo.modifications.append(.inserted(insertingIndex, toBeInserted))
+        cursorPosition = .init(
+            line: startLine + toBeInserted.count - 1,
+            character: max(0, cursorCol)
+        )
+    }
 
-                // 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
-                }
-            }
+    func recoverSuffixIfNeeded(
+        endOfReplacedContent end: CursorPosition,
+        toBeInserted: inout [String],
+        lastRemovedLine: String?
+    ) -> Int {
+        // If there is no line removed, there is no need to recover anything.
+        guard let lastRemovedLine, !lastRemovedLine.isEmptyOrNewLine else { return 0 }
 
-            return maxCount
-        }()
+        let lastRemovedLineCleaned = lastRemovedLine.droppedLineBreak()
 
-        /// The number of characters appended to the last line.
-        var appenderCount = 0
+        // If the replaced range covers the whole line, return immediately.
+        guard end.character >= 0, end.character - 1 < lastRemovedLineCleaned.count else { return 0 }
 
-        // appending suffix text not in range if needed.
-        if let lastRemovedLine,
-           leftoverCount > 0,
-           !lastRemovedLine.isEmptyOrNewLine,
-           end.character >= 0,
-           end.character - 1 < lastRemovedLine.count,
-           !toBeInserted.isEmpty
-        {
-            let leftoverRange = (lastRemovedLine.index(
-                lastRemovedLine.startIndex,
-                offsetBy: end.character,
-                limitedBy: lastRemovedLine.endIndex
-            ) ?? lastRemovedLine.endIndex)..<lastRemovedLine.endIndex
-            if toBeInserted[toBeInserted.endIndex - 1].hasSuffix("\n") {
-                toBeInserted[toBeInserted.endIndex - 1].removeLast(1)
-            }
-            let leftover = {
-                if lastRemovedLine.hasSuffix("\n") {
-                    return lastRemovedLine[leftoverRange].dropLast(1).suffix(leftoverCount)
-                }
-                return lastRemovedLine[leftoverRange].suffix(leftoverCount)
-            }()
+        // if we are not inserting anything, return immediately.
+        guard !toBeInserted.isEmpty,
+              let first = toBeInserted.first?.droppedLineBreak(), !first.isEmpty,
+              let last = toBeInserted.last?.droppedLineBreak(), !last.isEmpty
+        else { return 0 }
 
-            toBeInserted[toBeInserted.endIndex - 1].append(contentsOf: leftover)
+        // case 1: user keeps typing as the suggestion suggests.
 
-            appenderCount = leftover.count
+        if first.hasPrefix(lastRemovedLineCleaned) {
+            return 0
+        }
 
-            if !toBeInserted[toBeInserted.endIndex - 1].hasSuffix("\n") {
-                toBeInserted[toBeInserted.endIndex - 1].append("\n")
+        // case 2: user also typed the suffix of the suggestion (or auto-completed by Xcode)
+
+        // locate the split index, the prefix of which matches the suggestion prefix.
+        var splitIndex: String.Index?
+
+        for offset in end.character..<lastRemovedLineCleaned.count {
+            let proposedIndex = lastRemovedLineCleaned.index(
+                lastRemovedLineCleaned.startIndex,
+                offsetBy: offset
+            )
+            let prefix = lastRemovedLineCleaned[..<proposedIndex]
+            if first.hasPrefix(prefix) {
+                splitIndex = proposedIndex
             }
         }
 
-        let cursorCol = toBeInserted[toBeInserted.endIndex - 1].count - 1 - appenderCount
-        let insertingIndex = min(start.line, content.endIndex)
-        content.insert(contentsOf: toBeInserted, at: insertingIndex)
-        extraInfo.modifications.append(.inserted(insertingIndex, toBeInserted))
-        cursorPosition = .init(
-            line: startLine + toBeInserted.count - 1,
-            character: max(0, cursorCol)
-        )
+        // then check how many characters are not in the suffix of the suggestion.
+        guard let splitIndex else { return 0 }
+        
+        let suffix = lastRemovedLineCleaned[splitIndex...]
+        if last.hasSuffix(suffix) { return 0 }
+
+        let lastInsertingLine = toBeInserted[toBeInserted.endIndex - 1]
+            .droppedLineBreak()
+            .appending(suffix)
+            .recoveredLineBreak()
+
+        toBeInserted[toBeInserted.endIndex - 1] = lastInsertingLine
+
+        return suffix.count
     }
 }
 
@@ -302,6 +282,20 @@ extension String {
     var isEmptyOrNewLine: Bool {
         isEmpty || self == "\n"
     }
+
+    func droppedLineBreak() -> String {
+        if hasSuffix("\n") {
+            return String(dropLast(1))
+        }
+        return self
+    }
+
+    func recoveredLineBreak() -> String {
+        if hasSuffix("\n") {
+            return self
+        }
+        return self + "\n"
+    }
 }
 
 func longestCommonPrefix(of a: String, and b: String) -> String {