Async TreeSitterClient

Package.swift

@@ -19,7 +19,7 @@ let package = Package(
         ),
         .package(
             url: "https://github.com/CodeEditApp/CodeEditLanguages.git",
-            exact: "0.1.13"
+            exact: "0.1.14"
         ),
         .package(
             url: "https://github.com/lukepistrol/SwiftLintPlugin",

Sources/CodeEditTextView/Controller/STTextViewController+Highlighter.swift

@@ -31,7 +31,7 @@ extension STTextViewController {
                 return self?.textView.textContentStorage?.textStorage?.mutableString.substring(with: range)
             }
 
-            provider = TreeSitterClient(codeLanguage: language, textProvider: textProvider)
+            provider = TreeSitterClient(textProvider: textProvider)
         }
 
         if let provider = provider {

Sources/CodeEditTextView/Enums/CaptureName.swift

@@ -33,7 +33,7 @@ public enum CaptureName: String, CaseIterable {
     /// - Parameter string: A string to get the capture name from
     /// - Returns: A `CaptureNames` case
     static func fromString(_ string: String?) -> CaptureName? {
-        allCases.first { $0.rawValue == string }
+        CaptureName(rawValue: string ?? "")
     }
 
     var alternate: CaptureName {

Sources/CodeEditTextView/Extensions/HighlighterTextView+createReadBlock.swift

@@ -0,0 +1,25 @@
+//
+//  HighlighterTextView+createReadBlock.swift
+//  
+//
+//  Created by Khan Winter on 5/20/23.
+//
+
+import Foundation
+import SwiftTreeSitter
+
+extension HighlighterTextView {
+    func createReadBlock() -> Parser.ReadBlock {
+        return { byteOffset, _ in
+            let limit = self.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..<end)
+            return self.stringForRange(range)?.data(using: String.nativeUTF16Encoding)
+        }
+    }
+}

Sources/CodeEditTextView/Extensions/Parser+createTree.swift

@@ -0,0 +1,20 @@
+//
+//  Parser+createTree.swift
+//  
+//
+//  Created by Khan Winter on 5/20/23.
+//
+
+import Foundation
+import SwiftTreeSitter
+
+extension Parser {
+    /// Creates a tree-sitter tree.
+    /// - Parameters:
+    ///   - 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(readBlock: @escaping Parser.ReadBlock) -> Tree? {
+        return parse(tree: nil, readBlock: readBlock)
+    }
+}

Sources/CodeEditTextView/Extensions/STTextView+/STTextView+HighlighterTextView.swift

@@ -0,0 +1,23 @@
+//
+//  STTextView+HighlighterTextView.swift
+//  CodeEditTextView
+//
+//  Created by Khan Winter on 6/2/23.
+//
+
+import Foundation
+import STTextView
+
+/// A default implementation for `STTextView` to be passed to `HighlightProviding` objects.
+extension STTextView: HighlighterTextView {
+    public var documentRange: NSRange {
+        return NSRange(
+            location: 0,
+            length: textContentStorage?.textStorage?.length ?? 0
+        )
+    }
+
+    public func stringForRange(_ nsRange: NSRange) -> String? {
+        return textContentStorage?.textStorage?.mutableString.substring(with: nsRange)
+    }
+}

Sources/CodeEditTextView/Extensions/STTextView+/STTextView+VisibleRange.swift

@@ -21,20 +21,43 @@ extension STTextView {
 
         // Get visible rect
         guard let bounds = enclosingScrollView?.documentVisibleRect else {
-            return textLayoutManager.documentRange.nsRange(using: textContentStorage)
+            return nil
         }
 
         // Calculate min & max points w/ a small amount of padding vertically.
-        let minPoint = CGPoint(x: bounds.minX,
-                               y: bounds.minY - 100)
-        let maxPoint = CGPoint(x: bounds.maxX,
-                               y: bounds.maxY + 100)
+        let minPoint = CGPoint(
+            x: bounds.minX,
+            y: bounds.minY - 200
+        )
+        let maxPoint = CGPoint(
+            x: bounds.maxX,
+            y: bounds.maxY + 200
+        )
 
         // Get text fragments for both the min and max points
-        guard let start = textLayoutManager.textLayoutFragment(for: minPoint)?.rangeInElement.location,
-              let end = textLayoutManager.textLayoutFragment(for: maxPoint)?.rangeInElement.endLocation else {
-            return textLayoutManager.documentRange.nsRange(using: textContentStorage)
+        guard let start = textLayoutManager.textLayoutFragment(for: minPoint)?.rangeInElement.location else {
+            return nil
         }
+
+        // End point can be tricky sometimes. If the document is smaller than the scroll view it can sometimes return
+        // nil for the `maxPoint` layout fragment. So we attempt to grab the last fragment.
+        var end: NSTextLocation?
+
+        if let endFragment = textLayoutManager.textLayoutFragment(for: maxPoint) {
+            end = endFragment.rangeInElement.location
+        } else {
+            textLayoutManager.ensureLayout(for: NSTextRange(location: textLayoutManager.documentRange.endLocation))
+            textLayoutManager.enumerateTextLayoutFragments(
+                from: textLayoutManager.documentRange.endLocation,
+                options: [.reverse, .ensuresLayout, .ensuresExtraLineFragment]
+            ) { layoutFragment in
+                end = layoutFragment.rangeInElement.endLocation
+                return false
+            }
+        }
+
+        guard let end else { return nil }
+
         guard start.compare(end) != .orderedDescending else {
             return NSTextRange(location: end, end: start)?.nsRange(using: textContentStorage)
         }

Sources/CodeEditTextView/Highlighting/HighlightProviding.swift

@@ -11,19 +11,17 @@ import STTextView
 import AppKit
 
 /// The protocol a class must conform to to be used for highlighting.
-public protocol HighlightProviding {
+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 at editor initialization.
-    func setUp(textView: HighlighterTextView)
-
-    /// Updates the highlighter's code language.
+    /// 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 setLanguage(codeLanguage: CodeLanguage)
+    func setUp(textView: HighlighterTextView, codeLanguage: CodeLanguage)
 
     /// 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.

Sources/CodeEditTextView/Highlighting/Highlighter.swift

@@ -67,11 +67,13 @@ class Highlighter: NSObject {
     ///   - textView: The text view to highlight.
     ///   - treeSitterClient: The tree-sitter client to handle tree updates and highlight queries.
     ///   - theme: The theme to use for highlights.
-    init(textView: STTextView,
-         highlightProvider: HighlightProviding?,
-         theme: EditorTheme,
-         attributeProvider: ThemeAttributesProviding,
-         language: CodeLanguage) {
+    init(
+        textView: STTextView,
+        highlightProvider: HighlightProviding?,
+        theme: EditorTheme,
+        attributeProvider: ThemeAttributesProviding,
+        language: CodeLanguage
+    ) {
         self.textView = textView
         self.highlightProvider = highlightProvider
         self.theme = theme
@@ -80,15 +82,13 @@ class Highlighter: NSObject {
 
         super.init()
 
-        highlightProvider?.setLanguage(codeLanguage: language)
-
         guard textView.textContentStorage?.textStorage != nil else {
             assertionFailure("Text view does not have a textStorage")
             return
         }
 
         textView.textContentStorage?.textStorage?.delegate = self
-        highlightProvider?.setUp(textView: textView)
+        highlightProvider?.setUp(textView: textView, codeLanguage: language)
 
         if let scrollView = textView.enclosingScrollView {
             NotificationCenter.default.addObserver(self,
@@ -113,16 +113,15 @@ class Highlighter: NSObject {
     /// Sets the language and causes a re-highlight of the entire text.
     /// - Parameter language: The language to update to.
     public func setLanguage(language: CodeLanguage) {
-        highlightProvider?.setLanguage(codeLanguage: language)
+        highlightProvider?.setUp(textView: textView, codeLanguage: language)
         invalidate()
     }
 
     /// Sets the highlight provider. Will cause a re-highlight of the entire text.
     /// - Parameter provider: The provider to use for future syntax highlights.
     public func setHighlightProvider(_ provider: HighlightProviding) {
         self.highlightProvider = provider
-        highlightProvider?.setLanguage(codeLanguage: language)
-        highlightProvider?.setUp(textView: textView)
+        highlightProvider?.setUp(textView: textView, codeLanguage: language)
         invalidate()
     }
 
@@ -172,14 +171,11 @@ private extension Highlighter {
             guard let attributeProvider = self?.attributeProvider,
                   let textView = self?.textView else { return }
 
-            // Mark these indices as not pending and valid
             self?.pendingSet.remove(integersIn: rangeToHighlight)
-            self?.validSet.formUnion(IndexSet(integersIn: rangeToHighlight))
-
-            // If this range does not exist in the visible set, we can exit.
-            if !(self?.visibleSet ?? .init()).contains(integersIn: rangeToHighlight) {
+            guard self?.visibleSet.intersects(integersIn: rangeToHighlight) ?? false else {
                 return
             }
+            self?.validSet.formUnion(IndexSet(integersIn: rangeToHighlight))
 
             // Try to create a text range for invalidating. If this fails we fail silently
             guard let textContentManager = textView.textLayoutManager.textContentManager,
@@ -251,7 +247,9 @@ private extension Highlighter {
 private extension Highlighter {
     /// Updates the view to highlight newly visible text when the textview is scrolled or bounds change.
     @objc func visibleTextChanged(_ notification: Notification) {
-        visibleSet = IndexSet(integersIn: textView.visibleTextRange ?? NSRange())
+        if let newVisibleRange = textView.visibleTextRange {
+            visibleSet = IndexSet(integersIn: newVisibleRange)
+        }
 
         // Any indices that are both *not* valid and in the visible text range should be invalidated
         let newlyInvalidSet = visibleSet.subtracting(validSet)

Sources/CodeEditTextView/Highlighting/HighlighterTextView.swift

@@ -10,21 +10,9 @@ import AppKit
 import STTextView
 
 /// The object `HighlightProviding` objects are given when asked for highlights.
-public protocol HighlighterTextView {
+public protocol HighlighterTextView: AnyObject {
     /// The entire range of the document.
     var documentRange: NSRange { get }
     /// A substring for the requested range.
     func stringForRange(_ nsRange: NSRange) -> String?
 }
-
-/// A default implementation for `STTextView` to be passed to `HighlightProviding` objects.
-extension STTextView: HighlighterTextView {
-    public var documentRange: NSRange {
-        return NSRange(location: 0,
-                       length: textContentStorage?.textStorage?.length ?? 0)
-    }
-
-    public func stringForRange(_ nsRange: NSRange) -> String? {
-        return textContentStorage?.textStorage?.mutableString.substring(with: nsRange)
-    }
-}