Add Injected Language Support

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 {

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
+    }
+}

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))
+            }
+        }
+    }
+}

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)
+        )
+    }
+}

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

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])

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.

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 {

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<UInt32>] {
+        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<LanguageLayer>,
+        touchedLayers: inout Set<LanguageLayer>,
+        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
+    }
+}