Filter Ambiguous TreeSitter Highlights

Author: thecoolwinter

Sources/CodeEditSourceEditor/Highlighting/HighlightRange.swift

@@ -8,7 +8,7 @@
 import Foundation
 
 /// This struct represents a range to highlight, as well as the capture name for syntax coloring.
-public struct HighlightRange: Sendable {
+public struct HighlightRange: Hashable, Sendable {
     public let range: NSRange
     public let capture: CaptureName?
     public let modifiers: CaptureModifierSet
@@ -19,3 +19,13 @@ public struct HighlightRange: Sendable {
         self.modifiers = modifiers
     }
 }
+
+extension HighlightRange: CustomDebugStringConvertible {
+    public var debugDescription: String {
+        if capture == nil && modifiers.isEmpty {
+            "\(range) (empty)"
+        } else {
+            "\(range) (\(capture?.stringValue ?? "No Capture")) \(modifiers.values.map({ $0.stringValue }))"
+        }
+    }
+}

Sources/CodeEditSourceEditor/TreeSitter/TreeSitterClient+Highlight.swift

@@ -87,17 +87,35 @@ extension TreeSitterClient {
         cursor: QueryCursor,
         includedRange: NSRange
     ) -> [HighlightRange] {
+        guard let readCallback else { return [] }
+        var ranges: [NSRange: Int] = [:]
         return cursor
+            .resolve(with: .init(textProvider: readCallback)) // Resolve our cursor against the query
             .flatMap { $0.captures }
-            .compactMap {
-                // Sometimes `cursor.setRange` just doesn't work :( so we have to do a redundant check for a valid range
-                // in the included range
-                let intersectionRange = $0.range.intersection(includedRange) ?? .zero
-                // Check that the capture name is one CESE can parse. If not, ignore it completely.
-                if intersectionRange.length > 0, let captureName = CaptureName.fromString($0.name) {
-                    return HighlightRange(range: intersectionRange, capture: captureName)
+            .reversed() // SwiftTreeSitter returns captures in the reverse order of what we need to filter with.
+            .compactMap { capture in
+                let range = capture.range
+                let index = capture.index
+
+                // Lower indexed captures are favored over higher, this is why we reverse it above
+                if let existingLevel = ranges[range], existingLevel <= index {
+                    return nil
+                }
+
+                guard let captureName = CaptureName.fromString(capture.name) else {
+                    return nil
                 }
-                return nil
+
+                // Update the filter level to the current index since it's lower and a 'valid' capture
+                ranges[range] = index
+
+                // Validate range and capture name
+                let intersectionRange = range.intersection(includedRange) ?? .zero
+                guard intersectionRange.length > 0 else {
+                    return nil
+                }
+
+                return HighlightRange(range: intersectionRange, capture: captureName)
             }
     }
 }

Sources/CodeEditSourceEditor/TreeSitter/TreeSitterClient.swift

@@ -210,7 +210,7 @@ public final class TreeSitterClient: HighlightProviding {
         completion: @escaping @MainActor (Result<[HighlightRange], Error>) -> Void
     ) {
         let operation = { [weak self] in
-            return self?.queryHighlightsForRange(range: range) ?? []
+            return (self?.queryHighlightsForRange(range: range) ?? []).sorted { $0.range.location < $1.range.location }
         }
 
         let longQuery = range.length > Constants.maxSyncQueryLength