diff --git a/Sources/CodeEditTextView/Controller/STTextViewController.swift b/Sources/CodeEditTextView/Controller/STTextViewController.swift index 8f8986716..cc7eeea53 100644 --- a/Sources/CodeEditTextView/Controller/STTextViewController.swift +++ b/Sources/CodeEditTextView/Controller/STTextViewController.swift @@ -320,7 +320,7 @@ public class STTextViewController: NSViewController, STTextViewDelegate, ThemeAt return self?.textView.textContentStorage.textStorage?.mutableString.substring(with: range) } - provider = try? TreeSitterClient(codeLanguage: language, textProvider: textProvider) + provider = TreeSitterClient(codeLanguage: language, textProvider: textProvider) } if let provider = provider { diff --git a/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+Comparable.swift b/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+Comparable.swift new file mode 100644 index 000000000..15c7e3b12 --- /dev/null +++ b/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+Comparable.swift @@ -0,0 +1,18 @@ +// +// NSRange+Comparable.swift +// +// +// Created by Khan Winter on 3/15/23. +// + +import Foundation + +extension NSRange: Comparable { + public static func == (lhs: NSRange, rhs: NSRange) -> Bool { + return lhs.location == rhs.location && lhs.length == rhs.length + } + + public static func < (lhs: NSRange, rhs: NSRange) -> Bool { + return lhs.location < rhs.location + } +} diff --git a/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+InputEdit.swift b/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+InputEdit.swift index 9994a92df..ee8018de5 100644 --- a/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+InputEdit.swift +++ b/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+InputEdit.swift @@ -29,3 +29,35 @@ extension InputEdit { newEndPoint: newEndPoint) } } + +extension NSRange { + // swiftlint:disable line_length + /// Modifies the range to account for an edit. + /// Largely based on code from + /// [tree-sitter](https://github.com/tree-sitter/tree-sitter/blob/ddeaa0c7f534268b35b4f6cb39b52df082754413/lib/src/subtree.c#L691-L720) + mutating func applyInputEdit(_ edit: InputEdit) { + // swiftlint:enable line_length + let endIndex = NSMaxRange(self) + let isPureInsertion = edit.oldEndByte == edit.startByte + + // Edit is after the range + if (edit.startByte/2) > endIndex { + return + } else if edit.oldEndByte/2 < location { + // If the edit is entirely before this range + self.location += (Int(edit.newEndByte) - Int(edit.oldEndByte))/2 + } else if edit.startByte/2 < location { + // If the edit starts in the space before this range and extends into this range + length -= Int(edit.oldEndByte)/2 - location + location = Int(edit.newEndByte)/2 + } else if edit.startByte/2 == location && isPureInsertion { + // If the edit is *only* an insertion right at the beginning of the range + location = Int(edit.newEndByte)/2 + } else { + // Otherwise, the edit is entirely within this range + if edit.startByte/2 < endIndex || (edit.startByte/2 == endIndex && isPureInsertion) { + length = (Int(edit.newEndByte)/2 - location) + (length - (Int(edit.oldEndByte)/2 - location)) + } + } + } +} diff --git a/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+TSRange.swift b/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+TSRange.swift new file mode 100644 index 000000000..a0273784e --- /dev/null +++ b/Sources/CodeEditTextView/Extensions/NSRange+/NSRange+TSRange.swift @@ -0,0 +1,18 @@ +// +// NSRange+TSRange.swift +// +// +// Created by Khan Winter on 2/26/23. +// + +import Foundation +import SwiftTreeSitter + +extension NSRange { + var tsRange: TSRange { + return TSRange( + points: .zero..<(.zero), + bytes: (UInt32(self.location) * 2)..<(UInt32(self.location + self.length) * 2) + ) + } +} diff --git a/Sources/CodeEditTextView/Extensions/Tree+prettyPrint.swift b/Sources/CodeEditTextView/Extensions/Tree+prettyPrint.swift new file mode 100644 index 000000000..a021733c9 --- /dev/null +++ b/Sources/CodeEditTextView/Extensions/Tree+prettyPrint.swift @@ -0,0 +1,71 @@ +// +// Tree+prettyPrint.swift +// +// +// Created by Khan Winter on 3/16/23. +// + +import SwiftTreeSitter + +#if DEBUG +extension Tree { + func prettyPrint() { + guard let cursor = self.rootNode?.treeCursor else { + print("NO ROOT NODE") + return + } + guard cursor.currentNode != nil else { + print("NO CURRENT NODE") + return + } + + func p(_ cursor: TreeCursor, depth: Int) { + guard let node = cursor.currentNode else { + return + } + + let visible = node.isNamed + + if visible { + print(String(repeating: " ", count: depth * 2), terminator: "") + if let fieldName = cursor.currentFieldName { + print(fieldName, ": ", separator: "", terminator: "") + } + print("(", node.nodeType ?? "NONE", " ", node.range, " ", separator: "", terminator: "") + } + + if cursor.goToFirstChild() { + while true { + if cursor.currentNode != nil && cursor.currentNode!.isNamed { + print("") + } + + p(cursor, depth: depth + 1) + + if !cursor.gotoNextSibling() { + break + } + } + + if !cursor.gotoParent() { + fatalError("Could not go to parent, this tree may be invalid.") + } + } + + if visible { + print(")", terminator: "") + } + } + + if cursor.currentNode?.childCount == 0 { + if !cursor.currentNode!.isNamed { + print("{\(cursor.currentNode!.nodeType ?? "NONE")}") + } else { + print("\"\(cursor.currentNode!.nodeType ?? "NONE")\"") + } + } else { + p(cursor, depth: 1) + } + } +} +#endif diff --git a/Sources/CodeEditTextView/Filters/STTextViewController+TextFormation.swift b/Sources/CodeEditTextView/Filters/STTextViewController+TextFormation.swift index 55bba4a2e..0e0ad7cac 100644 --- a/Sources/CodeEditTextView/Filters/STTextViewController+TextFormation.swift +++ b/Sources/CodeEditTextView/Filters/STTextViewController+TextFormation.swift @@ -71,7 +71,7 @@ extension STTextViewController { /// - whitespaceProvider: The whitespace providers to use. /// - indentationUnit: The unit of indentation to use. private func setUpNewlineTabFilters(whitespaceProvider: WhitespaceProviders, indentationUnit: String) { - let newlineFilter: Filter = NewlineFilter(whitespaceProviders: whitespaceProvider) + let newlineFilter: Filter = NewlineProcessingFilter(whitespaceProviders: whitespaceProvider) let tabReplacementFilter: Filter = TabReplacementFilter(indentationUnit: indentationUnit) textFilters.append(contentsOf: [newlineFilter, tabReplacementFilter]) diff --git a/Sources/CodeEditTextView/Highlighting/HighlightProviding.swift b/Sources/CodeEditTextView/Highlighting/HighlightProviding.swift index cd376ef9f..199c2cc42 100644 --- a/Sources/CodeEditTextView/Highlighting/HighlightProviding.swift +++ b/Sources/CodeEditTextView/Highlighting/HighlightProviding.swift @@ -17,6 +17,9 @@ public protocol HighlightProviding { /// - Note: This does not need to be *globally* unique, merely unique across all the highlighters used. var identifier: String { get } + /// Called once at editor initialization. + func setUp(textView: HighlighterTextView) + /// Updates the highlighter's code language. /// - Parameters: /// - codeLanguage: The langugage that should be used by the highlighter. diff --git a/Sources/CodeEditTextView/Highlighting/Highlighter.swift b/Sources/CodeEditTextView/Highlighting/Highlighter.swift index fe4990bca..ebda49b97 100644 --- a/Sources/CodeEditTextView/Highlighting/Highlighter.swift +++ b/Sources/CodeEditTextView/Highlighting/Highlighter.swift @@ -88,6 +88,7 @@ class Highlighter: NSObject { } textView.textContentStorage.textStorage?.delegate = self + highlightProvider?.setUp(textView: textView) if let scrollView = textView.enclosingScrollView { NotificationCenter.default.addObserver(self, @@ -121,6 +122,7 @@ class Highlighter: NSObject { public func setHighlightProvider(_ provider: HighlightProviding) { self.highlightProvider = provider highlightProvider?.setLanguage(codeLanguage: language) + highlightProvider?.setUp(textView: textView) invalidate() } @@ -282,7 +284,7 @@ extension Highlighter: NSTextStorageDelegate { delta: delta) { [weak self] invalidatedIndexSet in let indexSet = invalidatedIndexSet .union(IndexSet(integersIn: editedRange)) - // Only invalidate indices that aren't visible. + // Only invalidate indices that are visible. .intersection(self?.visibleSet ?? .init()) for range in indexSet.rangeView { diff --git a/Sources/CodeEditTextView/TreeSitter/TreeSitterClient+Edit.swift b/Sources/CodeEditTextView/TreeSitter/TreeSitterClient+Edit.swift new file mode 100644 index 000000000..b7a31329e --- /dev/null +++ b/Sources/CodeEditTextView/TreeSitter/TreeSitterClient+Edit.swift @@ -0,0 +1,163 @@ +// +// TreeSitterClient+Edit.swift +// +// +// Created by Khan Winter on 3/10/23. +// + +import Foundation +import SwiftTreeSitter +import CodeEditLanguages + +extension TreeSitterClient { + + /// Calculates a series of ranges that have been invalidated by a given edit. + /// - Parameters: + /// - textView: The text view to use for text. + /// - edit: The edit to act on. + /// - language: The language to use. + /// - readBlock: A callback for fetching blocks of text. + /// - Returns: An array of distinct `NSRanges` that need to be re-highlighted. + func findChangedByteRanges( + textView: HighlighterTextView, + edit: InputEdit, + layer: LanguageLayer, + readBlock: @escaping Parser.ReadBlock + ) -> [NSRange] { + let (oldTree, newTree) = calculateNewState( + tree: layer.tree, + parser: layer.parser, + edit: edit, + readBlock: readBlock + ) + if oldTree == nil && newTree == nil { + // There was no existing tree, make a new one and return all indexes. + layer.tree = createTree(parser: layer.parser, readBlock: readBlock) + return [NSRange(textView.documentRange.intRange)] + } + + let ranges = changedByteRanges(oldTree, rhs: newTree).map { $0.range } + + layer.tree = newTree + + return ranges + } + + /// Applies the edit to the current `tree` and returns the old tree and a copy of the current tree with the + /// processed edit. + /// - Parameters: + /// - tree: The tree before an edit used to parse the new tree. + /// - parser: The parser used to parse the new tree. + /// - edit: The edit to apply. + /// - readBlock: The block to use to read text. + /// - Returns: (The old state, the new state). + internal func calculateNewState( + tree: Tree?, + parser: Parser, + edit: InputEdit, + readBlock: @escaping Parser.ReadBlock + ) -> (Tree?, Tree?) { + guard let oldTree = tree else { + return (nil, nil) + } + semaphore.wait() + + // Apply the edit to the old tree + oldTree.edit(edit) + + let newTree = parser.parse(tree: oldTree, readBlock: readBlock) + + semaphore.signal() + + return (oldTree.copy(), newTree) + } + + /// Calculates the changed byte ranges between two trees. + /// - Parameters: + /// - lhs: The first (older) tree. + /// - rhs: The second (newer) tree. + /// - Returns: Any changed ranges. + internal func changedByteRanges(_ lhs: Tree?, rhs: Tree?) -> [Range] { + switch (lhs, rhs) { + case (let t1?, let t2?): + return t1.changedRanges(from: t2).map({ $0.bytes }) + case (nil, let t2?): + let range = t2.rootNode?.byteRange + + return range.flatMap({ [$0] }) ?? [] + case (_, nil): + return [] + } + } + + /// Performs an injections query on the given language layer. + /// Updates any existing layers with new ranges and adds new layers if needed. + /// - Parameters: + /// - textView: The text view to use. + /// - layer: The language layer to perform the query on. + /// - layerSet: The set of layers that exist in the document. + /// Used for efficient lookup of existing `(language, range)` pairs + /// - touchedLayers: The set of layers that existed before updating injected layers. + /// Will have items removed as they are found. + /// - readBlock: A completion block for reading from text storage efficiently. + /// - Returns: An index set of any updated indexes. + @discardableResult + internal func updateInjectedLanguageLayers( + textView: HighlighterTextView, + layer: LanguageLayer, + layerSet: inout Set, + touchedLayers: inout Set, + readBlock: @escaping Parser.ReadBlock + ) -> IndexSet { + guard let tree = layer.tree, + let rootNode = tree.rootNode, + let cursor = layer.languageQuery?.execute(node: rootNode, in: tree) else { + return IndexSet() + } + + cursor.matchLimit = Constants.treeSitterMatchLimit + + let languageRanges = self.injectedLanguagesFrom(cursor: cursor) { range, _ in + return textView.stringForRange(range) + } + + var updatedRanges = IndexSet() + + for (languageName, ranges) in languageRanges { + guard let treeSitterLanguage = TreeSitterLanguage(rawValue: languageName) else { + continue + } + + if treeSitterLanguage == primaryLayer { + continue + } + + for range in ranges { + // Temp layer object for + let layer = LanguageLayer( + id: treeSitterLanguage, + parser: Parser(), + supportsInjections: false, + ranges: [range.range] + ) + + if layerSet.contains(layer) { + // If we've found this layer, it means it should exist after an edit. + touchedLayers.remove(layer) + } else { + // New range, make a new layer! + if let addedLayer = addLanguageLayer(layerId: treeSitterLanguage, readBlock: readBlock) { + addedLayer.ranges = [range.range] + addedLayer.parser.includedRanges = addedLayer.ranges.map { $0.tsRange } + addedLayer.tree = createTree(parser: addedLayer.parser, readBlock: readBlock) + + layerSet.insert(addedLayer) + updatedRanges.insert(range: range.range) + } + } + } + } + + return updatedRanges + } +} diff --git a/Sources/CodeEditTextView/TreeSitter/TreeSitterClient+Highlight.swift b/Sources/CodeEditTextView/TreeSitter/TreeSitterClient+Highlight.swift new file mode 100644 index 000000000..de8764f0a --- /dev/null +++ b/Sources/CodeEditTextView/TreeSitter/TreeSitterClient+Highlight.swift @@ -0,0 +1,92 @@ +// +// TreeSitterClient+Highlight.swift +// +// +// Created by Khan Winter on 3/10/23. +// + +import Foundation +import SwiftTreeSitter +import CodeEditLanguages + +extension TreeSitterClient { + + /// Queries the given language layer for any highlights. + /// - Parameters: + /// - layer: The layer to query. + /// - textView: A text view to use for contextual data. + /// - range: The range to query for. + /// - Returns: Any ranges to highlight. + internal func queryLayerHighlights( + layer: LanguageLayer, + textView: HighlighterTextView, + range: NSRange + ) -> [HighlightRange] { + // Make sure we don't change the tree while we copy it. + self.semaphore.wait() + + guard let tree = layer.tree?.copy() else { + self.semaphore.signal() + return [] + } + + self.semaphore.signal() + + guard let rootNode = tree.rootNode else { + return [] + } + + // This needs to be on the main thread since we're going to use the `textProvider` in + // the `highlightsFromCursor` method, which uses the textView's text storage. + guard let cursor = layer.languageQuery?.execute(node: rootNode, in: tree) else { + return [] + } + cursor.setRange(range) + cursor.matchLimit = Constants.treeSitterMatchLimit + + return highlightsFromCursor(cursor: ResolvingQueryCursor(cursor: cursor)) + } + + /// Resolves a query cursor to the highlight ranges it contains. + /// **Must be called on the main thread** + /// - Parameter cursor: The cursor to resolve. + /// - Returns: Any highlight ranges contained in the cursor. + internal func highlightsFromCursor(cursor: ResolvingQueryCursor) -> [HighlightRange] { + cursor.prepare(with: self.textProvider) + return cursor + .flatMap { $0.captures } + .compactMap { + // Some languages add an "@spell" capture to indicate a portion of text that should be spellchecked + // (usually comments). But this causes other captures in the same range to be overriden. So we ignore + // that specific capture type. + if $0.name != "spell" && $0.name != "injection.content" { + return HighlightRange(range: $0.range, capture: CaptureName.fromString($0.name ?? "")) + } + return nil + } + } + + /// Returns all injected languages from a given cursor. The cursor must be new, + /// having not been used for normal highlight matching. + /// - Parameters: + /// - cursor: The cursor to use for finding injected languages. + /// - textProvider: A callback for efficiently fetching text. + /// - Returns: A map of each language to all the ranges they have been injected into. + internal func injectedLanguagesFrom( + cursor: QueryCursor, + textProvider: @escaping ResolvingQueryCursor.TextProvider + ) -> [String: [NamedRange]] { + var languages: [String: [NamedRange]] = [:] + + for match in cursor { + if let injection = match.injection(with: textProvider) { + if languages[injection.name] == nil { + languages[injection.name] = [] + } + languages[injection.name]?.append(injection) + } + } + + return languages + } +} diff --git a/Sources/CodeEditTextView/TreeSitter/TreeSitterClient+LanguageLayer.swift b/Sources/CodeEditTextView/TreeSitter/TreeSitterClient+LanguageLayer.swift new file mode 100644 index 000000000..c5ab2cc15 --- /dev/null +++ b/Sources/CodeEditTextView/TreeSitter/TreeSitterClient+LanguageLayer.swift @@ -0,0 +1,54 @@ +// +// TreeSitterClient+LanguageLayer.swift +// +// +// Created by Khan Winter on 3/8/23. +// + +import Foundation +import CodeEditLanguages +import SwiftTreeSitter + +extension TreeSitterClient { + class LanguageLayer: Hashable { + /// Initialize a language layer + /// - Parameters: + /// - id: The ID of the layer. + /// - parser: A parser to use for the layer. + /// - supportsInjections: Set to true when the langauge supports the `injections` query. + /// - tree: The tree-sitter tree generated while editing/parsing a document. + /// - languageQuery: The language query used for fetching the associated `queries.scm` file + /// - ranges: All ranges this layer acts on. Must be kept in order and w/o overlap. + init( + id: TreeSitterLanguage, + parser: Parser, + supportsInjections: Bool, + tree: Tree? = nil, + languageQuery: Query? = nil, + ranges: [NSRange] + ) { + self.id = id + self.parser = parser + self.supportsInjections = supportsInjections + self.tree = tree + self.languageQuery = languageQuery + self.ranges = ranges + } + + let id: TreeSitterLanguage + let parser: Parser + let supportsInjections: Bool + var tree: Tree? + var languageQuery: Query? + var ranges: [NSRange] + + static func == (lhs: TreeSitterClient.LanguageLayer, rhs: TreeSitterClient.LanguageLayer) -> Bool { + return lhs.id == rhs.id && lhs.ranges == rhs.ranges + } + + func hash(into hasher: inout Hasher) { + hasher.combine(id) + hasher.combine(ranges) + } + } +} diff --git a/Sources/CodeEditTextView/TreeSitter/TreeSitterClient.swift b/Sources/CodeEditTextView/TreeSitter/TreeSitterClient.swift index eea44f160..66276576f 100644 --- a/Sources/CodeEditTextView/TreeSitter/TreeSitterClient.swift +++ b/Sources/CodeEditTextView/TreeSitter/TreeSitterClient.swift @@ -18,46 +18,91 @@ import SwiftTreeSitter /// However, the `setText` method will re-compile the entire corpus so should be used sparingly. final class TreeSitterClient: HighlightProviding { + // MARK: - Properties/Constants + public var identifier: String { "CodeEdit.TreeSitterClient" } - internal var parser: Parser - internal var tree: Tree? - internal var languageQuery: Query? - - private var textProvider: ResolvingQueryCursor.TextProvider + internal var primaryLayer: TreeSitterLanguage + internal var layers: [LanguageLayer] = [] - /// The queue to do tree-sitter work on for large edits/queries - private let queue: DispatchQueue = DispatchQueue(label: "CodeEdit.CodeEditTextView.TreeSitter", - qos: .userInteractive) + internal var textProvider: ResolvingQueryCursor.TextProvider /// Used to ensure safe use of the shared tree-sitter tree state in different sync/async contexts. - private let semaphore: DispatchSemaphore = DispatchSemaphore(value: 1) + internal let semaphore: DispatchSemaphore = DispatchSemaphore(value: 1) + + internal enum Constants { + /// The maximum amount of limits a cursor can match during a query. + /// Used to ensure performance in large files, even though we generally limit the query to the visible range. + /// Neovim encountered this issue and uses 64 for their limit. Helix uses 256 due to issues with some + /// languages when using 64. + /// See: https://github.com/neovim/neovim/issues/14897 + /// And: https://github.com/helix-editor/helix/pull/4830 + static let treeSitterMatchLimit = 256 + } + + // MARK: - Init/Config /// Initializes the `TreeSitterClient` with the given parameters. /// - Parameters: /// - codeLanguage: The language to set up the parser with. /// - textProvider: The text provider callback to read any text. - public init(codeLanguage: CodeLanguage, textProvider: @escaping ResolvingQueryCursor.TextProvider) throws { - parser = Parser() - languageQuery = TreeSitterModel.shared.query(for: codeLanguage.id) - tree = nil - + public init(codeLanguage: CodeLanguage, textProvider: @escaping ResolvingQueryCursor.TextProvider) { self.textProvider = textProvider + self.primaryLayer = codeLanguage.id + setLanguage(codeLanguage: codeLanguage) + } + + /// Sets the primary language for the client. Will reset all layers, will not do any parsing work. + /// - Parameter codeLanguage: The new primary language. + public func setLanguage(codeLanguage: CodeLanguage) { + // Remove all trees and languages, everything needs to be re-parsed. + layers.removeAll() + + primaryLayer = codeLanguage.id + layers = [ + LanguageLayer( + id: codeLanguage.id, + parser: Parser(), + supportsInjections: codeLanguage.additionalHighlights?.contains("injections") ?? false, + tree: nil, + languageQuery: TreeSitterModel.shared.query(for: codeLanguage.id), + ranges: [] + ) + ] if let treeSitterLanguage = codeLanguage.language { - try parser.setLanguage(treeSitterLanguage) + try? layers[0].parser.setLanguage(treeSitterLanguage) } } - func setLanguage(codeLanguage: CodeLanguage) { - if let treeSitterLanguage = codeLanguage.language { - try? parser.setLanguage(treeSitterLanguage) - } + // MARK: - HighlightProviding + + /// Set up and parse the initial language tree and all injected layers. + func setUp(textView: HighlighterTextView) { + let readBlock = createReadBlock(textView: textView) + + layers[0].tree = createTree( + parser: layers[0].parser, + readBlock: readBlock + ) + + var layerSet = Set(arrayLiteral: layers[0]) + var touchedLayers = Set() + + var idx = 0 + while idx < layers.count { + updateInjectedLanguageLayers( + textView: textView, + layer: layers[idx], + layerSet: &layerSet, + touchedLayers: &touchedLayers, + readBlock: readBlock + ) - // Get rid of the current tree, it needs to be re-parsed. - tree = nil + idx += 1 + } } /// Notifies the highlighter of an edit and in exchange gets a set of indices that need to be re-highlighted. @@ -67,151 +112,179 @@ final class TreeSitterClient: HighlightProviding { /// - range: The range of the edit. /// - delta: The length of the edit, can be negative for deletions. /// - completion: The function to call with an `IndexSet` containing all Indices to invalidate. - func applyEdit(textView: HighlighterTextView, - range: NSRange, - delta: Int, - completion: @escaping ((IndexSet) -> Void)) { - guard let edit = InputEdit(range: range, delta: delta, oldEndPoint: .zero) else { - return - } + func applyEdit( + textView: HighlighterTextView, + range: NSRange, + delta: Int, + completion: @escaping ((IndexSet) -> Void) + ) { + guard let edit = InputEdit(range: range, delta: delta, oldEndPoint: .zero) else { return } + let readBlock = createReadBlock(textView: textView) + var rangeSet = IndexSet() - let readFunction: Parser.ReadBlock = { byteOffset, _ in - let limit = textView.documentRange.length - let location = byteOffset / 2 - let end = min(location + (1024), limit) - if location > end { - assertionFailure("location is greater than end") - return nil + // Helper data structure for finding existing layers in O(1) when adding injected layers + var layerSet = Set(minimumCapacity: layers.count) + // Tracks which layers were not touched at some point during the edit. Any layers left in this set + // after the second loop are removed. + var touchedLayers = Set(minimumCapacity: layers.count) + + // Loop through all layers, apply edits & find changed byte ranges. + for layerIdx in (0.. Void)) { - // Make sure we dont accidentally change the tree while we copy it. - self.semaphore.wait() - guard let tree = self.tree?.copy() else { - // In this case, we don't have a tree to work with already, so we need to make it and try to - // return some highlights - createTree(textView: textView) - - // This is slightly redundant but we're only doing one check. - guard let treeRetry = self.tree?.copy() else { - // Now we can return nothing for real. - self.semaphore.signal() - completion([]) - return + // Loop again and apply injections query, add any ranges not previously found + // using while loop because `updateInjectedLanguageLayers` can add to `layers` during the loop + var idx = 0 + while idx < layers.count { + let layer = layers[idx] + + if layer.supportsInjections { + rangeSet.formUnion( + updateInjectedLanguageLayers( + textView: textView, + layer: layer, + layerSet: &layerSet, + touchedLayers: &touchedLayers, + readBlock: readBlock + ) + ) } - self.semaphore.signal() - _queryColorsFor(tree: treeRetry, range: range, completion: completion) - return + idx += 1 } - self.semaphore.signal() + // Delete any layers that weren't touched at some point during the edit. + layers.removeAll(where: { touchedLayers.contains($0) }) - _queryColorsFor(tree: tree, range: range, completion: completion) + completion(rangeSet) } - private func _queryColorsFor(tree: Tree, - range: NSRange, - completion: @escaping (([HighlightRange]) -> Void)) { - guard let rootNode = tree.rootNode else { - completion([]) - return + /// Initiates a highlight query. + /// - Parameters: + /// - textView: The text view to use. + /// - range: The range to limit the highlights to. + /// - completion: Called when the query completes. + func queryHighlightsFor( + textView: HighlighterTextView, + range: NSRange, + completion: @escaping ((([HighlightRange]) -> Void)) + ) { + var highlights: [HighlightRange] = [] + var injectedSet = IndexSet(integersIn: range) + + for layer in layers where layer.id != primaryLayer { + // Query injected only if a layer's ranges intersects with `range` + for layerRange in layer.ranges { + if let rangeIntersection = range.intersection(layerRange) { + highlights.append(contentsOf: queryLayerHighlights( + layer: layer, + textView: textView, + range: rangeIntersection + )) + + injectedSet.remove(integersIn: rangeIntersection) + } + } } - // This needs to be on the main thread since we're going to use the `textProvider` in - // the `highlightsFromCursor` method, which uses the textView's text storage. - guard let cursor = self.languageQuery?.execute(node: rootNode, in: tree) else { - completion([]) - return + // Query primary for any ranges that weren't used in the injected layers. + for range in injectedSet.rangeView { + highlights.append(contentsOf: queryLayerHighlights( + layer: layers[0], + textView: textView, + range: NSRange(range) + )) } - cursor.setRange(range) - - let highlights = self.highlightsFromCursor(cursor: ResolvingQueryCursor(cursor: cursor)) completion(highlights) } - /// Creates a tree. - /// - Parameter textView: The text provider to use. - private func createTree(textView: HighlighterTextView) { - self.tree = self.parser.parse(textView.stringForRange(textView.documentRange) ?? "") - } - - /// Resolves a query cursor to the highlight ranges it contains. - /// **Must be called on the main thread** - /// - Parameter cursor: The cursor to resolve. - /// - Returns: Any highlight ranges contained in the cursor. - private func highlightsFromCursor(cursor: ResolvingQueryCursor) -> [HighlightRange] { - cursor.prepare(with: self.textProvider) - return cursor - .flatMap { $0.captures } - .map { HighlightRange(range: $0.range, capture: CaptureName.fromString($0.name ?? "")) } - } -} + // MARK: - Helpers -extension TreeSitterClient { - /// Applies the edit to the current `tree` and returns the old tree and a copy of the current tree with the - /// processed edit. + /// Attempts to create a language layer and load a highlights file. + /// Adds the layer to the `layers` array if successful. /// - Parameters: - /// - edit: The edit to apply. - /// - readBlock: The block to use to read text. - /// - Returns: (The old state, the new state). - private func calculateNewState(edit: InputEdit, - readBlock: @escaping Parser.ReadBlock) -> (Tree?, Tree?) { - guard let oldTree = self.tree else { - return (nil, nil) + /// - layerId: A language ID to add as a layer. + /// - readBlock: Completion called for efficient string lookup. + internal func addLanguageLayer( + layerId: TreeSitterLanguage, + readBlock: @escaping Parser.ReadBlock + ) -> LanguageLayer? { + guard let language = CodeLanguage.allLanguages.first(where: { $0.id == layerId }), + let parserLanguage = language.language + else { + return nil } - self.semaphore.wait() - // Apply the edit to the old tree - oldTree.edit(edit) - - self.tree = self.parser.parse(tree: oldTree, readBlock: readBlock) - - self.semaphore.signal() + let newLayer = LanguageLayer( + id: layerId, + parser: Parser(), + supportsInjections: language.additionalHighlights?.contains("injections") ?? false, + tree: nil, + languageQuery: TreeSitterModel.shared.query(for: layerId), + ranges: [] + ) + + do { + try newLayer.parser.setLanguage(parserLanguage) + } catch { + return nil + } - return (oldTree.copy(), self.tree?.copy()) + layers.append(newLayer) + return newLayer } - /// Calculates the changed byte ranges between two trees. + /// Creates a tree-sitter tree. /// - Parameters: - /// - lhs: The first (older) tree. - /// - rhs: The second (newer) tree. - /// - Returns: Any changed ranges. - private func changedByteRanges(_ lhs: Tree?, rhs: Tree?) -> [Range] { - switch (lhs, rhs) { - case (let t1?, let t2?): - return t1.changedRanges(from: t2).map({ $0.bytes }) - case (nil, let t2?): - let range = t2.rootNode?.byteRange - - return range.flatMap({ [$0] }) ?? [] - case (_, nil): - return [] + /// - parser: The parser object to use to parse text. + /// - readBlock: A callback for fetching blocks of text. + /// - Returns: A tree if it could be parsed. + internal func createTree(parser: Parser, readBlock: @escaping Parser.ReadBlock) -> Tree? { + return parser.parse(tree: nil, readBlock: readBlock) + } + + internal func createReadBlock(textView: HighlighterTextView) -> Parser.ReadBlock { + return { byteOffset, _ in + let limit = textView.documentRange.length + let location = byteOffset / 2 + let end = min(location + (1024), limit) + if location > end { + // Ignore and return nothing, tree-sitter's internal tree can be incorrect in some situations. + return nil + } + let range = NSRange(location..