Improve performance of code highlighting in chat

Author: intitni

Core/Sources/ChatGPTChatTab/ChatPanel.swift

@@ -277,39 +277,52 @@ struct ChatHistory: View {
     var body: some View {
         WithViewStore(chat, observe: \.history) { viewStore in
             ForEach(viewStore.state, id: \.id) { message in
-                let text = message.text
-
-                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
-                    )
+                ChatHistoryItem(chat: chat, message: message).id(message.id)
+            }
+        }
+    }
+}
+
+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 .tool:
-                    FunctionMessage(id: message.id, text: text)
-                case .ignored:
-                    EmptyView()
-                }
+            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()
             }
-        }
+        }.environment(\.codeHighlightCacheController, codeHighlightCacheController)
     }
 }
 
@@ -563,29 +576,6 @@ struct ChatPanel_EmptyChat_Preview: PreviewProvider {
     }
 }
 
-struct ChatCodeSyntaxHighlighter: CodeSyntaxHighlighter {
-    let brightMode: Bool
-    let font: NSFont
-    let colorChange: Color?
-
-    init(brightMode: Bool, font: NSFont, colorChange: Color?) {
-        self.brightMode = brightMode
-        self.font = font
-        self.colorChange = colorChange
-    }
-
-    func highlightCode(_ content: String, language: String?) -> Text {
-        let content = highlightedCodeBlock(
-            code: content,
-            language: language ?? "",
-            scenario: "chat",
-            brightMode: brightMode,
-            font: font
-        )
-        return Text(AttributedString(content))
-    }
-}
-
 struct ChatPanel_InputText_Preview: PreviewProvider {
     static var previews: some View {
         ChatPanel(chat: .init(

Core/Sources/ChatGPTChatTab/CodeBlockHighlighter.swift

@@ -0,0 +1,67 @@
+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
+    }
+}
+
+struct CodeHighlightCacheEnvironmentKey: EnvironmentKey {
+    static var defaultValue: CodeBlockHighlighterCacheController = .init()
+}
+
+extension EnvironmentValues {
+    var codeHighlightCacheController: CodeBlockHighlighterCacheController {
+        get { self[CodeHighlightCacheEnvironmentKey.self] }
+        set { self[CodeHighlightCacheEnvironmentKey.self] = newValue }
+    }
+}
+
+struct ChatCodeSyntaxHighlighter: CodeSyntaxHighlighter {
+    let brightMode: Bool
+    let font: NSFont
+    let colorChange: Color?
+    var cacheController: CodeBlockHighlighterCacheController
+
+    init(
+        brightMode: Bool,
+        font: NSFont,
+        colorChange: Color?,
+        cacheController: CodeBlockHighlighterCacheController
+    ) {
+        self.brightMode = brightMode
+        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)
+        }
+
+        let content = highlightedCodeBlock(
+            code: code,
+            language: language ?? "",
+            scenario: "chat",
+            brightMode: brightMode,
+            font: font
+        )
+        let string = AttributedString(content)
+        Task { @MainActor in
+            cacheController.set(key, string)
+        }
+        return Text(string)
+    }
+}
+

Core/Sources/ChatGPTChatTab/Styles.swift

@@ -83,6 +83,7 @@ struct ThemedMarkdownText: View {
     @AppStorage(\.chatFontSize) var chatFontSize
     @AppStorage(\.chatCodeFont) var chatCodeFont
     @Environment(\.colorScheme) var colorScheme
+    @Environment(\.codeHighlightCacheController) var codeHighlightCacheController
 
     let text: String
 
@@ -125,7 +126,8 @@ struct ThemedMarkdownText: View {
                     font: chatCodeFont.value.nsFont,
                     colorChange: colorScheme == .dark
                         ? codeForegroundColorDark.value?.swiftUIColor
-                        : codeForegroundColorLight.value?.swiftUIColor
+                        : codeForegroundColorLight.value?.swiftUIColor,
+                    cacheController: codeHighlightCacheController
                 )
             )
     }