internal: ANSI SGR parser and line syntax tokenizer for 101-500

internal/ansi.go

@@ -0,0 +1,164 @@
+package internal
+
+import (
+	"log/slog"
+	"strconv"
+	"strings"
+)
+
+// AnsiColorKind says how an AnsiColor is given.
+type AnsiColorKind uint8
+
+const (
+	AnsiDefault AnsiColorKind = iota
+	AnsiIndexed
+	AnsiRGB
+)
+
+// AnsiColor is a terminal color: default, a 256-color index, or true color.
+type AnsiColor struct {
+	Kind    AnsiColorKind
+	Index   uint8
+	R, G, B uint8
+}
+
+// AnsiStyle is the SGR state a span is drawn with.
+type AnsiStyle struct {
+	FG, BG                                AnsiColor
+	Bold, Dim, Italic, Underline, Inverse bool
+	Strike                                bool
+}
+
+// AnsiSpan is a run of text in one style.
+type AnsiSpan struct {
+	Text  string
+	Style AnsiStyle
+}
+
+// ParseAnsi splits s into styled spans from st; unknown escapes stay literal.
+func ParseAnsi(s string, st AnsiStyle) ([]AnsiSpan, AnsiStyle) {
+	var spans []AnsiSpan
+	var b strings.Builder
+	flush := func() {
+		if b.Len() > 0 {
+			spans = append(spans, AnsiSpan{b.String(), st})
+			b.Reset()
+		}
+	}
+	for i := 0; i < len(s); {
+		if s[i] != 0x1b || i+1 >= len(s) || s[i+1] != '[' {
+			b.WriteByte(s[i])
+			i++
+			continue
+		}
+		j := i + 2
+		for j < len(s) && (s[j] < 0x40 || s[j] > 0x7e) {
+			j++
+		}
+		if j >= len(s) {
+			b.WriteString(s[i:])
+			break
+		}
+		if s[j] != 'm' {
+			slog.Debug("mujica: ansi unsupported sequence kept", "seq", strconv.Quote(s[i:j+1]))
+			b.WriteString(s[i : j+1])
+			i = j + 1
+			continue
+		}
+		flush()
+		st = applySGR(st, s[i+2:j])
+		i = j + 1
+	}
+	flush()
+	return spans, st
+}
+
+func applySGR(st AnsiStyle, params string) AnsiStyle {
+	if params == "" {
+		return AnsiStyle{}
+	}
+	codes := strings.Split(params, ";")
+	for k := 0; k < len(codes); k++ {
+		n, err := strconv.Atoi(codes[k])
+		if err != nil {
+			slog.Debug("mujica: ansi bad SGR parameter", "params", params)
+			return st
+		}
+		switch {
+		case n == 0:
+			st = AnsiStyle{}
+		case n == 1:
+			st.Bold = true
+		case n == 2:
+			st.Dim = true
+		case n == 3:
+			st.Italic = true
+		case n == 4:
+			st.Underline = true
+		case n == 7:
+			st.Inverse = true
+		case n == 9:
+			st.Strike = true
+		case n == 22:
+			st.Bold, st.Dim = false, false
+		case n == 23:
+			st.Italic = false
+		case n == 24:
+			st.Underline = false
+		case n == 27:
+			st.Inverse = false
+		case n == 29:
+			st.Strike = false
+		case n >= 30 && n <= 37:
+			st.FG = AnsiColor{Kind: AnsiIndexed, Index: uint8(n - 30)}
+		case n >= 90 && n <= 97:
+			st.FG = AnsiColor{Kind: AnsiIndexed, Index: uint8(n - 90 + 8)}
+		case n == 39:
+			st.FG = AnsiColor{}
+		case n >= 40 && n <= 47:
+			st.BG = AnsiColor{Kind: AnsiIndexed, Index: uint8(n - 40)}
+		case n >= 100 && n <= 107:
+			st.BG = AnsiColor{Kind: AnsiIndexed, Index: uint8(n - 100 + 8)}
+		case n == 49:
+			st.BG = AnsiColor{}
+		case n == 38 || n == 48:
+			c, used, ok := extendedColor(codes[k+1:])
+			if !ok {
+				slog.Debug("mujica: ansi bad extended color", "params", params)
+				return st
+			}
+			if n == 38 {
+				st.FG = c
+			} else {
+				st.BG = c
+			}
+			k += used
+		}
+	}
+	return st
+}
+
+// extendedColor reads "5;n" or "2;r;g;b" after 38 or 48.
+func extendedColor(rest []string) (AnsiColor, int, bool) {
+	num := func(i int) (uint8, bool) {
+		if i >= len(rest) {
+			return 0, false
+		}
+		v, err := strconv.Atoi(rest[i])
+		return uint8(v), err == nil && v >= 0 && v <= 255
+	}
+	if len(rest) == 0 {
+		return AnsiColor{}, 0, false
+	}
+	switch rest[0] {
+	case "5":
+		v, ok := num(1)
+		return AnsiColor{Kind: AnsiIndexed, Index: v}, 2, ok
+	case "2":
+		r, ok1 := num(1)
+		g, ok2 := num(2)
+		b, ok3 := num(3)
+		return AnsiColor{Kind: AnsiRGB, R: r, G: g, B: b}, 4, ok1 && ok2 && ok3
+	}
+	return AnsiColor{}, 0, false
+}

