Add Syntax Highlighting To Source Previews

Author: thecoolwinter

Sources/CodeEditSourceEditor/CodeSuggestion/Model/SuggestionViewModel.swift

@@ -16,6 +16,8 @@ final class SuggestionViewModel: ObservableObject {
 
     weak var delegate: CodeSuggestionDelegate?
 
+    private var syntaxHighlightedCache: [Int: NSAttributedString] = [:]
+
     func showCompletions(
         textView: TextViewController,
         delegate: CodeSuggestionDelegate,
@@ -56,6 +58,7 @@ final class SuggestionViewModel: ObservableObject {
                     }
 
                     self.items = completionItems.items
+                    self.syntaxHighlightedCache = [:]
                     showWindowOnParent(targetParentWindow, cursorRect)
                 }
             } catch {
@@ -109,4 +112,26 @@ final class SuggestionViewModel: ObservableObject {
         items.removeAll()
         activeTextView = nil
     }
+
+    func syntaxHighlights(forIndex index: Int) -> NSAttributedString? {
+        if let cached = syntaxHighlightedCache[index] {
+            return cached
+        }
+
+        if let sourcePreview = items[index].sourcePreview,
+           let theme = activeTextView?.theme,
+           let font = activeTextView?.font,
+           let language = activeTextView?.language {
+            let string = TreeSitterClient.quickHighlight(
+                string: sourcePreview,
+                theme: theme,
+                font: font,
+                language: language
+            )
+            syntaxHighlightedCache[index] = string
+            return string
+        }
+
+        return nil
+    }
 }

Sources/CodeEditSourceEditor/CodeSuggestion/TableView/CodeSuggestionPreviewView.swift

@@ -10,9 +10,9 @@ import SwiftUI
 final class CodeSuggestionPreviewView: NSVisualEffectView {
     private let spacing: CGFloat = 5
 
-    var sourcePreview: String? {
+    var sourcePreview: NSAttributedString? {
         didSet {
-            sourcePreviewLabel.stringValue = sourcePreview ?? ""
+            sourcePreviewLabel.attributedStringValue = sourcePreview ?? NSAttributedString(string: "")
             sourcePreviewLabel.isHidden = sourcePreview == nil
         }
     }

Sources/CodeEditSourceEditor/CodeSuggestion/TableView/SuggestionViewController.swift

@@ -333,7 +333,7 @@ extension SuggestionViewController: NSTableViewDataSource, NSTableViewDelegate {
             // Update our preview view
             let selectedItem = model.items[tableView.selectedRow]
 
-            previewView.sourcePreview = selectedItem.sourcePreview
+            previewView.sourcePreview = model.syntaxHighlights(forIndex: tableView.selectedRow)
             previewView.documentation = selectedItem.documentation
             previewView.pathComponents = selectedItem.pathComponents ?? []
             previewView.targetRange = selectedItem.targetPosition

Sources/CodeEditSourceEditor/TreeSitter/TreeSitterClient+Temporary.swift

@@ -5,10 +5,67 @@
 //  Created by Khan Winter on 7/24/25.
 //
 
-import Foundation
+import AppKit
+import SwiftTreeSitter
+import CodeEditLanguages
 
 extension TreeSitterClient {
-//    func temporaryHighlight() -> NSAttributedString {
-//
-//    }
+    static func quickHighlight(
+        string: String,
+        theme: EditorTheme,
+        font: NSFont,
+        language: CodeLanguage
+    ) -> NSAttributedString? {
+        guard let parserLanguage = language.language, let query = TreeSitterModel.shared.query(for: language.id) else {
+            return nil
+        }
+
+        do {
+            let parser = Parser()
+            try parser.setLanguage(parserLanguage)
+            guard let syntaxTree = parser.parse(string) else {
+                return nil
+            }
+            let queryCursor = query.execute(in: syntaxTree)
+            var ranges: [NSRange: Int] = [:]
+            let highlights: [HighlightRange] = queryCursor
+                .resolve(with: .init(string: string))
+                .flatMap { $0.captures }
+                .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
+                    }
+
+                    // Update the filter level to the current index since it's lower and a 'valid' capture
+                    ranges[range] = index
+
+                    return HighlightRange(range: range, capture: captureName)
+                }
+
+            var string = NSMutableAttributedString(string: string)
+
+            for highlight in highlights {
+                string.setAttributes(
+                    [
+                        .font: theme.fontFor(for: highlight.capture, from: font),
+                        .foregroundColor: theme.colorFor(highlight.capture)
+                    ],
+                    range: highlight.range
+                )
+            }
+
+            return string
+        } catch {
+            return nil
+        }
+    }
 }