Skip to content

Commit a8fc8e6

Browse files
committed
perf(markdown): cache syntax highlighting results for code blocks
During streaming, the FastRenderer re-renders the full accumulated content on every incoming chunk. Code blocks that were already fully received were being re-tokenized through chroma's expensive regex engine on every render, causing O(n²) performance degradation. Add a package-level cache keyed by (lang, code) for syntax highlighting token results. Once a code block has been tokenized, subsequent renders reuse the cached tokens. The cache is cleared on theme changes via ResetStyles(). Streaming benchmark results (BenchmarkStreamingFastRenderer): Time: 3,311ms → ~350ms (~10x faster) Allocs: 26M → ~2.3M (~11x fewer) Memory: 2,497MB → ~906MB (~2.8x less) Assisted-By: docker-agent
1 parent d871092 commit a8fc8e6

3 files changed

Lines changed: 272 additions & 24 deletions

File tree

pkg/tui/components/markdown/fast_renderer.go

Lines changed: 70 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -173,10 +173,14 @@ func ResetStyles() {
173173
globalStylesOnce = sync.Once{}
174174
globalStylesMu.Unlock()
175175

176-
// Also clear chroma syntax highlighting cache
176+
// Also clear chroma syntax highlighting caches
177177
chromaStyleCacheMu.Lock()
178178
chromaStyleCache = make(map[chroma.TokenType]ansiStyle)
179179
chromaStyleCacheMu.Unlock()
180+
181+
syntaxHighlightCacheMu.Lock()
182+
syntaxHighlightCache.clear()
183+
syntaxHighlightCacheMu.Unlock()
180184
}
181185

182186
func getGlobalStyles() *cachedStyles {
@@ -2158,41 +2162,56 @@ type token struct {
21582162
style ansiStyle
21592163
}
21602164

2165+
// syntaxCacheKey builds a cache key for syntax highlighting results.
2166+
type syntaxCacheKey struct {
2167+
lang string
2168+
code string
2169+
}
2170+
21612171
var (
21622172
lexerCache = make(map[string]chroma.Lexer)
21632173
lexerCacheMu sync.RWMutex
21642174

21652175
// Cache for chroma token type to ansiStyle conversion (with code bg)
21662176
chromaStyleCache = make(map[chroma.TokenType]ansiStyle)
21672177
chromaStyleCacheMu sync.RWMutex
2178+
2179+
// Cache for syntax highlighting results to avoid re-tokenizing unchanged code blocks.
2180+
// Uses an LRU cache bounded to 128 entries to prevent unbounded memory growth
2181+
// in long-running TUI sessions with many unique code blocks.
2182+
syntaxHighlightCache = newLRUCache[syntaxCacheKey, []token](syntaxHighlightCacheSize)
2183+
syntaxHighlightCacheMu sync.RWMutex
2184+
)
2185+
2186+
const (
2187+
// syntaxHighlightCacheSize is the maximum number of syntax-highlighted code blocks
2188+
// to keep in cache. This bounds memory usage while retaining recently viewed blocks.
2189+
syntaxHighlightCacheSize = 128
21682190
)
21692191

21702192
func (p *parser) syntaxHighlight(code, lang string) []token {
2171-
var lexer chroma.Lexer
2172-
2173-
if lang != "" {
2174-
// Try cache first
2175-
lexerCacheMu.RLock()
2176-
lexer = lexerCache[lang]
2177-
lexerCacheMu.RUnlock()
2178-
2179-
if lexer == nil {
2180-
lexer = lexers.Get(lang)
2181-
if lexer == nil {
2182-
// Try with file extension
2183-
lexer = lexers.Match("file." + lang)
2184-
}
2185-
if lexer != nil {
2186-
lexer = chroma.Coalesce(lexer)
2187-
lexerCacheMu.Lock()
2188-
lexerCache[lang] = lexer
2189-
lexerCacheMu.Unlock()
2190-
}
2191-
}
2193+
cacheKey := syntaxCacheKey{lang: lang, code: code}
2194+
2195+
syntaxHighlightCacheMu.RLock()
2196+
if cached, ok := syntaxHighlightCache.get(cacheKey); ok {
2197+
syntaxHighlightCacheMu.RUnlock()
2198+
return cached
21922199
}
2200+
syntaxHighlightCacheMu.RUnlock()
2201+
2202+
tokens := p.doSyntaxHighlight(code, lang)
2203+
2204+
syntaxHighlightCacheMu.Lock()
2205+
syntaxHighlightCache.put(cacheKey, tokens)
2206+
syntaxHighlightCacheMu.Unlock()
2207+
2208+
return tokens
2209+
}
21932210

2211+
// doSyntaxHighlight performs the actual syntax highlighting without caching.
2212+
func (p *parser) doSyntaxHighlight(code, lang string) []token {
2213+
lexer := p.getLexer(lang)
21942214
if lexer == nil {
2195-
// No highlighting - return plain text with code background
21962215
return []token{{text: code, style: p.getCodeStyle(chroma.None)}}
21972216
}
21982217

@@ -2212,10 +2231,37 @@ func (p *parser) syntaxHighlight(code, lang string) []token {
22122231
style: p.getCodeStyle(tok.Type),
22132232
})
22142233
}
2215-
22162234
return tokens
22172235
}
22182236