internal/ansi_test.go

@@ -0,0 +1,40 @@
+package internal
+
+import "testing"
+
+func TestParseAnsi(t *testing.T) {
+	spans, end := ParseAnsi("ok \x1b[1;31mfail\x1b[0m \x1b[38;2;10;20;30mrgb\x1b[38;5;200mix", AnsiStyle{})
+	want := []AnsiSpan{
+		{"ok ", AnsiStyle{}},
+		{"fail", AnsiStyle{Bold: true, FG: AnsiColor{Kind: AnsiIndexed, Index: 1}}},
+		{" ", AnsiStyle{}},
+		{"rgb", AnsiStyle{FG: AnsiColor{Kind: AnsiRGB, R: 10, G: 20, B: 30}}},
+		{"ix", AnsiStyle{FG: AnsiColor{Kind: AnsiIndexed, Index: 200}}},
+	}
+	if len(spans) != len(want) {
+		t.Fatalf("spans %+v", spans)
+	}
+	for i := range want {
+		if spans[i] != want[i] {
+			t.Errorf("span %d = %+v, want %+v", i, spans[i], want[i])
+		}
+	}
+	if end.FG.Index != 200 {
+		t.Errorf("end style %+v", end)
+	}
+}
+
+func TestParseAnsiCarriesStyleAcrossLines(t *testing.T) {
+	_, st := ParseAnsi("\x1b[92mgreen", AnsiStyle{})
+	spans, _ := ParseAnsi("still", st)
+	if spans[0].Style.FG != (AnsiColor{Kind: AnsiIndexed, Index: 10}) {
+		t.Fatalf("style lost: %+v", spans[0].Style)
+	}
+}
+
+func TestParseAnsiKeepsUnknownSequences(t *testing.T) {
+	spans, _ := ParseAnsi("a\x1b[2Kb", AnsiStyle{})
+	if len(spans) != 1 || spans[0].Text != "a\x1b[2Kb" {
+		t.Fatalf("spans %+v", spans)
+	}
+}

internal/highlight.go

