Improve code highlighting performance in chat

Author: intitni

Core/Sources/ChatGPTChatTab/ChatPanel.swift

@@ -286,43 +286,39 @@ struct ChatHistory: View {
 struct ChatHistoryItem: View {
     let chat: StoreOf<Chat>
     let message: DisplayedChatMessage
-    @State var height: CGFloat = 0
-    @State var codeHighlightCacheController = CodeBlockHighlighterCacheController()
 
     var body: some View {
         let text = message.text
 
-        Group {
-            switch message.role {
-            case .user:
-                UserMessage(id: message.id, text: text, chat: chat)
-                    .listRowInsets(EdgeInsets(
-                        top: 0,
-                        leading: -8,
-                        bottom: 0,
-                        trailing: -8
-                    ))
-                    .padding(.vertical, 4)
-            case .assistant:
-                BotMessage(
-                    id: message.id,
-                    text: text,
-                    references: message.references,
-                    chat: chat
-                )
+        switch message.role {
+        case .user:
+            UserMessage(id: message.id, text: text, chat: chat)
                 .listRowInsets(EdgeInsets(
                     top: 0,
                     leading: -8,
                     bottom: 0,
                     trailing: -8
                 ))
                 .padding(.vertical, 4)
-            case .tool:
-                FunctionMessage(id: message.id, text: text)
-            case .ignored:
-                EmptyView()
-            }
-        }.environment(\.codeHighlightCacheController, codeHighlightCacheController)
+        case .assistant:
+            BotMessage(
+                id: message.id,
+                text: text,
+                references: message.references,
+                chat: chat
+            )
+            .listRowInsets(EdgeInsets(
+                top: 0,
+                leading: -8,
+                bottom: 0,
+                trailing: -8
+            ))
+            .padding(.vertical, 4)
+        case .tool:
+            FunctionMessage(id: message.id, text: text)
+        case .ignored:
+            EmptyView()
+        }
     }
 }
 

Core/Sources/ChatGPTChatTab/CodeBlockHighlighter.swift

@@ -1,67 +1,104 @@
+import Combine
+import DebounceFunction
 import Foundation
 import MarkdownUI
 import SharedUIComponents
 import SwiftUI
 
-final class CodeBlockHighlighterCacheController {
-    private var cache: [String: AttributedString] = [:]
-    
-    func get(_ key: String) -> AttributedString? {
-        cache[key]
-    }
-
-    func set(_ key: String, _ value: AttributedString) {
-        cache[key] = value
-    }
-}
+/// Use this instead of the built in ``CodeBlockView`` to highlight code blocks asynchronously,
+/// so that the UI doesn't freeze when rendering large code blocks.
+struct AsyncCodeBlockView: View {
+    class Storage: ObservableObject {
+        static let queue = DispatchQueue(
+            label: "chat-code-block-highlight",
+            qos: .userInteractive
+        )
 
-struct CodeHighlightCacheEnvironmentKey: EnvironmentKey {
-    static var defaultValue: CodeBlockHighlighterCacheController = .init()
-}
+        @Published var highlighted: AttributedString?
+        var debounceFunction: DebounceFunction<AsyncCodeBlockView>?
+        
+        init() {
+            self.debounceFunction = .init(duration: 0.5, block: { [weak self] view in
+                self?.highlight(for: view)
+            })
+        }
 
-extension EnvironmentValues {
-    var codeHighlightCacheController: CodeBlockHighlighterCacheController {
-        get { self[CodeHighlightCacheEnvironmentKey.self] }
-        set { self[CodeHighlightCacheEnvironmentKey.self] = newValue }
+        func highlight(debounce: Bool, for view: AsyncCodeBlockView) {
+            if debounce {
+                Task { await debounceFunction?(view) }
+            } else {
+                highlight(for: view)
+            }
+        }
+        
+        func highlight(for view: AsyncCodeBlockView) {
+            let content = view.content
+            let language = view.fenceInfo ?? ""
+            let brightMode = view.colorScheme != .dark
+            let font = view.font
+            Self.queue.async {
+                let content = highlightedCodeBlock(
+                    code: content,
+                    language: language,
+                    scenario: "chat",
+                    brightMode: brightMode,
+                    font: font
+                )
+                let string = AttributedString(content)
+                DispatchQueue.main.async {
+                    self.highlighted = string
+                }
+            }
+        }
     }
-}
 
-struct ChatCodeSyntaxHighlighter: CodeSyntaxHighlighter {
-    let brightMode: Bool
+    let fenceInfo: String?
+    let content: String
     let font: NSFont
-    let colorChange: Color?
-    var cacheController: CodeBlockHighlighterCacheController
 
-    init(
-        brightMode: Bool,
-        font: NSFont,
-        colorChange: Color?,
-        cacheController: CodeBlockHighlighterCacheController
-    ) {
-        self.brightMode = brightMode
+    @Environment(\.colorScheme) var colorScheme
+    @StateObject var storage = Storage()
+    @AppStorage(\.syncChatCodeHighlightTheme) var syncCodeHighlightTheme
+    @AppStorage(\.codeForegroundColorLight) var codeForegroundColorLight
+    @AppStorage(\.codeBackgroundColorLight) var codeBackgroundColorLight
+    @AppStorage(\.codeForegroundColorDark) var codeForegroundColorDark
+    @AppStorage(\.codeBackgroundColorDark) var codeBackgroundColorDark
+
+    init(fenceInfo: String?, content: String, font: NSFont) {
+        self.fenceInfo = fenceInfo
+        self.content = content.hasSuffix("\n") ? String(content.dropLast()) : content
         self.font = font
-        self.colorChange = colorChange
-        self.cacheController = cacheController
     }
 
-    func highlightCode(_ code: String, language: String?) -> Text {
-        let key = "\(language ?? "unknown") - \(code)"
-        if let text = cacheController.get(key) {
-            return Text(text)
+    var body: some View {
+        Group {
+            if let highlighted = storage.highlighted {
+                Text(highlighted)
+            } else {
+                Text(content).font(.init(font))
+            }
         }
-
-        let content = highlightedCodeBlock(
-            code: code,
-            language: language ?? "",
-            scenario: "chat",
-            brightMode: brightMode,
-            font: font
-        )
-        let string = AttributedString(content)
-        Task { @MainActor in
-            cacheController.set(key, string)
+        .onAppear {
+            storage.highlight(debounce: false, for: self)
+        }
+        .onChange(of: colorScheme) { _ in
+            storage.highlight(debounce: false, for: self)
+        }
+        .onChange(of: syncCodeHighlightTheme) { _ in
+            storage.highlight(debounce: true, for: self)
+        }
+        .onChange(of: codeForegroundColorLight) { _ in
+            storage.highlight(debounce: true, for: self)
+        }
+        .onChange(of: codeBackgroundColorLight) { _ in
+            storage.highlight(debounce: true, for: self)
+        }
+        .onChange(of: codeForegroundColorDark) { _ in
+            storage.highlight(debounce: true, for: self)
+        }
+        .onChange(of: codeBackgroundColorDark) { _ in
+            storage.highlight(debounce: true, for: self)
         }
-        return Text(string)
     }
 }