enkly-ai/pkg/renderkit/highlight.go

package renderkit

import (
	stdhtml "html"
	"strings"
	"unicode/utf8"
)

// A small, dependency-free syntax highlighter for fenced code blocks.
//
// It is a single-pass lexer, not a parser: it recognises comments, strings,
// numbers, keywords, built-in types, constants, function-call names and (for
// shell) $variables, and wraps each in <span class="tok-...">. That is
// enough to make code readable without pulling in a highlighting library or
// a CDN (this stack must work fully offline).
//
// Safety: input is gomarkdown's already-escaped block body. It is unescaped,
// tokenised, and every emitted piece is re-escaped, so no model-controlled
// text can ever become markup. The visible text is unchanged (the copy
// button still copies the original code).

// maxHighlightBytes skips highlighting for huge blocks; they are rare and
// the lexer is per-byte.
const maxHighlightBytes = 128 * 1024

type langSpec struct {
	keywords, types, consts map[string]bool
	lineComments            []string
	blockStart, blockEnd    string
	quotes                  string
	tripleQuotes            bool
	foldCase                bool // keywords are case-insensitive (SQL)
	shellVars               bool
	jsonKeys                bool
}

func wordSet(words string) map[string]bool {
	m := make(map[string]bool)
	for _, w := range strings.Fields(words) {
		m[w] = true
	}
	return m
}

var (
	goSpec = &langSpec{
		keywords: wordSet(`break case chan const continue default defer else fallthrough for func go goto if
			import interface map package range return select struct switch type var`),
		types: wordSet(`string int int8 int16 int32 int64 uint uint8 uint16 uint32 uint64 uintptr byte rune
			float32 float64 complex64 complex128 bool error any`),
		consts:       wordSet(`true false nil iota`),
		lineComments: []string{"//"}, blockStart: "/*", blockEnd: "*/",
		quotes: "\"'`",
	}
	jsSpec = &langSpec{
		keywords: wordSet(`async await break case catch class const continue debugger default delete do else
			export extends finally for from function if import in instanceof let new of return static super
			switch this throw try typeof var void while with yield interface type enum implements namespace
			declare readonly public private protected abstract as`),
		types:        wordSet(`string number boolean any unknown never object Promise Array Map Set`),
		consts:       wordSet(`true false null undefined NaN Infinity`),
		lineComments: []string{"//"}, blockStart: "/*", blockEnd: "*/",
		quotes: "\"'`",
	}
	pySpec = &langSpec{
		keywords: wordSet(`and as assert async await break class continue def del elif else except finally for
			from global if import in is lambda nonlocal not or pass raise return try while with yield match case`),
		types:        wordSet(`int str float bool list dict set tuple bytes object`),
		consts:       wordSet(`True False None`),
		lineComments: []string{"#"},
		quotes:       "\"'", tripleQuotes: true,
	}
	shSpec = &langSpec{
		keywords: wordSet(`if then else elif fi for while until do done case esac in function select time
			return exit break continue export local readonly declare unset source alias`),
		types: wordSet(`echo cd ls cat grep sed awk find git go make docker curl mkdir rm cp mv chmod chown
			touch sudo npm pip python gofmt`),
		lineComments: []string{"#"},
		quotes:       "\"'", shellVars: true,
	}
	jsonSpec = &langSpec{
		consts: wordSet(`true false null`),
		quotes: `"`, jsonKeys: true,
	}
	sqlSpec = &langSpec{
		keywords: wordSet(`select from where insert into values update set delete create table alter drop index
			join left right inner outer full cross on group by order having limit offset as and or not in is
			like between exists union all distinct case when then else end primary key foreign references
			default unique constraint with returning begin commit rollback extension if`),
		types: wordSet(`int integer bigint smallint serial bigserial text varchar char boolean bool date
			timestamp timestamptz numeric decimal float double real jsonb json uuid vector`),
		consts:       wordSet(`true false null`),
		lineComments: []string{"--"}, blockStart: "/*", blockEnd: "*/",
		quotes: "\"'", foldCase: true,
	}
	cSpec = &langSpec{
		keywords: wordSet(`auto break case const continue default do else enum extern for goto if inline
			register return sizeof static struct switch typedef union volatile while class namespace
			template typename public private protected virtual new delete using try catch throw`),
		types: wordSet(`void char short int long float double signed unsigned bool size_t int8_t int16_t
			int32_t int64_t uint8_t uint16_t uint32_t uint64_t string`),
		consts:       wordSet(`true false NULL nullptr`),
		lineComments: []string{"//"}, blockStart: "/*", blockEnd: "*/",
		quotes: "\"'",
	}
	javaSpec = &langSpec{
		keywords: wordSet(`abstract assert break case catch class const continue default do else enum extends
			final finally for if implements import instanceof interface native new package private protected
			public return static super switch synchronized this throw throws try void volatile while`),
		types:  wordSet(`boolean byte char short int long float double String Object List Map Set`),
		consts: wordSet(`true false null`), lineComments: []string{"//"},
		blockStart: "/*", blockEnd: "*/", quotes: "\"'",
	}
	// No ' in quotes: Rust lifetimes ('a) would swallow the rest of the line.
	rustSpec = &langSpec{
		keywords: wordSet(`as async await break const continue crate dyn else enum extern fn for if impl in
			let loop match mod move mut pub ref return self Self static struct super trait type unsafe use
			where while`),
		types: wordSet(`i8 i16 i32 i64 i128 isize u8 u16 u32 u64 u128 usize f32 f64 bool char str String
			Vec Option Result Box`),
		consts: wordSet(`true false None Some Ok Err`), lineComments: []string{"//"},
		blockStart: "/*", blockEnd: "*/", quotes: `"`,
	}

	langByName = map[string]*langSpec{
		"go": goSpec, "golang": goSpec,
		"js": jsSpec, "javascript": jsSpec, "jsx": jsSpec, "ts": jsSpec, "typescript": jsSpec, "tsx": jsSpec,
		"py": pySpec, "python": pySpec,
		"sh": shSpec, "bash": shSpec, "shell": shSpec, "zsh": shSpec, "console": shSpec,
		"json": jsonSpec,
		"sql":  sqlSpec, "postgres": sqlSpec, "psql": sqlSpec,
		"c": cSpec, "cpp": cSpec, "c++": cSpec, "cc": cSpec, "h": cSpec, "hpp": cSpec,
		"java": javaSpec,
		"rust": rustSpec, "rs": rustSpec,
	}
)

