Rework Async tree-sitter Model, Fix Strong Ref Cycle (#225)

Package.resolved

@@ -5,17 +5,17 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/CodeEditApp/CodeEditLanguages.git",
       "state" : {
-        "revision" : "af29ab4a15474a0a38ef88ef65c20e58a0812e43",
-        "version" : "0.1.17"
+        "revision" : "620b463c88894741e20d4711c9435b33547de5d2",
+        "version" : "0.1.18"
       }
     },
     {
       "identity" : "codeedittextview",
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/CodeEditApp/CodeEditTextView.git",
       "state" : {
-        "revision" : "c867fed329b2b4ce91a13742e20626f50cf233bb",
-        "version" : "0.7.0"
+        "revision" : "6abce20f1827a3665a5159195157f592352e38b4",
+        "version" : "0.7.1"
       }
     },
     {
@@ -41,8 +41,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/apple/swift-collections.git",
       "state" : {
-        "revision" : "a902f1823a7ff3c9ab2fba0f992396b948eda307",
-        "version" : "1.0.5"
+        "revision" : "d029d9d39c87bed85b1c50adee7c41795261a192",
+        "version" : "1.0.6"
       }
     },
     {
@@ -57,10 +57,10 @@
     {
       "identity" : "swifttreesitter",
       "kind" : "remoteSourceControl",
-      "location" : "https://github.com/ChimeHQ/SwiftTreeSitter",
+      "location" : "https://github.com/ChimeHQ/SwiftTreeSitter.git",
       "state" : {
-        "revision" : "df25a52f72ebc5b50ae20d26d1363793408bb28b",
-        "version" : "0.7.1"
+        "revision" : "2599e95310b3159641469d8a21baf2d3d200e61f",
+        "version" : "0.8.0"
       }
     },
     {

Package.swift

@@ -22,7 +22,7 @@ let package = Package(
         // tree-sitter languages
         .package(
             url: "https://github.com/CodeEditApp/CodeEditLanguages.git",
-            exact: "0.1.17"
+            exact: "0.1.18"
         ),
         // SwiftLint
         .package(

Sources/CodeEditSourceEditor/CodeEditSourceEditor.swift

@@ -212,7 +212,7 @@ public struct CodeEditSourceEditor: NSViewControllerRepresentable {
     @MainActor
     public class Coordinator: NSObject {
         var parent: CodeEditSourceEditor
-        var controller: TextViewController?
+        weak var controller: TextViewController?
         var isUpdatingFromRepresentable: Bool = false
         var isUpdateFromTextView: Bool = false
 

Sources/CodeEditSourceEditor/Controller/TextViewController+HighlightBracket.swift

@@ -15,11 +15,11 @@ extension TextViewController {
         for range in textView.selectionManager.textSelections.map({ $0.range }) {
             if range.isEmpty,
                range.location > 0, // Range is not the beginning of the document
-               let preceedingCharacter = textView.textStorage.substring(
+               let precedingCharacter = textView.textStorage.substring(
                 from: NSRange(location: range.location - 1, length: 1) // The preceding character exists
                ) {
                 for pair in BracketPairs.allValues {
-                    if preceedingCharacter == pair.0 {
+                    if precedingCharacter == pair.0 {
                         // Walk forwards
                         if let characterIndex = findClosingPair(
                             pair.0,
@@ -34,7 +34,7 @@ extension TextViewController {
                                 highlightCharacter(range.location - 1)
                             }
                         }
-                    } else if preceedingCharacter == pair.1 && range.location - 1 > 0 {
+                    } else if precedingCharacter == pair.1 && range.location - 1 > 0 {
                         // Walk backwards
                         if let characterIndex = findClosingPair(
                             pair.1,

Sources/CodeEditSourceEditor/Controller/TextViewController+Highlighter.swift

@@ -32,11 +32,8 @@ extension TextViewController {
         if let highlightProvider = highlightProvider {
             provider = highlightProvider
         } else {
-            let textProvider: ResolvingQueryCursor.TextProvider = { [weak self] range, _ -> String? in
-                return self?.textView.textStorage.mutableString.substring(with: range)
-            }
-
-            provider = TreeSitterClient(textProvider: textProvider)
+            self.treeSitterClient = TreeSitterClient()
+            provider = self.treeSitterClient!
         }
 
         if let provider = provider {

Sources/CodeEditSourceEditor/Controller/TextViewController.swift

@@ -149,7 +149,12 @@ public class TextViewController: NSViewController {
         }
     }
 
-    internal var highlighter: Highlighter?
+    var highlighter: Highlighter?
+
+    /// The tree sitter client managed by the source editor.
+    ///
+    /// This will be `nil` if another highlighter provider is passed to the source editor.
+    internal(set) public var treeSitterClient: TreeSitterClient?
 
     private var fontCharWidth: CGFloat { (" " as NSString).size(withAttributes: [.font: font]).width }
 

Sources/CodeEditSourceEditor/Enums/CaptureName.swift

@@ -6,7 +6,7 @@
 //
 
 /// A collection of possible capture names for `tree-sitter` with their respected raw values.
-public enum CaptureName: String, CaseIterable {
+public enum CaptureName: String, CaseIterable, Sendable {
     case include
     case constructor
     case keyword

Sources/CodeEditSourceEditor/Extensions/NSRange+/NSRange+InputEdit.swift

@@ -6,27 +6,30 @@
 //
 
 import Foundation
+import CodeEditTextView
 import SwiftTreeSitter
 
 extension InputEdit {
-    init?(range: NSRange, delta: Int, oldEndPoint: Point) {
+    init?(range: NSRange, delta: Int, oldEndPoint: Point, textView: TextView) {
         let newEndLocation = NSMaxRange(range) + delta
 
         if newEndLocation < 0 {
             assertionFailure("Invalid range/delta")
             return nil
         }
 
-        // TODO: - Ask why Neon only uses .zero for these
-        let startPoint: Point = .zero
-        let newEndPoint: Point = .zero
+        let newRange = NSRange(location: range.location, length: range.length + delta)
+        let startPoint = textView.pointForLocation(newRange.location) ?? .zero
+        let newEndPoint = textView.pointForLocation(newEndLocation) ?? .zero
 
-        self.init(startByte: UInt32(range.location * 2),
-                  oldEndByte: UInt32(NSMaxRange(range) * 2),
-                  newEndByte: UInt32(newEndLocation * 2),
-                  startPoint: startPoint,
-                  oldEndPoint: oldEndPoint,
-                  newEndPoint: newEndPoint)
+        self.init(
+            startByte: UInt32(range.location * 2),
+            oldEndByte: UInt32(NSMaxRange(range) * 2),
+            newEndByte: UInt32(newEndLocation * 2),
+            startPoint: startPoint,
+            oldEndPoint: oldEndPoint,
+            newEndPoint: newEndPoint
+        )
     }
 }
 

Sources/CodeEditSourceEditor/Extensions/TextView+/TextView+Point.swift

@@ -0,0 +1,18 @@
+//
+//  TextView+Point.swift
+//  CodeEditSourceEditor
+//
+//  Created by Khan Winter on 1/18/24.
+//
+
+import Foundation
+import CodeEditTextView
+import SwiftTreeSitter
+
+extension TextView {
+    func pointForLocation(_ location: Int) -> Point? {
+        guard let linePosition = layoutManager.textLineForOffset(location) else { return nil }
+        let column = location - linePosition.range.location
+        return Point(row: linePosition.index, column: column)
+    }
+}

Sources/CodeEditSourceEditor/Extensions/HighlighterTextView+createReadBlock.swift → Sources/CodeEditSourceEditor/Extensions/TextView+/TextView+createReadBlock.swift

@@ -6,20 +6,27 @@
 //
 
 import Foundation
+import CodeEditTextView
 import SwiftTreeSitter
 
-extension HighlighterTextView {
+extension TextView {
     func createReadBlock() -> Parser.ReadBlock {
-        return { byteOffset, _ in
-            let limit = self.documentRange.length
+        return { [weak self] byteOffset, _ in
+            let limit = self?.documentRange.length ?? 0
             let location = byteOffset / 2
             let end = min(location + (1024), limit)
-            if location > end {
+            if location > end || self == nil {
                 // Ignore and return nothing, tree-sitter's internal tree can be incorrect in some situations.
                 return nil
             }
             let range = NSRange(location..<end)
-            return self.stringForRange(range)?.data(using: String.nativeUTF16Encoding)
+            return self?.stringForRange(range)?.data(using: String.nativeUTF16Encoding)
+        }
+    }
+
+    func createReadCallback() -> SwiftTreeSitter.Predicate.TextProvider {
+        return { [weak self] range, _ in
+            return self?.stringForRange(range)
         }
     }
 }

Sources/CodeEditSourceEditor/Highlighting/HighlightProviding.swift

@@ -6,42 +6,43 @@
 //
 
 import Foundation
+import CodeEditTextView
 import CodeEditLanguages
 import AppKit
 
 /// The protocol a class must conform to to be used for highlighting.
 public protocol HighlightProviding: AnyObject {
-    /// A unique identifier for the highlighter object.
-    /// Example: `"CodeEdit.TreeSitterHighlighter"`
-    /// - Note: This does not need to be *globally* unique, merely unique across all the highlighters used.
-    var identifier: String { get }
-
     /// Called once to set up the highlight provider with a data source and language.
     /// - Parameters:
     ///   - textView: The text view to use as a text source.
-    ///   - codeLanguage: The langugage that should be used by the highlighter.
-    func setUp(textView: HighlighterTextView, codeLanguage: CodeLanguage)
+    ///   - codeLanguage: The language that should be used by the highlighter.
+    func setUp(textView: TextView, codeLanguage: CodeLanguage)
+
+    /// Notifies the highlighter that an edit is going to happen in the given range.
+    /// - Parameters:
+    ///   - textView: The text view to use.
+    ///   - range: The range of the incoming edit.
+    func willApplyEdit(textView: TextView, range: NSRange)
 
     /// Notifies the highlighter of an edit and in exchange gets a set of indices that need to be re-highlighted.
     /// The returned `IndexSet` should include all indexes that need to be highlighted, including any inserted text.
     /// - Parameters:
-    ///   - textView:The text view to use.
+    ///   - textView: The text view to use.
     ///   - 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))
+    /// - Returns: an `IndexSet` containing all Indices to invalidate.
+    func applyEdit(textView: TextView, range: NSRange, delta: Int, completion: @escaping (IndexSet) -> Void)
 
     /// Queries the highlight provider for any ranges to apply highlights to. The highlight provider should return an
     /// array containing all ranges to highlight, and the capture type for the range. Any ranges or indexes
     /// excluded from the returned array will be treated as plain text and highlighted as such.
     /// - Parameters:
     ///   - textView: The text view to use.
-    ///   - range: The range to operate on.
-    ///   - completion: Function to call with all ranges to highlight
-    func queryHighlightsFor(textView: HighlighterTextView,
-                            range: NSRange,
-                            completion: @escaping (([HighlightRange]) -> Void))
+    ///   - range: The range to query.
+    /// - Returns: All highlight ranges for the queried ranges.
+    func queryHighlightsFor(textView: TextView, range: NSRange, completion: @escaping ([HighlightRange]) -> Void)
+}
+
+extension HighlightProviding {
+    public func willApplyEdit(textView: TextView, range: NSRange) { }
 }

Sources/CodeEditSourceEditor/Highlighting/HighlightRange.swift

@@ -7,13 +7,8 @@
 
 import Foundation
 
-/// This class represents a range to highlight, as well as the capture name for syntax coloring.
-public class HighlightRange {
-    init(range: NSRange, capture: CaptureName?) {
-        self.range = range
-        self.capture = capture
-    }
-
+/// This struct represents a range to highlight, as well as the capture name for syntax coloring.
+public struct HighlightRange: Sendable {
     let range: NSRange
     let capture: CaptureName?
 }

Sources/CodeEditSourceEditor/Highlighting/Highlighter+NSTextStorageDelegate.swift

@@ -0,0 +1,36 @@
+//
+//  Highlighter+NSTextStorageDelegate.swift
+//  CodeEditSourceEditor
+//
+//  Created by Khan Winter on 1/18/24.
+//
+
+import AppKit
+
+extension Highlighter: NSTextStorageDelegate {
+    /// Processes an edited range in the text.
+    /// Will query tree-sitter for any updated indices and re-highlight only the ranges that need it.
+    func textStorage(
+        _ textStorage: NSTextStorage,
+        didProcessEditing editedMask: NSTextStorageEditActions,
+        range editedRange: NSRange,
+        changeInLength delta: Int
+    ) {
+        // This method is called whenever attributes are updated, so to avoid re-highlighting the entire document
+        // each time an attribute is applied, we check to make sure this is in response to an edit.
+        guard editedMask.contains(.editedCharacters) else { return }
+
+        self.storageDidEdit(editedRange: editedRange, delta: delta)
+    }
+
+    func textStorage(
+        _ textStorage: NSTextStorage,
+        willProcessEditing editedMask: NSTextStorageEditActions,
+        range editedRange: NSRange,
+        changeInLength delta: Int
+    ) {
+        guard editedMask.contains(.editedCharacters) else { return }
+
+        self.storageWillEdit(editedRange: editedRange)
+    }
+}