feat: introduce a LRU compiled style cache for the HTML formatter (#938)

alecthomas · web-flow · commit 6dd9f269efaa · 2024-02-27T17:05:08.000+11:00
```
🐚 ~/dev/chroma $ benchcmp before.txt after.txt
benchmark                    old ns/op     new ns/op     delta
BenchmarkHTMLFormatter-8     160560        77797         -51.55%

benchmark                    old allocs     new allocs     delta
BenchmarkHTMLFormatter-8     1267           459            -63.77%

benchmark                    old bytes     new bytes     delta
BenchmarkHTMLFormatter-8     52568         25067         -52.32%
```
diff --git a/formatters/html/html.go b/formatters/html/html.go
@@ -7,6 +7,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 
 	"github.com/alecthomas/chroma/v2"
 )
@@ -133,6 +134,7 @@ func New(options ...Option) *Formatter {
 		baseLineNumber: 1,
 		preWrapper:     defaultPreWrapper,
 	}
+	f.styleCache = newStyleCache(f)
 	for _, option := range options {
 		option(f)
 	}
@@ -189,6 +191,7 @@ var (
 
 // Formatter that generates HTML.
 type Formatter struct {
+	styleCache            *styleCache
 	standalone            bool
 	prefix                string
 	Classes               bool // Exported field to detect when classes are being used
@@ -221,12 +224,7 @@ func (f *Formatter) Format(w io.Writer, style *chroma.Style, iterator chroma.Ite
 //
 // OTOH we need to be super careful about correct escaping...
 func (f *Formatter) writeHTML(w io.Writer, style *chroma.Style, tokens []chroma.Token) (err error) { // nolint: gocyclo
-	css := f.styleToCSS(style)
-	if !f.Classes {
-		for t, style := range css {
-			css[t] = compressStyle(style)
-		}
-	}
+	css := f.styleCache.get(style)
 	if f.standalone {
 		fmt.Fprint(w, "<html>\n")
 		if f.Classes {
@@ -420,7 +418,7 @@ func (f *Formatter) tabWidthStyle() string {
 
 // WriteCSS writes CSS style definitions (without any surrounding HTML).
 func (f *Formatter) WriteCSS(w io.Writer, style *chroma.Style) error {
-	css := f.styleToCSS(style)
+	css := f.styleCache.get(style)
 	// Special-case background as it is mapped to the outer ".chroma" class.
 	if _, err := fmt.Fprintf(w, "/* %s */ .%sbg { %s }\n", chroma.Background, f.prefix, css[chroma.Background]); err != nil {
 		return err
@@ -563,3 +561,60 @@ func compressStyle(s string) string {
 	}
 	return strings.Join(out, ";")
 }
+
+const styleCacheLimit = 16
+
+type styleCacheEntry struct {
+	style *chroma.Style
+	cache map[chroma.TokenType]string
+}
+
+type styleCache struct {
+	mu sync.Mutex
+	// LRU cache of compiled (and possibly compressed) styles. This is a slice
+	// because the cache size is small, and a slice is sufficiently fast for
+	// small N.
+	cache []styleCacheEntry
+	f     *Formatter
+}
+
+func newStyleCache(f *Formatter) *styleCache {
+	return &styleCache{f: f}
+}
+
+func (l *styleCache) get(style *chroma.Style) map[chroma.TokenType]string {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	// Look for an existing entry.
+	for i := len(l.cache) - 1; i >= 0; i-- {
+		entry := l.cache[i]
+		if entry.style == style {
+			// Top of the cache, no need to adjust the order.
+			if i == len(l.cache)-1 {
+				return entry.cache
+			}
+			// Move this entry to the end of the LRU
+			copy(l.cache[i:], l.cache[i+1:])
+			l.cache[len(l.cache)-1] = entry
+			return entry.cache
+		}
+	}
+
+	// No entry, create one.
+	cached := l.f.styleToCSS(style)
+	if !l.f.Classes {
+		for t, style := range cached {
+			cached[t] = compressStyle(style)
+		}
+	}
+	for t, style := range cached {
+		cached[t] = compressStyle(style)
+	}
+	// Evict the oldest entry.
+	if len(l.cache) >= styleCacheLimit {
+		l.cache = l.cache[0:copy(l.cache, l.cache[1:])]
+	}
+	l.cache = append(l.cache, styleCacheEntry{style: style, cache: cached})
+	return cached
+}
diff --git a/formatters/html/html_test.go b/formatters/html/html_test.go
@@ -222,7 +222,7 @@ func TestTableLinkeableLineNumbers(t *testing.T) {
 
 	assert.Contains(t, buf.String(), `id="line1"><a class="lnlinks" href="#line1">1</a>`)
 	assert.Contains(t, buf.String(), `id="line5"><a class="lnlinks" href="#line5">5</a>`)
-	assert.Contains(t, buf.String(), `/* LineLink */ .chroma .lnlinks { outline: none; text-decoration: none; color: inherit }`, buf.String())
+	assert.Contains(t, buf.String(), `/* LineLink */ .chroma .lnlinks { outline:none;text-decoration:none;color:inherit }`, buf.String())
 }
 
 func TestTableLineNumberSpacing(t *testing.T) {
@@ -351,12 +351,25 @@ func TestReconfigureOptions(t *testing.T) {
 }
 
 func TestWriteCssWithAllClasses(t *testing.T) {
-	formatter := New()
-	formatter.allClasses = true
+	formatter := New(WithAllClasses(true))
 
 	var buf bytes.Buffer
 	err := formatter.WriteCSS(&buf, styles.Fallback)
 
 	assert.NoError(t, err)
 	assert.NotContains(t, buf.String(), ".chroma . {", "Generated css doesn't contain invalid css")
 }
+
+func TestStyleCache(t *testing.T) {
+	f := New()
+
+	assert.True(t, len(styles.Registry) > styleCacheLimit)
+
+	for _, style := range styles.Registry {
+		var buf bytes.Buffer
+		err := f.WriteCSS(&buf, style)
+		assert.NoError(t, err)
+	}
+
+	assert.Equal(t, styleCacheLimit, len(f.styleCache.cache))
+}
