Fix Bugs With Multiple Highlight Providers

Sources/CodeEditSourceEditor/Highlighting/Highlighter.swift

@@ -163,7 +163,7 @@ class Highlighter: NSObject {
         let difference = newIds.difference(from: existingIds).inferringMoves()
 
         var highlightProviders = self.highlightProviders // Make a mutable copy
-        var moveMap: [Int: HighlightProviderState] = [:]
+        var moveMap: [Int: (Int, HighlightProviderState)] = [:]
 
         for change in difference {
             switch change {
@@ -174,7 +174,8 @@ class Highlighter: NSObject {
                     guard let movedProvider = moveMap[offset] else {
                         continue
                     }
-                    highlightProviders.insert(movedProvider, at: offset)
+                    highlightProviders.insert(movedProvider.1, at: offset)
+                    styleContainer.setPriority(providerId: movedProvider.0, priority: offset)
                     continue
                 }
                 // Set up a new provider and insert it with a unique ID
@@ -188,12 +189,12 @@ class Highlighter: NSObject {
                     language: language
                 )
                 highlightProviders.insert(state, at: offset)
-                styleContainer.addProvider(providerIdCounter, documentLength: textView.length)
+                styleContainer.addProvider(providerIdCounter, priority: offset, documentLength: textView.length)
                 state.invalidate() // Invalidate this new one
             case let .remove(offset, _, associatedOffset):
-                guard associatedOffset == nil else {
+                if let associatedOffset {
                     // Moved, add it to the move map
-                    moveMap[associatedOffset!] = highlightProviders.remove(at: offset)
+                    moveMap[associatedOffset] = (offset, highlightProviders.remove(at: offset))
                     continue
                 }
                 // Removed entirely

Sources/CodeEditSourceEditor/Highlighting/StyledRangeContainer/StyledRangeContainer+runsIn.swift

@@ -0,0 +1,72 @@
+//
+//  StyledRangeContainer+runsIn.swift
+//  CodeEditSourceEditor
+//
+//  Created by Khan Winter on 7/18/25.
+//
+
+import Foundation
+
+extension StyledRangeContainer {
+    /// 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
+    /// provider 1 gives `.string` to `1..<2`), the provider with a lower identifier will be prioritized. In the example
+    /// case, the final value would be `0..<1=.comment` and `1..<2=.string`.
+    ///
+    /// - Parameter range: The range to query.
+    /// - Returns: An array of continuous styled runs.
+    func runsIn(range: NSRange) -> [RangeStoreRun<StyleElement>] {
+        func combineLowerPriority(_ lhs: inout RangeStoreRun<StyleElement>, _ rhs: RangeStoreRun<StyleElement>) {
+            lhs.value = lhs.value?.combineLowerPriority(rhs.value) ?? rhs.value
+        }
+
+        func combineHigherPriority(_ lhs: inout RangeStoreRun<StyleElement>, _ rhs: RangeStoreRun<StyleElement>) {
+            lhs.value = lhs.value?.combineHigherPriority(rhs.value) ?? rhs.value
+        }
+
+        // Ordered by priority, lower = higher priority.
+        var allRuns = _storage.values
+            .sorted(by: { $0.priority < $1.priority })
+            .map { $0.store.runs(in: range.intRange) }
+
+        var runs: [RangeStoreRun<StyleElement>] = []
+        var minValue = allRuns.compactMap { $0.last }.enumerated().min(by: { $0.1.length < $1.1.length })
+        var counter = 0
+
+        while let value = minValue {
+            // Get minimum length off the end of each array
+            let minRunIdx = value.offset
+            var minRun = value.element
+
+            for idx in (0..<allRuns.count).reversed() where idx != minRunIdx {
+                guard let last = allRuns[idx].last else { continue }
+
+                if idx < minRunIdx {
+                    combineHigherPriority(&minRun, last)
+                } else {
+                    combineLowerPriority(&minRun, last)
+                }
+
+                if last.length == minRun.length {
+                    allRuns[idx].removeLast()
+                } else {
+                    // safe due to guard a few lines above.
+                    allRuns[idx][allRuns[idx].count - 1].subtractLength(minRun)
+                }
+            }
+
+            if !allRuns[minRunIdx].isEmpty {
+                allRuns[minRunIdx].removeLast()
+            }
+
+            assert(minRun.length > 0, "Empty or negative runs are not allowed.")
+            runs.append(minRun)
+            minValue = allRuns.compactMap { $0.last }.enumerated().min(by: { $0.1.length < $1.1.length })
+
+            counter += 1
+        }
+
+        return runs.reversed()
+    }
+}

Sources/CodeEditSourceEditor/Highlighting/StyledRangeContainer/StyledRangeContainer.swift

@@ -70,7 +70,7 @@ class StyledRangeContainer {
         }
     }
 
-    var _storage: [ProviderID: RangeStore<StyleElement>] = [:]
+    var _storage: [ProviderID: (store: RangeStore<StyleElement>, priority: Int)] = [:]
     weak var delegate: StyledRangeContainerDelegate?
 
     /// Initialize the container with a list of provider identifiers. Each provider is given an id, they should be
@@ -80,17 +80,21 @@ class StyledRangeContainer {
     ///   - providers: An array of identifiers given to providers.
     init(documentLength: Int, providers: [ProviderID]) {
         for provider in providers {
-            _storage[provider] = RangeStore<StyleElement>(documentLength: documentLength)
+            _storage[provider] = (store: RangeStore<StyleElement>(documentLength: documentLength), priority: provider)
         }
     }
 
-    func addProvider(_ id: ProviderID, documentLength: Int) {
+    func addProvider(_ id: ProviderID, priority: Int, documentLength: Int) {
         assert(!_storage.keys.contains(id), "Provider already exists")
-        _storage[id] = RangeStore<StyleElement>(documentLength: documentLength)
+        _storage[id] = (store: RangeStore<StyleElement>(documentLength: documentLength), priority: priority)
+    }
+
+    func setPriority(providerId: ProviderID, priority: Int) {
+        _storage[providerId]?.priority = priority
     }
 
     func removeProvider(_ id: ProviderID) {
-        guard let provider = _storage[id] else { return }
+        guard let provider = _storage[id]?.store else { return }
         applyHighlightResult(
             provider: id,
             highlights: [],
@@ -99,64 +103,9 @@ class StyledRangeContainer {
         _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
-    /// provider 1 gives `.string` to `1..<2`), the provider with a lower identifier will be prioritized. In the example
-    /// case, the final value would be `0..<1=.comment` and `1..<2=.string`.
-    ///
-    /// - Parameter range: The range to query.
-    /// - Returns: An array of continuous styled runs.
-    func runsIn(range: NSRange) -> [RangeStoreRun<StyleElement>] {
-        func combineLowerPriority(_ lhs: inout RangeStoreRun<StyleElement>, _ rhs: RangeStoreRun<StyleElement>) {
-            lhs.value = lhs.value?.combineLowerPriority(rhs.value) ?? rhs.value
-        }
-
-        func combineHigherPriority(_ lhs: inout RangeStoreRun<StyleElement>, _ rhs: RangeStoreRun<StyleElement>) {
-            lhs.value = lhs.value?.combineHigherPriority(rhs.value) ?? rhs.value
-        }
-
-        // Ordered by priority, lower = higher priority.
-        var allRuns = _storage.sorted(by: { $0.key < $1.key }).map { $0.value.runs(in: range.intRange) }
-        var runs: [RangeStoreRun<StyleElement>] = []
-
-        var minValue = allRuns.compactMap { $0.last }.enumerated().min(by: { $0.1.length < $1.1.length })
-
-        while let value = minValue {
-            // Get minimum length off the end of each array
-            let minRunIdx = value.offset
-            var minRun = value.element
-
-            for idx in (0..<allRuns.count).reversed() where idx != minRunIdx {
-                guard let last = allRuns[idx].last else { continue }
-                if idx < minRunIdx {
-                    combineHigherPriority(&minRun, last)
-                } else {
-                    combineLowerPriority(&minRun, last)
-                }
-
-                if last.length == minRun.length {
-                    allRuns[idx].removeLast()
-                } else {
-                    // safe due to guard a few lines above.
-                    allRuns[idx][allRuns[idx].count - 1].subtractLength(minRun)
-                }
-            }
-
-            if !allRuns[minRunIdx].isEmpty {
-                allRuns[minRunIdx].removeLast()
-            }
-
-            runs.append(minRun)
-            minValue = allRuns.compactMap { $0.last }.enumerated().min(by: { $0.1.length < $1.1.length })
-        }
-
-        return runs.reversed()
-    }
-
     func storageUpdated(editedRange: NSRange, changeInLength delta: Int) {
         for key in _storage.keys {
-            _storage[key]?.storageUpdated(editedRange: editedRange, changeInLength: delta)
+            _storage[key]?.store.storageUpdated(editedRange: editedRange, changeInLength: delta)
         }
     }
 }
@@ -171,7 +120,7 @@ extension StyledRangeContainer: HighlightProviderStateDelegate {
     ///   - rangeToHighlight: The range to apply the highlights to.
     func applyHighlightResult(provider: ProviderID, highlights: [HighlightRange], rangeToHighlight: NSRange) {
         assert(rangeToHighlight != .notFound, "NSNotFound is an invalid highlight range")
-        guard var storage = _storage[provider] else {
+        guard var storage = _storage[provider]?.store else {
             assertionFailure("No storage found for the given provider: \(provider)")
             return
         }
@@ -193,12 +142,12 @@ extension StyledRangeContainer: HighlightProviderStateDelegate {
             lastIndex = highlight.range.max
         }
 
-        if lastIndex != rangeToHighlight.upperBound {
+        if lastIndex < rangeToHighlight.upperBound {
             runs.append(.empty(length: rangeToHighlight.upperBound - lastIndex))
         }
 
         storage.set(runs: runs, for: rangeToHighlight.intRange)
-        _storage[provider] = storage
+        _storage[provider]?.store = storage
         delegate?.styleContainerDidUpdate(in: rangeToHighlight)
     }
 }

Sources/CodeEditSourceEditor/RangeStore/RangeStore.swift

@@ -39,23 +39,31 @@ struct RangeStore<Element: RangeStoreElement>: Sendable {
     /// - Parameter range: The range to query.
     /// - Returns: A continuous array of runs representing the queried range.
     func runs(in range: Range<Int>) -> [Run] {
+        let length = _guts.count(in: OffsetMetric())
         assert(range.lowerBound >= 0, "Negative lowerBound")
-        assert(range.upperBound <= _guts.count(in: OffsetMetric()), "upperBound outside valid range")
+        assert(range.upperBound <= length, "upperBound outside valid range")
         if let cache, cache.range == range {
             return cache.runs
         }
 
         var runs = [Run]()
-
         var index = findIndex(at: range.lowerBound).index
-        var offset: Int? = range.lowerBound - _guts.offset(of: index, in: OffsetMetric())
+        var offset: Int = range.lowerBound - _guts.offset(of: index, in: OffsetMetric())
+        var remainingLength = range.upperBound - range.lowerBound
 
-        while index < _guts.endIndex, _guts.offset(of: index, in: OffsetMetric()) < range.upperBound {
+        while index < _guts.endIndex,
+              _guts.offset(of: index, in: OffsetMetric()) < range.upperBound,
+              remainingLength > 0 {
             let run = _guts[index]
-            runs.append(Run(length: run.length - (offset ?? 0), value: run.value))
+            let runLength = min(run.length - offset, remainingLength)
+            runs.append(Run(length: runLength, value: run.value))
 
+            remainingLength -= runLength
+            if remainingLength <= 0 {
+                break // Avoid even checking the storage for the next index
+            }
             index = _guts.index(after: index)
-            offset = nil
+            offset = 0
         }
 
         return runs
@@ -83,6 +91,9 @@ struct RangeStore<Element: RangeStoreElement>: Sendable {
             storageUpdated(replacedCharactersIn: upperBound..<upperBound, withCount: missingCharacters)
         }
 
+        // This is quite slow in debug builds but is a *really* important assertion for internal state.
+        assert(!runs.contains(where: { $0.length < 0 }), "Runs cannot have negative length.")
+
         _guts.replaceSubrange(
             range,
             in: OffsetMetric(),

Tests/CodeEditSourceEditorTests/Highlighting/StyledRangeContainerTests.swift

@@ -11,7 +11,10 @@ final class StyledRangeContainerTests: XCTestCase {
 
         // Have to do string conversion due to missing Comparable conformance pre-macOS 14
         XCTAssertEqual(store._storage.keys.sorted(), providers)
-        XCTAssert(store._storage.values.allSatisfy({ $0.length == 100 }), "One or more providers have incorrect length")
+        XCTAssert(
+            store._storage.values.allSatisfy({ $0.store.length == 100 }),
+            "One or more providers have incorrect length"
+        )
     }
 
     @MainActor
@@ -26,10 +29,10 @@ final class StyledRangeContainerTests: XCTestCase {
         )
 
         XCTAssertNotNil(store._storage[providers[0]])
-        XCTAssertEqual(store._storage[providers[0]]!.count, 3)
-        XCTAssertNil(store._storage[providers[0]]!.runs(in: 0..<100)[0].value?.capture)
-        XCTAssertEqual(store._storage[providers[0]]!.runs(in: 0..<100)[1].value?.capture, .comment)
-        XCTAssertNil(store._storage[providers[0]]!.runs(in: 0..<100)[2].value?.capture)
+        XCTAssertEqual(store._storage[providers[0]]!.store.count, 3)
+        XCTAssertNil(store._storage[providers[0]]!.store.runs(in: 0..<100)[0].value?.capture)
+        XCTAssertEqual(store._storage[providers[0]]!.store.runs(in: 0..<100)[1].value?.capture, .comment)
+        XCTAssertNil(store._storage[providers[0]]!.store.runs(in: 0..<100)[2].value?.capture)
 
         XCTAssertEqual(
             store.runsIn(range: NSRange(location: 0, length: 100)),