@@ -0,0 +1,160 @@
+package internal
+
+import (
+	"strings"
+	"unicode"
+)
+
+// TokenKind classifies a piece of source code.
+type TokenKind uint8
+
+const (
+	TokPlain TokenKind = iota
+	TokKeyword
+	TokString
+	TokNumber
+	TokComment
+	TokType
+	TokPunct
+)
+
+// Token is a run of source text of one kind.
+type Token struct {
+	Kind TokenKind
+	Text string
+}
+
+// lang is a language's lexical rules.
+type lang struct {
+	keywords, types map[string]bool
+	line            string
+	blockOpen       string
+	blockClose      string
+	quotes          string
+}
+
+func words(s string) map[string]bool {
+	m := map[string]bool{}
+	for _, w := range strings.Fields(s) {
+		m[w] = true
+	}
+	return m
+}
+
+var langs = map[string]lang{
+	"go": {keywords: words("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 nil true false"),
+		types: words("bool byte complex64 complex128 error float32 float64 int int8 int16 int32 int64 rune string uint uint8 uint16 uint32 uint64 uintptr any"),
+		line:  "//", blockOpen: "/*", blockClose: "*/", quotes: "\"'`"},
+	"javascript": {keywords: words("async await break case catch class const continue default delete do else export extends finally for from function if import in instanceof let new return static super switch this throw try typeof var void while yield null undefined true false"),
+		types: words("string number boolean any unknown never object interface type enum"),
+		line:  "//", blockOpen: "/*", blockClose: "*/", quotes: "\"'`"},
+	"python": {keywords: words("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 None True False"),
+		types: words("int float str bool list dict set tuple bytes object"),
+		line:  "#", quotes: "\"'"},
+	"json":  {keywords: words("true false null"), quotes: "\""},
+	"shell": {keywords: words("if then else elif fi for in do done while case esac function return export local echo cd"), line: "#", quotes: "\"'"},
+	"sql": {keywords: words("select from where and or not insert into values update set delete create table index on join left right inner outer group by order having limit offset as null is in like distinct union all primary key foreign references"),
+		line: "--", blockOpen: "/*", blockClose: "*/", quotes: "'\""},
+}
+
+// Highlight tokenizes one line of code; inBlock carries an open block comment between lines.
+func Highlight(language, line string, inBlock bool) ([]Token, bool) {
+	if language == "typescript" {
+		language = "javascript"
+	}
+	l, ok := langs[language]
+	if !ok {
+		return []Token{{TokPlain, line}}, false
+	}
+	var out []Token
+	emit := func(k TokenKind, s string) {
+		if s == "" {
+			return
+		}
+		if n := len(out); n > 0 && out[n-1].Kind == k {
+			out[n-1].Text += s
+			return
+		}
+		out = append(out, Token{k, s})
+	}
+	i := 0
+	for i < len(line) {
+		rest := line[i:]
+		switch {
+		case inBlock:
+			end := strings.Index(rest, l.blockClose)
+			if end < 0 {
+				emit(TokComment, rest)
+				return out, true
+			}
+			emit(TokComment, rest[:end+len(l.blockClose)])
+			i += end + len(l.blockClose)
+			inBlock = false
+		case l.line != "" && strings.HasPrefix(rest, l.line):
+			emit(TokComment, rest)
+			return out, false
+		case l.blockOpen != "" && strings.HasPrefix(rest, l.blockOpen):
+			emit(TokComment, l.blockOpen)
+			i += len(l.blockOpen)
+			inBlock = true
+		case strings.ContainsRune(l.quotes, rune(rest[0])):
+			n := quotedLen(rest)
+			emit(TokString, rest[:n])
+			i += n
+		case rest[0] >= '0' && rest[0] <= '9':
+			n := spanLen(rest, func(r rune) bool { return unicode.IsDigit(r) || unicode.IsLetter(r) || r == '.' || r == '_' })
+			emit(TokNumber, rest[:n])
+			i += n
+		case isIdent(rune(rest[0])):
+			n := spanLen(rest, func(r rune) bool { return isIdent(r) || unicode.IsDigit(r) })
+			w := rest[:n]
+			switch {
+			case l.keywords[w] || l.keywords[strings.ToLower(w)] && language == "sql":
+				emit(TokKeyword, w)
+			case l.types[w]:
+				emit(TokType, w)
+			default:
+				emit(TokPlain, w)
+			}
+			i += n
+		case strings.ContainsRune("{}()[];,.:=+-*/<>!&|%^~?", rune(rest[0])):
+			emit(TokPunct, rest[:1])
+			i++
+		default:
+			n := spanLen(rest, func(r rune) bool { return r < 128 && unicode.IsSpace(r) })
+			if n == 0 {
+				n = len(string([]rune(rest)[0]))
+			}
+			emit(TokPlain, rest[:n])
+			i += n
+		}
+	}
+	return out, inBlock
+}
+
+func isIdent(r rune) bool { return r == '_' || r == '$' || unicode.IsLetter(r) }
+
+func spanLen(s string, ok func(rune) bool) int {
+	for i, r := range s {
+		if !ok(r) {
+			return i
+		}
+	}
+	return len(s)
+}
+
+// quotedLen measures a string literal at s[0], honoring backslash escapes.
+func quotedLen(s string) int {
+	q := s[0]
+	for i := 1; i < len(s); i++ {
+		switch s[i] {
+		case '\\':
+			if q != '`' {
+				i++
+			}
+		case q:
+			return i + 1
+		}
+	}
+	return len(s)
+}

internal/highlight_test.go

@@ -0,0 +1,65 @@
+package internal
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestHighlightGo(t *testing.T) {
+	toks, open := Highlight("go", `func f(s string) int { return 42 // done`, false)
+	if open {
+		t.Fatal("block left open")
+	}
+	got := map[string]TokenKind{}
+	for _, tk := range toks {
+		got[strings.TrimSpace(tk.Text)] = tk.Kind
+	}
+	for text, want := range map[string]TokenKind{"func": TokKeyword, "string": TokType, "42": TokNumber, "// done": TokComment} {
+		if got[text] != want {
+			t.Errorf("%q kind %d, want %d (tokens %+v)", text, got[text], want, toks)
+		}
+	}
+}
+
+func TestHighlightStringsAndEscapes(t *testing.T) {
+	toks, _ := Highlight("javascript", `const s = "a\"b"; x`, false)
+	var str string
+	for _, tk := range toks {
+		if tk.Kind == TokString {
+			str = tk.Text
+		}
+	}
+	if str != `"a\"b"` {
+		t.Fatalf("string token %q", str)
+	}
+}
+
+func TestHighlightBlockCommentAcrossLines(t *testing.T) {
+	_, open := Highlight("go", "x /* start", false)
+	if !open {
+		t.Fatal("block not open")
+	}
+	toks, open := Highlight("go", "still */ y", true)
+	if open || toks[0].Kind != TokComment || toks[0].Text != "still */" {
+		t.Fatalf("tokens %+v open %v", toks, open)
+	}
+}
+
+func TestHighlightUnknownLanguageIsPlain(t *testing.T) {
+	toks, _ := Highlight("cobol", "MOVE A TO B", false)
+	if len(toks) != 1 || toks[0].Kind != TokPlain {
+		t.Fatalf("tokens %+v", toks)
+	}
+}
+
+func TestHighlightKeepsText(t *testing.T) {
+	line := `select * from "users" where id = 7 -- note`
+	toks, _ := Highlight("sql", line, false)
+	var b strings.Builder
+	for _, tk := range toks {
+		b.WriteString(tk.Text)
+	}
+	if b.String() != line {
+		t.Fatalf("round trip %q", b.String())
+	}
+}