2237+
// getLexer returns a cached chroma lexer for the given language, or nil if unknown.
2238+
func (p *parser) getLexer(lang string) chroma.Lexer {
2239+
if lang == "" {
2240+
return nil
2241+
}
2242+
2243+
lexerCacheMu.RLock()
2244+
lexer := lexerCache[lang]
2245+
lexerCacheMu.RUnlock()
2246+
if lexer != nil {
2247+
return lexer
2248+
}
2249+
2250+
lexer = lexers.Get(lang)
2251+
if lexer == nil {
2252+
lexer = lexers.Match("file." + lang)
2253+
}
2254+
if lexer == nil {
2255+
return nil
2256+
}
2257+
2258+
lexer = chroma.Coalesce(lexer)
2259+
lexerCacheMu.Lock()
2260+
lexerCache[lang] = lexer
2261+
lexerCacheMu.Unlock()
2262+
return lexer
2263+
}
2264+
22192265
func (p *parser) getCodeStyle(tokenType chroma.TokenType) ansiStyle {
22202266
chromaStyleCacheMu.RLock()
22212267
style, ok := chromaStyleCache[tokenType]
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package markdown
2+
3+
import "container/list"
4+
5+
// lruCache is a simple LRU (Least Recently Used) cache with a fixed maximum size.
6+
// It is NOT safe for concurrent use; callers must provide their own synchronization.
7+
type lruCache[K comparable, V any] struct {
8+
maxSize int
9+
items map[K]*list.Element
10+
order *list.List // front = most recently used
11+
}
12+
13+
type lruEntry[K comparable, V any] struct {
14+
key K
15+
value V
16+
}
17+
18+
// newLRUCache creates an LRU cache that holds at most maxSize entries.
19+
func newLRUCache[K comparable, V any](maxSize int) *lruCache[K, V] {
20+
return &lruCache[K, V]{
21+
maxSize: maxSize,
22+
items: make(map[K]*list.Element, maxSize),
23+
order: list.New(),
24+
}
25+
}
26+
27+
// get retrieves a value from the cache, promoting it to most-recently-used.
28+
// Returns the value and true if found, or the zero value and false otherwise.
29+
func (c *lruCache[K, V]) get(key K) (V, bool) {
30+
if elem, ok := c.items[key]; ok {
31+
c.order.MoveToFront(elem)
32+
return elem.Value.(*lruEntry[K, V]).value, true
33+
}
34+
var zero V
35+
return zero, false
36+
}
37+
38+
// put adds or updates a key-value pair in the cache.
39+
// If the cache is at capacity, the least recently used entry is evicted.
40+
func (c *lruCache[K, V]) put(key K, value V) {
41+
if elem, ok := c.items[key]; ok {
42+
// Update existing entry
43+
c.order.MoveToFront(elem)
44+
elem.Value.(*lruEntry[K, V]).value = value
45+
return
46+
}
47+
48+
// Evict if at capacity
49+
if c.order.Len() >= c.maxSize {
50+
c.evictOldest()
51+
}
52+
53+
entry := &lruEntry[K, V]{key: key, value: value}
54+
elem := c.order.PushFront(entry)
55+
c.items[key] = elem
56+
}
57+
58+
// clear removes all entries from the cache.
59+
func (c *lruCache[K, V]) clear() {
60+
c.items = make(map[K]*list.Element, c.maxSize)
61+
c.order.Init()
62+
}
63+
64+
// evictOldest removes the least recently used entry.
65+
func (c *lruCache[K, V]) evictOldest() {
66+
oldest := c.order.Back()
67+
if oldest == nil {
68+
return
69+
}
70+
c.order.Remove(oldest)
71+
delete(c.items, oldest.Value.(*lruEntry[K, V]).key)
72+
}
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
package markdown
2+
3+
import "testing"
4+
5+
func TestLRUCache_BasicGetPut(t *testing.T) {
6+
c := newLRUCache[string, int](3)
7+
8+
c.put("a", 1)
9+
c.put("b", 2)
10+
c.put("c", 3)
11+
12+
v, ok := c.get("a")
13+
if !ok || v != 1 {
14+
t.Fatalf("expected (1, true), got (%d, %v)", v, ok)
15+
}
16+
v, ok = c.get("b")
17+
if !ok || v != 2 {
18+
t.Fatalf("expected (2, true), got (%d, %v)", v, ok)
19+
}
20+
v, ok = c.get("c")
21+
if !ok || v != 3 {
22+
t.Fatalf("expected (3, true), got (%d, %v)", v, ok)
23+
}
24+
}
25+
26+
func TestLRUCache_Miss(t *testing.T) {
27+
c := newLRUCache[string, int](2)
28+
29+
_, ok := c.get("missing")
30+
if ok {
31+
t.Fatal("expected miss for non-existent key")
32+
}
33+
}
34+
35+
func TestLRUCache_Eviction(t *testing.T) {
36+
c := newLRUCache[string, int](2)
37+
38+
c.put("a", 1)
39+
c.put("b", 2)
40+
// Cache is full: [b, a] (b is most recent)
41+
42+
c.put("c", 3)
43+
// "a" should be evicted as least recently used: [c, b]
44+
45+
_, ok := c.get("a")
46+
if ok {
47+
t.Fatal("expected 'a' to be evicted")
48+
}
49+
50+
v, ok := c.get("b")
51+
if !ok || v != 2 {
52+
t.Fatalf("expected (2, true), got (%d, %v)", v, ok)
53+
}
54+
v, ok = c.get("c")
55+
if !ok || v != 3 {
56+
t.Fatalf("expected (3, true), got (%d, %v)", v, ok)
57+
}
58+
}
59+
60+
func TestLRUCache_GetPromotesEntry(t *testing.T) {
61+
c := newLRUCache[string, int](2)
62+
63+
c.put("a", 1)
64+
c.put("b", 2)
65+
// [b, a]
66+
67+
// Access "a" to promote it
68+
c.get("a")
69+
// Now [a, b]
70+
71+
// Add "c" - should evict "b" (now least recently used)
72+
c.put("c", 3)
73+
74+
_, ok := c.get("b")
75+
if ok {
76+
t.Fatal("expected 'b' to be evicted after 'a' was promoted")
77+
}
78+
79+
v, ok := c.get("a")
80+
if !ok || v != 1 {
81+
t.Fatalf("expected (1, true), got (%d, %v)", v, ok)
82+
}
83+
}
84+
85+
func TestLRUCache_UpdateExistingKey(t *testing.T) {
86+
c := newLRUCache[string, int](2)
87+
88+
c.put("a", 1)
89+
c.put("b", 2)
90+
91+
// Update "a"
92+
c.put("a", 10)
93+
94+
v, ok := c.get("a")
95+
if !ok || v != 10 {
96+
t.Fatalf("expected (10, true), got (%d, %v)", v, ok)
97+
}
98+
99+
// "a" was promoted by the update, so adding "c" should evict "b"
100+
c.put("c", 3)
101+
_, ok = c.get("b")
102+
if ok {
103+
t.Fatal("expected 'b' to be evicted")
104+
}
105+
}
106+
107+
func TestLRUCache_Clear(t *testing.T) {
108+
c := newLRUCache[string, int](3)
109+
110+
c.put("a", 1)
111+
c.put("b", 2)
112+
113+
c.clear()
114+
115+
_, ok := c.get("a")
116+
if ok {
117+
t.Fatal("expected empty cache after clear")
118+
}
119+
_, ok = c.get("b")
120+
if ok {
121+
t.Fatal("expected empty cache after clear")
122+
}
123+
124+
// Should work normally after clear
125+
c.put("c", 3)
126+
v, ok := c.get("c")
127+
if !ok || v != 3 {
128+
t.Fatalf("expected (3, true), got (%d, %v)", v, ok)
129+
}
130+
}

0 commit comments

Comments
 (0)