Highlight Invalidation Fix

Author: thecoolwinter

Sources/CodeEditSourceEditor/Highlighting/HighlighProviding/HighlightProviderState.swift

@@ -116,6 +116,16 @@ class HighlightProviderState {
 
         highlightInvalidRanges()
     }
+
+    /// Accumulates all pending ranges and calls `queryHighlights`.
+    func highlightInvalidRanges() {
+        var ranges: [NSRange] = []
+        while let nextRange = getNextRange() {
+            ranges.append(nextRange)
+            pendingSet.insert(range: nextRange)
+        }
+        queryHighlights(for: ranges)
+    }
 }
 
 extension HighlightProviderState {
@@ -143,16 +153,6 @@ extension HighlightProviderState {
 }
 
 private extension HighlightProviderState {
-    /// Accumulates all pending ranges and calls `queryHighlights`.
-    func highlightInvalidRanges() {
-        var ranges: [NSRange] = []
-        while let nextRange = getNextRange() {
-            ranges.append(nextRange)
-            pendingSet.insert(range: nextRange)
-        }
-        queryHighlights(for: ranges)
-    }
-
     /// Gets the next `NSRange` to highlight based on the invalid set, visible set, and pending set.
     /// - Returns: An `NSRange` to highlight if it could be fetched.
     func getNextRange() -> NSRange? {

Sources/CodeEditSourceEditor/Highlighting/Highlighter.swift

@@ -118,6 +118,10 @@ class Highlighter: NSObject {
         highlightProviders.forEach { $0.invalidate() }
     }
 
+    public func invalidate(_ set: IndexSet) {
+        highlightProviders.forEach { $0.invalidate(set) }
+    }
+
     /// Sets the language and causes a re-highlight of the entire text.
     /// - Parameter language: The language to update to.
     public func setLanguage(language: CodeLanguage) {
@@ -154,11 +158,6 @@ extension Highlighter: NSTextStorageDelegate {
         // This method is called whenever attributes are updated, so to avoid re-highlighting the entire document
         // each time an attribute is applied, we check to make sure this is in response to an edit.
         guard editedMask.contains(.editedCharacters), let textView else { return }
-        if delta > 0 {
-            visibleRangeProvider.visibleSet.insert(range: editedRange)
-        }
-
-        visibleRangeProvider.updateVisibleSet(textView: textView)
 
         let styleContainerRange: Range<Int>
         let newLength: Int
@@ -176,6 +175,12 @@ extension Highlighter: NSTextStorageDelegate {
             withCount: newLength
         )
 
+        if delta > 0 {
+            visibleRangeProvider.visibleSet.insert(range: editedRange)
+        }
+
+        visibleRangeProvider.updateVisibleSet(textView: textView)
+
         let providerRange = NSRange(location: editedRange.location, length: editedRange.length - delta)
         highlightProviders.forEach { $0.storageDidUpdate(range: providerRange, delta: delta) }
     }
@@ -196,26 +201,30 @@ extension Highlighter: NSTextStorageDelegate {
 extension Highlighter: StyledRangeContainerDelegate {
     func styleContainerDidUpdate(in range: NSRange) {
         guard let textView, let attributeProvider else { return }
+//        textView.layoutManager.beginTransaction()
         textView.textStorage.beginEditing()
 
         let storage = textView.textStorage
 
         var offset = range.location
         for run in styleContainer.runsIn(range: range) {
-            let range = NSRange(location: offset, length: run.length)
+            guard let range = NSRange(location: offset, length: run.length).intersection(range) else {
+                continue
+            }
             storage?.setAttributes(attributeProvider.attributesFor(run.capture), range: range)
-            offset += run.length
+            offset += range.length
         }
 
         textView.textStorage.endEditing()
-        textView.layoutManager.invalidateLayoutForRange(range)
+//        textView.layoutManager.endTransaction()
+//        textView.layoutManager.invalidateLayoutForRange(range)
     }
 }
 
 // MARK: - VisibleRangeProviderDelegate
 
 extension Highlighter: VisibleRangeProviderDelegate {
     func visibleSetDidUpdate(_ newIndices: IndexSet) {
-        highlightProviders.forEach { $0.invalidate(newIndices) }
+        highlightProviders.forEach { $0.highlightInvalidRanges() }
     }
 }

Sources/CodeEditSourceEditor/TreeSitter/TreeSitterClient.swift

@@ -52,7 +52,7 @@ public final class TreeSitterClient: HighlightProviding {
     package var pendingEdits: Atomic<[InputEdit]> = Atomic([])
 
     /// Optional flag to force every operation to be done on the caller's thread.
-    var forceSyncOperation: Bool = false
+    package var forceSyncOperation: Bool = false
 
     public init() { }
 

Tests/CodeEditSourceEditorTests/Highlighting/HighlighterTests.swift

@@ -59,4 +59,42 @@ final class HighlighterTests: XCTestCase {
             "Highlighter did not query again after cancelling the first request"
         )
     }
+
+    @MainActor
+    func test_highlightsDoNotInvalidateEntireTextView() {
+        class SentryStorageDelegate: NSObject, NSTextStorageDelegate {
+            var editedIndices: IndexSet = IndexSet()
+
+            func textStorage(
+                _ textStorage: NSTextStorage,
+                didProcessEditing editedMask: NSTextStorageEditActions,
+                range editedRange: NSRange,
+                changeInLength delta: Int) {
+                editedIndices.insert(integersIn: editedRange)
+            }
+        }
+
+        let highlightProvider = TreeSitterClient()
+        highlightProvider.forceSyncOperation = true
+        let attributeProvider = MockAttributeProvider()
+        let textView = Mock.textView()
+        textView.frame = NSRect(x: 0, y: 0, width: 1000, height: 1000)
+        textView.setText("func helloWorld() {\n\tprint(\"Hello World!\")\n}")
+
+        let highlighter = Mock.highlighter(
+            textView: textView,
+            highlightProvider: highlightProvider,
+            attributeProvider: attributeProvider
+        )
+
+        highlighter.invalidate()
+
+        let sentryStorage = SentryStorageDelegate()
+        textView.addStorageDelegate(sentryStorage)
+
+        let invalidSet = IndexSet(integersIn: NSRange(location: 0, length: 24))
+        highlighter.invalidate(invalidSet) // Invalidate first line
+
+        XCTAssertEqual(sentryStorage.editedIndices, invalidSet) // Should only cause highlights on the first line
+    }
 }