func isDigit(c byte) bool      { return c >= '0' && c <= '9' }
func isIdentStart(c byte) bool { return c == '_' || c == '$' || (c|0x20 >= 'a' && c|0x20 <= 'z') }
func isIdent(c byte) bool      { return isIdentStart(c) || isDigit(c) }

func spanned(b *strings.Builder, class, text string) {
	if class == "" {
		b.WriteString(stdhtml.EscapeString(text))
		return
	}
	b.WriteString(`<span class="tok-`)
	b.WriteString(class)
	b.WriteString(`">`)
	b.WriteString(stdhtml.EscapeString(text))
	b.WriteString(`</span>`)
}

// highlight takes a fenced block's language and its HTML-escaped body and
// returns the body with token spans. Unknown languages come back unchanged.
func highlight(lang, escapedBody string) string {
	lang = strings.ToLower(strings.TrimSpace(lang))
	if len(escapedBody) > maxHighlightBytes {
		return escapedBody
	}
	if lang == "diff" || lang == "patch" {
		return highlightDiff(stdhtml.UnescapeString(escapedBody))
	}
	spec := langByName[lang]
	if spec == nil {
		return escapedBody
	}
	return spec.lex(stdhtml.UnescapeString(escapedBody))
}

func highlightDiff(src string) string {
	var b strings.Builder
	lines := strings.Split(src, "\n")
	for i, ln := range lines {
		class := ""
		switch {
		case strings.HasPrefix(ln, "+++") || strings.HasPrefix(ln, "---"):
			class = "hunk"
		case strings.HasPrefix(ln, "@@"):
			class = "hunk"
		case strings.HasPrefix(ln, "+"):
			class = "add"
		case strings.HasPrefix(ln, "-"):
			class = "del"
		}
		spanned(&b, class, ln)
		if i < len(lines)-1 {
			b.WriteByte('\n')
		}
	}
	return b.String()
}

