internal/render/markdown/highlight/chroma.go

// Package highlight wires github.com/alecthomas/chroma/v2 into the
// tilstream markdown pipeline. The Highlighter type is safe for concurrent
// use and caches the style and lexer lookups, which showed up as a hot
// spot when re-rendering a site of ~1500 notes.
package highlight

import (
	"fmt"
	"io"
	"strings"
	"sync"

	"github.com/alecthomas/chroma/v2"
	"github.com/alecthomas/chroma/v2/formatters"
	chromaHTML "github.com/alecthomas/chroma/v2/formatters/html"
	"github.com/alecthomas/chroma/v2/lexers"
	"github.com/alecthomas/chroma/v2/styles"
)

// Options configure a Highlighter.
type Options struct {
	Style          string
	LineNumbers    bool
	HighlightLines [][2]int
	ClassPrefix    string
	Inline         bool
}

// Highlighter renders fenced code blocks to HTML via chroma.
type Highlighter struct {
	opts    Options
	style   *chroma.Style
	fmtr    *chromaHTML.Formatter
	lexers  sync.Map // map[string]chroma.Lexer
	classes []chromaHTML.Option
}

// NewHighlighter returns a Highlighter ready for concurrent use.
func NewHighlighter(o Options) (*Highlighter, error) {
	styleName := o.Style
	if styleName == "" {
		styleName = "monokai"
	}
	style := styles.Get(styleName)
	if style == nil {
		return nil, fmt.Errorf("highlight: unknown style %q", styleName)
	}
	var fopts []chromaHTML.Option
	if o.LineNumbers {
		fopts = append(fopts, chromaHTML.WithLineNumbers(true))
	}
	if o.ClassPrefix != "" {
		fopts = append(fopts, chromaHTML.ClassPrefix(o.ClassPrefix))
		fopts = append(fopts, chromaHTML.WithClasses(true))
	}
	if o.Inline {
		fopts = append(fopts, chromaHTML.InlineCode(true))
	}
	if len(o.HighlightLines) > 0 {
		fopts = append(fopts, chromaHTML.HighlightLines(o.HighlightLines))
	}
	return &Highlighter{
		opts:    o,
		style:   style,
		fmtr:    chromaHTML.New(fopts...),
		classes: fopts,
	}, nil
}

// Highlight writes HTML-ified code for src tagged with language lang.
// When lang is empty we fall back to chroma's analyser lexer.
func (h *Highlighter) Highlight(w io.Writer, lang string, src string) error {
	lexer := h.lexer(lang, src)
	it, err := lexer.Tokenise(nil, src)
	if err != nil {
		return fmt.Errorf("highlight: tokenise: %w", err)
	}
	return h.fmtr.Format(w, h.style, it)
}

// HighlightString is a convenience wrapper.
func (h *Highlighter) HighlightString(lang, src string) (string, error) {
	var b strings.Builder
	if err := h.Highlight(&b, lang, src); err != nil {
		return "", err
	}
	return b.String(), nil
}

func (h *Highlighter) lexer(lang, src string) chroma.Lexer {
	key := strings.ToLower(strings.TrimSpace(lang))
	if v, ok := h.lexers.Load(key); ok {
		return v.(chroma.Lexer)
	}
	var lx chroma.Lexer
	if key == "" {
		lx = lexers.Analyse(src)
	} else {
		lx = lexers.Get(key)
	}
	if lx == nil {
		lx = lexers.Fallback
	}
	lx = chroma.Coalesce(lx)
	h.lexers.Store(key, lx)
	return lx
}

// WriteCSS emits the class-based stylesheet that pairs with ClassPrefix.
// I dump this to a static asset at build time rather than inlining inside
// every page.
func (h *Highlighter) WriteCSS(w io.Writer) error {
	return h.fmtr.WriteCSS(w, h.style)
}

// RegisterTerminal lets tests verify the formatter registry is wired.
func RegisterTerminal() {
	_ = formatters.Register("noop", noopFormatter{})
}

type noopFormatter struct{}

func (noopFormatter) Format(io.Writer, *chroma.Style, chroma.Iterator) error { return nil }

// KnownLanguages returns a sorted list of chroma lexer names.
func KnownLanguages() []string {
	names := lexers.Names(false)
	return names
}

// MustStyle fetches a style by name or returns the default; used by tests
// to avoid error-handling boilerplate.
func MustStyle(name string) *chroma.Style {
	s := styles.Get(name)
	if s == nil {
		return styles.Fallback
	}
	return s
}