Feat: Highlighter Provider Diffing (#291)

Author: thecoolwinter

Example/CodeEditSourceEditorExample/CodeEditSourceEditorExample.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -57,7 +57,7 @@ public struct CodeEditSourceEditor: NSViewControllerRepresentable {
         editorOverscroll: CGFloat = 0,
         cursorPositions: Binding<[CursorPosition]>,
         useThemeBackground: Bool = true,
-        highlightProviders: [HighlightProviding] = [TreeSitterClient()],
+        highlightProviders: [any HighlightProviding] = [TreeSitterClient()],
         contentInsets: NSEdgeInsets? = nil,
         isEditable: Bool = true,
         isSelectable: Bool = true,
@@ -132,7 +132,7 @@ public struct CodeEditSourceEditor: NSViewControllerRepresentable {
         editorOverscroll: CGFloat = 0,
         cursorPositions: Binding<[CursorPosition]>,
         useThemeBackground: Bool = true,
-        highlightProviders: [HighlightProviding] = [TreeSitterClient()],
+        highlightProviders: [any HighlightProviding] = [TreeSitterClient()],
         contentInsets: NSEdgeInsets? = nil,
         isEditable: Bool = true,
         isSelectable: Bool = true,
@@ -179,7 +179,7 @@ public struct CodeEditSourceEditor: NSViewControllerRepresentable {
     private var editorOverscroll: CGFloat
     package var cursorPositions: Binding<[CursorPosition]>
     private var useThemeBackground: Bool
-    private var highlightProviders: [HighlightProviding]
+    private var highlightProviders: [any HighlightProviding]
     private var contentInsets: NSEdgeInsets?
     private var isEditable: Bool
     private var isSelectable: Bool
@@ -305,6 +305,10 @@ public struct CodeEditSourceEditor: NSViewControllerRepresentable {
             controller.useSystemCursor = useSystemCursor
         }
 
+        if !areHighlightProvidersEqual(controller: controller) {
+            controller.setHighlightProviders(highlightProviders)
+        }
+
         controller.bracketPairHighlight = bracketPairHighlight
     }
 
@@ -326,7 +330,12 @@ public struct CodeEditSourceEditor: NSViewControllerRepresentable {
         controller.tabWidth == tabWidth &&
         controller.letterSpacing == letterSpacing &&
         controller.bracketPairHighlight == bracketPairHighlight &&
-        controller.useSystemCursor == useSystemCursor
+        controller.useSystemCursor == useSystemCursor &&
+        areHighlightProvidersEqual(controller: controller)
+    }
+
+    private func areHighlightProvidersEqual(controller: TextViewController) -> Bool {
+        controller.highlightProviders.map { ObjectIdentifier($0) } == highlightProviders.map { ObjectIdentifier($0) }
     }
 }
 

Package.resolved

@@ -9,7 +9,7 @@ import Foundation
 import SwiftTreeSitter
 
 extension TextViewController {
-    internal func setUpHighlighter() {
+    package func setUpHighlighter() {
         if let highlighter {
             textView.removeStorageDelegate(highlighter)
             self.highlighter = nil
@@ -24,6 +24,17 @@ extension TextViewController {
         textView.addStorageDelegate(highlighter)
         self.highlighter = highlighter
     }
+
+    /// Sets new highlight providers. Recognizes when objects move in the array or are removed or inserted.
+    ///
+    /// This is in place of a setter on the ``highlightProviders`` variable to avoid wasting resources setting up
+    /// providers early.
+    ///
+    /// - Parameter newProviders: All the new providers.
+    package func setHighlightProviders(_ newProviders: [HighlightProviding]) {
+        highlighter?.setProviders(newProviders)
+        highlightProviders = newProviders
+    }
 }
 
 extension TextViewController: ThemeAttributesProviding {

Sources/CodeEditSourceEditor/CodeEditSourceEditor/CodeEditSourceEditor.swift

@@ -50,7 +50,8 @@ class HighlightProviderState {
     private weak var delegate: HighlightProviderStateDelegate?
 
     /// Calculates invalidated ranges given an edit.
-    private weak var highlightProvider: HighlightProviding?
+    /// Marked as package for deduplication when updating highlight providers.
+    package weak var highlightProvider: HighlightProviding?
 
     /// Provides a constantly updated visible index set.
     private weak var visibleRangeProvider: VisibleRangeProvider?

Sources/CodeEditSourceEditor/Controller/TextViewController+Highlighter.swift

@@ -78,6 +78,9 @@ class Highlighter: NSObject {
 
     private var visibleRangeProvider: VisibleRangeProvider
 
+    /// Counts upwards to provide unique IDs for new highlight providers.
+    private var providerIdCounter: Int
+
     // MARK: - Init
 
     init(
@@ -90,10 +93,12 @@ class Highlighter: NSObject {
         self.textView = textView
         self.attributeProvider = attributeProvider
 
-        visibleRangeProvider = VisibleRangeProvider(textView: textView)
+        self.visibleRangeProvider = VisibleRangeProvider(textView: textView)
 
         let providerIds = providers.indices.map({ $0 })
-        styleContainer = StyledRangeContainer(documentLength: textView.length, providers: providerIds)
+        self.styleContainer = StyledRangeContainer(documentLength: textView.length, providers: providerIds)
+
+        self.providerIdCounter = providers.count
 
         super.init()
 
@@ -138,6 +143,66 @@ class Highlighter: NSObject {
         highlightProviders.forEach { $0.setLanguage(language: language) }
     }
 
+    /// Updates the highlight providers the highlighter is using, removing any that don't appear in the given array,
+    /// and setting up any new ones.
+    ///
+    /// This is essential for working with SwiftUI, as we'd like to allow highlight providers to be added and removed
+    /// after the view is initialized. For instance after some sort of async registration method.
+    ///
+    /// - Note: Each provider will be identified by it's object ID.
+    /// - Parameter providers: All providers to use.
+    public func setProviders(_ providers: [HighlightProviding]) {
+        guard let textView else { return }
+
+        let existingIds: [ObjectIdentifier] = self.highlightProviders
+            .compactMap { $0.highlightProvider }
+            .map { ObjectIdentifier($0) }
+        let newIds: [ObjectIdentifier] = providers.map { ObjectIdentifier($0) }
+        // 2nd param is what we're moving *from*. We want to find how we to make existingIDs equal newIDs
+        let difference = newIds.difference(from: existingIds).inferringMoves()
+
+        var highlightProviders = self.highlightProviders // Make a mutable copy
+        var moveMap: [Int: HighlightProviderState] = [:]
+
+        for change in difference {
+            switch change {
+            case let .insert(offset, element, associatedOffset):
+                guard associatedOffset == nil,
+                      let newProvider = providers.first(where: { ObjectIdentifier($0) == element }) else {
+                    // Moved, grab the moved object from the move map
+                    guard let movedProvider = moveMap[offset] else {
+                        continue
+                    }
+                    highlightProviders.insert(movedProvider, at: offset)
+                    continue
+                }
+                // Set up a new provider and insert it with a unique ID
+                providerIdCounter += 1
+                let state = HighlightProviderState( // This will call setup on the highlight provider
+                    id: providerIdCounter,
+                    delegate: styleContainer,
+                    highlightProvider: newProvider,
+                    textView: textView,
+                    visibleRangeProvider: visibleRangeProvider,
+                    language: language
+                )
+                highlightProviders.insert(state, at: offset)
+                styleContainer.addProvider(providerIdCounter, documentLength: textView.length)
+                state.invalidate() // Invalidate this new one
+            case let .remove(offset, _, associatedOffset):
+                guard associatedOffset == nil else {
+                    // Moved, add it to the move map
+                    moveMap[associatedOffset!] = highlightProviders.remove(at: offset)
+                    continue
+                }
+                // Removed entirely
+                styleContainer.removeProvider(highlightProviders.remove(at: offset).id)
+            }
+        }
+
+        self.highlightProviders = highlightProviders
+    }
+
     deinit {
         self.attributeProvider = nil
         self.textView = nil

Sources/CodeEditSourceEditor/Highlighting/HighlighProviding/HighlightProviderState.swift

@@ -32,6 +32,21 @@ class StyledRangeContainer {
         }
     }
 
+    func addProvider(_ id: ProviderID, documentLength: Int) {
+        assert(!_storage.keys.contains(id), "Provider already exists")
+        _storage[id] = StyledRangeStore(documentLength: documentLength)
+    }
+
+    func removeProvider(_ id: ProviderID) {
+        guard let provider = _storage[id] else { return }
+        applyHighlightResult(
+            provider: id,
+            highlights: [],
+            rangeToHighlight: NSRange(location: 0, length: provider.length)
+        )
+        _storage.removeValue(forKey: id)
+    }
+
     /// Coalesces all styled runs into a single continuous array of styled runs.
     ///
     /// When there is an overlapping, conflicting style (eg: provider 2 gives `.comment` to the range `0..<2`, and

Sources/CodeEditSourceEditor/Highlighting/Highlighter.swift

@@ -17,6 +17,10 @@ final class StyledRangeStore {
     typealias Index = Rope<StyledRun>.Index
     var _guts = Rope<StyledRun>()
 
+    var length: Int {
+        _guts.count(in: OffsetMetric())
+    }
+
     /// A small performance improvement for multiple identical queries, as often happens when used
     /// in ``StyledRangeContainer``
     private var cache: (range: Range<Int>, runs: [Run])?