func (s *langSpec) lex(src string) string {
	var b strings.Builder
	b.Grow(len(src) + len(src)/2)
	i := 0
	for i < len(src) {
		c := src[i]
		rest := src[i:]

		if n := s.commentLen(src, i); n > 0 {
			spanned(&b, "cmt", rest[:n])
			i += n
			continue
		}

		if strings.IndexByte(s.quotes, c) >= 0 {
			n := s.stringLen(rest)
			class := "str"
			if s.jsonKeys {
				k := i + n
				for k < len(src) && (src[k] == ' ' || src[k] == '\t') {
					k++
				}
				if k < len(src) && src[k] == ':' {
					class = "key"
				}
			}
			spanned(&b, class, rest[:n])
			i += n
			continue
		}

		if isDigit(c) && (i == 0 || !isIdent(src[i-1])) {
			j := i + 1
			for j < len(src) && (isIdent(src[j]) || (src[j] == '.' && j+1 < len(src) && isDigit(src[j+1]))) {
				j++
			}
			spanned(&b, "num", src[i:j])
			i = j
			continue
		}

		if s.shellVars && c == '$' && i+1 < len(src) && (isIdentStart(src[i+1]) || src[i+1] == '{') {
			j := i + 1
			if src[j] == '{' {
				if k := strings.IndexAny(src[j:], "}\n"); k > 0 && src[j+k] == '}' {
					j += k + 1
				} else {
					j++
				}
			} else {
				for j < len(src) && isIdent(src[j]) {
					j++
				}
			}
			spanned(&b, "var", src[i:j])
			i = j
			continue
		}

		if isIdentStart(c) {
			j := i + 1
			for j < len(src) && isIdent(src[j]) {
				j++
			}
			spanned(&b, s.classify(src[i:j], src[j:]), src[i:j])
			i = j
			continue
		}

		_, size := utf8.DecodeRuneInString(rest)
		spanned(&b, "", rest[:size])
		i += size
	}
	return b.String()
}

func (s *langSpec) classify(word, after string) string {
	w := word
	if s.foldCase {
		w = strings.ToLower(word)
	}
	switch {
	case s.keywords[w]:
		return "kw"
	case s.types[w]:
		return "type"
	case s.consts[w]:
		return "const"
	}
	if len(s.types) > 0 || len(s.keywords) > 0 { // JSON has no functions
		for k := 0; k < len(after); k++ {
			if after[k] == ' ' || after[k] == '\t' {
				continue
			}
			if after[k] == '(' {
				return "fn"
			}
			break
		}
	}
	return ""
}

// commentLen returns the length of a comment starting at src[i], or 0.
func (s *langSpec) commentLen(src string, i int) int {
	rest := src[i:]
	for _, lc := range s.lineComments {
		if !strings.HasPrefix(rest, lc) {
			continue
		}
		// A "#" only starts a comment at a word boundary, so "a#b" and
		// "$#" aren't swallowed.
		if lc == "#" && i > 0 && src[i-1] != ' ' && src[i-1] != '\t' && src[i-1] != '\n' && src[i-1] != ';' {
			continue
		}
		if k := strings.IndexByte(rest, '\n'); k >= 0 {
			return k
		}
		return len(rest)
	}
	if s.blockStart != "" && strings.HasPrefix(rest, s.blockStart) {
		if k := strings.Index(rest[len(s.blockStart):], s.blockEnd); k >= 0 {
			return len(s.blockStart) + k + len(s.blockEnd)
		}
		return len(rest)
	}
	return 0
}

// stringLen returns the length of the string literal starting at rest[0].
// Unterminated ' and " literals end at the newline, so one stray quote
// can't colour the rest of the block.
func (s *langSpec) stringLen(rest string) int {
	q := rest[0]
	if s.tripleQuotes {
		if t := strings.Repeat(string(q), 3); strings.HasPrefix(rest, t) {
			if k := strings.Index(rest[3:], t); k >= 0 {
				return 3 + k + 3
			}
			return len(rest)
		}
	}
	if q == '`' { // raw / multi-line
		if k := strings.IndexByte(rest[1:], '`'); k >= 0 {
			return 1 + k + 1
		}
		return len(rest)
	}
	j := 1
	for j < len(rest) {
		switch rest[j] {
		case '\\':
			j += 2
			continue
		case q:
			return j + 1
		case '\n':
			return j
		}
		j++
	}
	if j > len(rest) {
		j = len(rest)
	}
	return j
}