Fix replace-all matching with reused full-file indexes

Sources/RepoPromptDomainRuntime/Diffing/DiffBatchGenerator.swift

@@ -43,7 +43,7 @@ package enum DiffBatchGenerator {
         outcomes.reserveCapacity(edits.count)
         allChunks.reserveCapacity(edits.count)
 
-        // Pre‑compute the high‑precision line‑index for first‑hit optimisation.
+        // Precompute matching data once so replace-all edits can keep full-file coordinates.
         var processed: [DiffGenerationUtility.LineData] = []
         processed.reserveCapacity(orig.count)
         for line in orig {
@@ -60,13 +60,14 @@ package enum DiffBatchGenerator {
 
                 let diff = try await DiffGenerationUtility.generateDiff(
                     fileContent: orig,
-                    lineIndexMap: start == 0 ? indexMap : nil, // optimisation only for first search
+                    lineIndexMap: start == 0 || edit.replaceAll ? indexMap : nil,
                     startSelector: nil,
                     endSelector: nil,
                     searchBlock: edit.search.isEmpty ? nil : edit.search,
                     newContent: sanitizedContent,
                     action: edit.search.isEmpty ? .rewrite : .modify,
                     diffPrecision: prec,
+                    processedFileContent: edit.replaceAll ? processed : nil,
                     searchStartLine: start,
                     mcpAmbiguityCheck: edit.replaceAll ? false : mcpAmbiguityCheck,
                     replaceAll: edit.replaceAll,

Sources/RepoPromptDomainRuntime/Diffing/DiffGenerationUtility.swift

@@ -78,6 +78,7 @@ package class DiffGenerationUtility {
         newContent: [String],
         action: FileAction,
         diffPrecision: DiffPrecision = .normal,
+        processedFileContent: [LineData]? = nil,
         searchStartLine: Int = 0,
         mcpAmbiguityCheck: Bool = false,
         replaceAll: Bool = false,
@@ -102,6 +103,7 @@ package class DiffGenerationUtility {
                     newContent: newContent,
                     diffPrecision: diffPrecision,
                     lineIndexMap: lineIndexMap,
+                    processedFileContent: processedFileContent,
                     searchStartLine: searchStartLine,
                     mcpAmbiguityCheck: mcpAmbiguityCheck,
                     replaceAll: replaceAll,
@@ -281,6 +283,7 @@ package class DiffGenerationUtility {
         newContent: [String],
         diffPrecision: DiffPrecision = .normal,
         lineIndexMap: [String: [Int]]? = nil,
+        processedFileContent: [LineData]? = nil,
         searchStartLine: Int = 0,
         mcpAmbiguityCheck: Bool = false,
         replaceAll: Bool = false,
@@ -308,34 +311,31 @@ package class DiffGenerationUtility {
 
         // Handle replace_all by finding multiple matches
         if replaceAll {
+            let processedFile = processedFileContent ?? fileContent.map {
+                processLine($0, precision: diffPrecision)
+            }
+            let fullIndexMap = lineIndexMap ?? buildLineIndexMapHigh(content: processedFile)
             var allChunks: [DiffChunk] = []
             allChunks.reserveCapacity(1)
             var currentStartLine = searchStartLine
 
             while currentStartLine < fileContent.count {
-                // ➊ Search only after `currentStartLine`
-                let fileSlice = fileContent[currentStartLine...]
-                let processedFileSlice = fileSlice.map { processLine($0, precision: diffPrecision) }
-
-                let sliceIndexMap = lineIndexMap ?? buildLineIndexMapHigh(content: processedFileSlice)
-
-                // ➋ Locate match *inside* the slice
-                let localMatch: Int
+                // ➊ Locate the next match using full-file coordinates.
+                let globalMatch: Int
                 do {
-                    localMatch = try await findBestMatchUsingNGrams(
+                    globalMatch = try await findBestMatchUsingNGrams(
                         selector: processedSearch,
-                        in: processedFileSlice,
-                        lineIndexMap: sliceIndexMap,
-                        mcpAmbiguityCheck: false // Skip ambiguity check for replace_all
+                        in: processedFile,
+                        lineIndexMap: fullIndexMap,
+                        mcpAmbiguityCheck: false, // Skip ambiguity check for replace_all
+                        minimumMatchIndex: currentStartLine
                     )
                 } catch DiffGenerationError.noMatchFound {
                     break // No more matches found
                 }
 
-                if localMatch == -1 { break }
+                if globalMatch == -1 { break }
 
-                // Convert to absolute indices in original file
-                let globalMatch = localMatch + currentStartLine
                 let globalEnd = globalMatch + sanitizedSearch.count
 
                 guard globalEnd <= fileContent.count else {
@@ -495,7 +495,13 @@ package class DiffGenerationUtility {
 
     /// Finds the best match for a selector in the given content using an n-gram similarity search,
     /// then refines that match in a smaller window.
-    package static func findBestMatchUsingNGrams(selector: [LineData], in content: [LineData], lineIndexMap: [String: [Int]], mcpAmbiguityCheck: Bool = false) async throws -> Int {
+    package static func findBestMatchUsingNGrams(
+        selector: [LineData],
+        in content: [LineData],
+        lineIndexMap: [String: [Int]],
+        mcpAmbiguityCheck: Bool = false,
+        minimumMatchIndex: Int = 0
+    ) async throws -> Int {
         // 1) Basic validation - ensure inputs are valid
         guard !selector.isEmpty else {
             throw DiffGenerationError.invalidSelector
@@ -512,7 +518,12 @@ package class DiffGenerationUtility {
         let quickIndex: Int? = if mcpAmbiguityCheck {
             try matchSelectorFastWithAmbiguityCheck(selector: selector, content: content, lineIndex: lineIndexMap)
         } else {
-            try matchSelectorFast(selector: selector, content: content, lineIndex: lineIndexMap)
+            try matchSelectorFast(
+                selector: selector,
+                content: content,
+                lineIndex: lineIndexMap,
+                minimumMatchIndex: minimumMatchIndex
+            )
         }
 
         if let quickIndex {
@@ -1244,7 +1255,8 @@ package class DiffGenerationUtility {
         content: [LineData],
         lineIndex: [String: [Int]],
         maxFuzzyKeys maxKeys: Int = 400,
-        fuzzyThreshold sim: Double = 0.90
+        fuzzyThreshold sim: Double = 0.90,
+        minimumMatchIndex: Int = 0
     ) throws -> Int? {
         // ── Guard rails ─────────────────────────────────────────────────────────
         guard !selector.isEmpty, !content.isEmpty else {
@@ -1274,7 +1286,9 @@ package class DiffGenerationUtility {
 
         /// ── Helper: candidate list for selector line 0 --------------------------
         func strictOrLoosePositions(for line: LineData) -> [Int] {
-            lineIndex[line.removedTagsHigh] ?? lineIndex[line.removedTags] ?? []
+            let strict = (lineIndex[line.removedTagsHigh] ?? []).filter { $0 >= minimumMatchIndex }
+            if !strict.isEmpty { return strict }
+            return (lineIndex[line.removedTags] ?? []).filter { $0 >= minimumMatchIndex }
         }
 
         var starts = strictOrLoosePositions(for: selector[0])
@@ -1294,13 +1308,15 @@ package class DiffGenerationUtility {
             var seen = 0
             for (k, pos) in lineIndex {
                 if seen >= maxKeys { break }
+                let positionsAtOrAfterMinimum = pos.filter { $0 >= minimumMatchIndex }
+                guard !positionsAtOrAfterMinimum.isEmpty else { continue }
                 seen += 1
                 let coeff = sKey.diceCoefficient(against: k)
                 if enableDetailedLogging {
                     print("  🔍 Fuzzy probe test \(seen)/\(maxKeys): \(sKey) ↔ \(k)  Dice: \(coeff)")
                 }
                 guard coeff >= fuzzyThresh else { continue }
-                for p in pos {
+                for p in positionsAtOrAfterMinimum {
                     starts.append(p)
                     fuzzyScoreMap[p] = max(fuzzyScoreMap[p] ?? 0, coeff)
                 }
@@ -1320,9 +1336,13 @@ package class DiffGenerationUtility {
                 var collected: [Int] = []
                 var scanned = 0
                 for (k, pos) in lineIndex where scanned < maxKeys {
+                    let positionsAtOrAfterMinimum = pos.filter { $0 >= minimumMatchIndex }
+                    guard !positionsAtOrAfterMinimum.isEmpty else { continue }
                     scanned += 1
                     if sKey.diceCoefficient(against: k) >= fuzzyThresh {
-                        collected += pos.compactMap { $0 > 0 ? $0 - 1 : nil }
+                        collected += positionsAtOrAfterMinimum.compactMap {
+                            $0 > minimumMatchIndex ? $0 - 1 : nil
+                        }
                     }
                 }
                 second = collected

Tests/RepoPromptTests/Diffing/DiffGenerationUtilityRoutingTests.swift

@@ -2,6 +2,129 @@
 import XCTest
 
 final class DiffGenerationUtilityRoutingTests: XCTestCase {
+    func testReplaceAllHonorsSearchStartLineWithFullFileIndex() async throws {
+        let fileContent = ["same", "skip", "same", "middle", "same"]
+        let processed = fileContent.map {
+            DiffGenerationUtility.processLine($0, precision: .high)
+        }
+        let lineIndexMap = DiffGenerationUtility.buildLineIndexMapHigh(content: processed)
+
+        let chunks = try await DiffGenerationUtility.generateDiff(
+            fileContent: fileContent,
+            lineIndexMap: lineIndexMap,
+            startSelector: nil,
+            endSelector: nil,
+            searchBlock: ["same"],
+            newContent: ["replacement"],
+            action: .modify,
+            diffPrecision: .high,
+            processedFileContent: processed,
+            searchStartLine: 2,
+            replaceAll: true
+        )
+        let result = try DiffChunkTextApplier.apply(
+            chunks: chunks,
+            to: fileContent.joined(separator: "\n")
+        )
+
+        XCTAssertEqual(chunks.map(\.startLine), [2, 4])
+        XCTAssertEqual(result, "same\nskip\nreplacement\nmiddle\nreplacement")
+    }
+
+    func testFuzzyProbeIgnoresPreBoundaryKeysBeforeSpendingBudget() throws {
+        let selectorText = "calculate total for selected invoice line items"
+        let selector = DiffGenerationUtility.processLine(selectorText, precision: .high)
+        let minimumMatchIndex = 401
+        let maxFuzzyKeys = 400
+        let prefix = (0 ..< minimumMatchIndex).map {
+            "calculate total for selected invoice line item \($0)"
+        }
+        let processedPrefix = prefix.map {
+            DiffGenerationUtility.processLine($0, precision: .high)
+        }
+        let prefixIndex = DiffGenerationUtility.buildLineIndexMapHigh(content: processedPrefix)
+        let fuzzyAliases = (0 ..< (maxFuzzyKeys * 10)).map {
+            "calculate total for selected invoice line item candidate \(String($0, radix: 36))"
+        }
+        let fuzzyCandidate = try XCTUnwrap(
+            fuzzyAliases.compactMap { candidate -> (String, [String: [Int]])? in
+                let candidateLine = DiffGenerationUtility.processLine(candidate, precision: .high)
+                let candidateKeys = DiffGenerationUtility
+                    .buildLineIndexMapHigh(content: [candidateLine])
+                    .keys
+                var candidateIndex = prefixIndex
+                for key in candidateKeys {
+                    candidateIndex[key] = [minimumMatchIndex]
+                }
+                let orderedKeys = Array(candidateIndex.keys)
+                let positions = candidateKeys.compactMap { key in
+                    orderedKeys.firstIndex(of: key)
+                }
+                guard !positions.isEmpty, positions.allSatisfy({ $0 >= 400 }) else {
+                    return nil
+                }
+                return (candidate, candidateIndex)
+            }.first,
+            "The test needs a fuzzy alias after the 400-key probe boundary"
+        )
+        let content = processedPrefix + [DiffGenerationUtility.processLine(fuzzyCandidate.0, precision: .high)]
+        let lineIndex = fuzzyCandidate.1
+
+        let result = try DiffGenerationUtility.matchSelectorFast(
+            selector: [selector],
+            content: content,
+            lineIndex: lineIndex,
+            maxFuzzyKeys: maxFuzzyKeys,
+            fuzzyThreshold: 0.80,
+            minimumMatchIndex: minimumMatchIndex
+        )
+
+        XCTAssertEqual(result, minimumMatchIndex)
+    }
+
+    func testBatchReplaceAllKeepsFullFileIndexCoordinatesAfterFirstMatch() async throws {
+        let original = [
+            "header",
+            "padding",
+            "TARGET",
+            "remove",
+            "gap one",
+            "gap two",
+            "gap three",
+            "TARGET",
+            "remove",
+            "footer"
+        ].joined(separator: "\n")
+        let request = ApplyEditsRequest(
+            path: "file.swift",
+            mode: .batch([
+                ApplyEditsOperation(
+                    search: "target\nremove",
+                    replace: "replacement",
+                    replaceAll: true
+                )
+            ]),
+            verbose: false
+        )
+
+        let result = try await ApplyEditsEngine.default.apply(request: request, to: original)
+
+        XCTAssertNil(result.note, "Case normalization should route through batch diff generation")
+        XCTAssertEqual(
+            result.updatedText,
+            [
+                "header",
+                "padding",
+                "replacement",
+                "gap one",
+                "gap two",
+                "gap three",
+                "replacement",
+                "footer"
+            ].joined(separator: "\n")
+        )
+    }
+
     func testReplaceAllBypassesDuplicateMatchAmbiguityAndAppliesCumulativeOffsets() async throws {
         let rows = [
             (