From ff23bd7cdfee473f904281409ebd7b6dd528b6da Mon Sep 17 00:00:00 2001 From: Willow Sparks Date: Mon, 9 Nov 2026 18:26:47 +0000 Subject: [PATCH] feat(treesitter): drive grammars and queries from a registry Replace the hardcoded filetype and grammar maps with one Definition per linked grammar: grammar id, micro filetypes, ordered query files, and language constructor. Query layers concatenate in order, so TSX and C++ reuse their base language highlights without a query preprocessor. Remove the Go keyword refinement so capture mapping stays language-neutral, and add table-driven span tests plus incremental differential tests for Go and Python. --- treesitter/internal/tsprovider/provider.go | 146 ++++++----- .../internal/tsprovider/provider_test.go | 239 ++++++++++++++---- 2 files changed, 269 insertions(+), 116 deletions(-) diff --git a/treesitter/internal/tsprovider/provider.go b/treesitter/internal/tsprovider/provider.go index 796a88e5..7fc39407 100644 --- a/treesitter/internal/tsprovider/provider.go +++ b/treesitter/internal/tsprovider/provider.go @@ -13,9 +13,23 @@ import ( "path/filepath" "sort" "strings" + "unsafe" + tslua "github.com/tree-sitter-grammars/tree-sitter-lua/bindings/go" + tstoml "github.com/tree-sitter-grammars/tree-sitter-toml/bindings/go" + tsyaml "github.com/tree-sitter-grammars/tree-sitter-yaml/bindings/go" sitter "github.com/tree-sitter/go-tree-sitter" + tsbash "github.com/tree-sitter/tree-sitter-bash/bindings/go" + tsc "github.com/tree-sitter/tree-sitter-c/bindings/go" + tscpp "github.com/tree-sitter/tree-sitter-cpp/bindings/go" + tscss "github.com/tree-sitter/tree-sitter-css/bindings/go" tsgo "github.com/tree-sitter/tree-sitter-go/bindings/go" + tshtml "github.com/tree-sitter/tree-sitter-html/bindings/go" + tsjavascript "github.com/tree-sitter/tree-sitter-javascript/bindings/go" + tsjson "github.com/tree-sitter/tree-sitter-json/bindings/go" + tspython "github.com/tree-sitter/tree-sitter-python/bindings/go" + tsrust "github.com/tree-sitter/tree-sitter-rust/bindings/go" + tstypescript "github.com/tree-sitter/tree-sitter-typescript/bindings/go" ) // Span is one half-open, single-line byte range. Line is 0-based; Start and @@ -27,17 +41,51 @@ type Span struct { Group string } -// filetypeGrammar names the tree-sitter grammar used for a micro filetype. -var filetypeGrammar = map[string]string{ - "go": "go", +// Definition describes one linked grammar. Queries names files under +// // and is concatenated in order, so a dialect can layer its +// own patterns over a base language. Filetypes are the micro filetypes from +// pkg/filetype/table.go that this grammar serves. +type Definition struct { + ID string + Filetypes []string + Queries []string + Language func() *sitter.Language } -// grammarLanguages links a grammar id to its compiled language. Adding a -// grammar means adding a module to go.mod and one entry here. -var grammarLanguages = map[string]func() *sitter.Language{ - "go": func() *sitter.Language { - return sitter.NewLanguage(tsgo.Language()) - }, +// definitions is the linked grammar registry. Adding a language means adding +// its module to go.mod, one entry here, and its query directory. +var definitions = []Definition{ + {ID: "go", Filetypes: []string{"go"}, Queries: []string{"highlights.scm"}, Language: language(tsgo.Language)}, + {ID: "python", Filetypes: []string{"python"}, Queries: []string{"highlights.scm"}, Language: language(tspython.Language)}, + {ID: "javascript", Filetypes: []string{"javascript"}, Queries: []string{"highlights.scm"}, Language: language(tsjavascript.Language)}, + // micro's "typescript" filetype covers both .ts and .tsx, so the TSX + // grammar serves it. TSX queries are layered over the JavaScript base. + {ID: "tsx", Filetypes: []string{"typescript"}, Queries: []string{"javascript.scm", "highlights.scm"}, Language: language(tstypescript.LanguageTSX)}, + {ID: "rust", Filetypes: []string{"rust"}, Queries: []string{"highlights.scm"}, Language: language(tsrust.Language)}, + {ID: "c", Filetypes: []string{"c"}, Queries: []string{"highlights.scm"}, Language: language(tsc.Language)}, + {ID: "cpp", Filetypes: []string{"c++"}, Queries: []string{"c.scm", "highlights.scm"}, Language: language(tscpp.Language)}, + {ID: "bash", Filetypes: []string{"shell"}, Queries: []string{"highlights.scm"}, Language: language(tsbash.Language)}, + {ID: "json", Filetypes: []string{"json"}, Queries: []string{"highlights.scm"}, Language: language(tsjson.Language)}, + {ID: "toml", Filetypes: []string{"toml"}, Queries: []string{"highlights.scm"}, Language: language(tstoml.Language)}, + {ID: "html", Filetypes: []string{"html"}, Queries: []string{"highlights.scm"}, Language: language(tshtml.Language)}, + {ID: "css", Filetypes: []string{"css"}, Queries: []string{"highlights.scm"}, Language: language(tscss.Language)}, + {ID: "lua", Filetypes: []string{"lua"}, Queries: []string{"highlights.scm"}, Language: language(tslua.Language)}, + {ID: "yaml", Filetypes: []string{"yaml"}, Queries: []string{"highlights.scm"}, Language: language(tsyaml.Language)}, +} + +func language(constructor func() unsafe.Pointer) func() *sitter.Language { + return func() *sitter.Language { return sitter.NewLanguage(constructor()) } +} + +func definitionForFiletype(filetype string) (Definition, bool) { + for _, definition := range definitions { + for _, candidate := range definition.Filetypes { + if candidate == filetype { + return definition, true + } + } + } + return Definition{}, false } // captureGroups maps a tree-sitter capture name to a micro highlight group. @@ -137,20 +185,6 @@ func MapCapture(name string) string { return "" } -// goKeywordGroups refines the grammar's single @keyword capture into the -// vocabulary the old go.yaml used: declarations are preproc, control-flow -// keywords are special, and the rest stay statement. -var goKeywordGroups = map[string]string{ - "package": "preproc", "import": "preproc", "const": "preproc", "var": "preproc", - "type": "preproc", "struct": "preproc", "func": "preproc", "interface": "preproc", - "map": "preproc", "chan": "preproc", "defer": "preproc", - "break": "special", "case": "special", "continue": "special", "default": "special", - "go": "special", "goto": "special", "range": "special", "return": "special", - "fallthrough": "special", - "else": "statement", "for": "statement", "if": "statement", - "switch": "statement", "select": "statement", -} - type grammar struct { id string parser *sitter.Parser @@ -192,7 +226,8 @@ type Engine struct { tick uint64 } -// New returns an engine that looks for //highlights.scm. +// New returns an engine that loads each grammar's query files from +// //. func New(grammarDir string) *Engine { return &Engine{ grammarDir: grammarDir, @@ -248,38 +283,38 @@ func (e *Engine) grammarFor(filetype string) (*grammar, error) { } func (e *Engine) loadGrammar(filetype string) (*grammar, error) { - grammarID, ok := filetypeGrammar[filetype] + definition, ok := definitionForFiletype(filetype) if !ok { return nil, fmt.Errorf("unsupported filetype %q: no linked tree-sitter grammar", filetype) } - newLanguage, ok := grammarLanguages[grammarID] - if !ok { - return nil, fmt.Errorf("filetype %q maps to unknown grammar %q", filetype, grammarID) - } - language := newLanguage() + language := definition.Language() if language == nil { - return nil, fmt.Errorf("grammar %q returned a nil language", grammarID) + return nil, fmt.Errorf("grammar %q returned a nil language", definition.ID) } - queryPath, err := e.findQuery(grammarID) - if err != nil { - return nil, err - } - querySource, err := os.ReadFile(queryPath) - if err != nil { - return nil, fmt.Errorf("read %s: %w", queryPath, err) + var sources []string + for _, name := range definition.Queries { + queryPath, err := e.findQuery(definition.ID, name) + if err != nil { + return nil, err + } + querySource, err := os.ReadFile(queryPath) + if err != nil { + return nil, fmt.Errorf("read %s: %w", queryPath, err) + } + sources = append(sources, string(querySource)) } - query, queryErr := sitter.NewQuery(language, string(querySource)) + query, queryErr := sitter.NewQuery(language, strings.Join(sources, "\n")) if queryErr != nil { - return nil, fmt.Errorf("%s: %s", queryPath, queryErr.Message) + return nil, fmt.Errorf("grammar %q query: %s", definition.ID, queryErr.Message) } parser := sitter.NewParser() if err := parser.SetLanguage(language); err != nil { query.Close() parser.Close() - return nil, fmt.Errorf("set language %q: %w", grammarID, err) + return nil, fmt.Errorf("set language %q: %w", definition.ID, err) } return &grammar{ - id: grammarID, + id: definition.ID, parser: parser, language: language, query: query, @@ -287,18 +322,20 @@ func (e *Engine) loadGrammar(filetype string) (*grammar, error) { }, nil } -func (e *Engine) findQuery(grammarID string) (string, error) { - candidates := []string{ - filepath.Join(e.grammarDir, grammarID, "highlights.scm"), - filepath.Join(e.grammarDir, grammarID+".scm"), - filepath.Join(e.grammarDir, "highlights.scm"), +func (e *Engine) findQuery(grammarID, name string) (string, error) { + candidates := []string{filepath.Join(e.grammarDir, grammarID, name)} + if name == "highlights.scm" { + candidates = append(candidates, + filepath.Join(e.grammarDir, grammarID+".scm"), + filepath.Join(e.grammarDir, "highlights.scm"), + ) } for _, path := range candidates { if info, err := os.Stat(path); err == nil && !info.IsDir() { return path, nil } } - return "", fmt.Errorf("no highlights.scm for grammar %q under %q", grammarID, e.grammarDir) + return "", fmt.Errorf("no %s for grammar %q under %q", name, grammarID, e.grammarDir) } func (e *Engine) touchSession(id int64) *session { @@ -456,11 +493,6 @@ func (g *grammar) spans(tree *sitter.Tree, text []byte) []Span { continue } node := capture.Node - if g.id == "go" && name == "keyword" { - if refined, ok := goKeywordGroups[nodeText(text, &node)]; ok { - group = refined - } - } start, end := node.StartPosition(), node.EndPosition() spec := strings.Count(name, ".") + 1 for _, span := range splitNode(lines, int(start.Row), int(start.Column), int(end.Row), int(end.Column)) { @@ -477,14 +509,6 @@ func (g *grammar) spans(tree *sitter.Tree, text []byte) []Span { return resolve(raw) } -func nodeText(text []byte, node *sitter.Node) string { - start, end := node.StartByte(), node.EndByte() - if start > end || end > uint(len(text)) { - return "" - } - return string(text[start:end]) -} - func splitNode(lines lineIndex, startRow, startColumn, endRow, endColumn int) []Span { if startRow < 0 { startRow = 0 diff --git a/treesitter/internal/tsprovider/provider_test.go b/treesitter/internal/tsprovider/provider_test.go index 63b3b3ab..0ddf69a7 100644 --- a/treesitter/internal/tsprovider/provider_test.go +++ b/treesitter/internal/tsprovider/provider_test.go @@ -4,6 +4,7 @@ import ( "bytes" "math/rand" "os" + "strings" "testing" ) @@ -38,6 +39,136 @@ func haveSpan(spans []Span, line, start, end int, group string) bool { return false } +func lineSpans(spans []Span, line int) []Span { + var out []Span + for _, span := range spans { + if span.Line == line { + out = append(out, span) + } + } + return out +} + +// positionOf returns the 0-based line and byte column of the first occurrence +// of needle in text. +func positionOf(t *testing.T, text, needle string) (int, int) { + t.Helper() + at := strings.Index(text, needle) + if at < 0 { + t.Fatalf("%q missing from sample %q", needle, text) + } + line := strings.Count(text[:at], "\n") + lineStart := strings.LastIndexByte(text[:at], '\n') + 1 + return line, at - lineStart +} + +func TestDefinitionRegistry(t *testing.T) { + want := map[string]string{ + "go": "go", + "python": "python", + "javascript": "javascript", + "typescript": "tsx", + "rust": "rust", + "c": "c", + "c++": "cpp", + "shell": "bash", + "json": "json", + "toml": "toml", + "html": "html", + "css": "css", + "lua": "lua", + "yaml": "yaml", + } + for filetype, id := range want { + definition, ok := definitionForFiletype(filetype) + if !ok { + t.Errorf("filetype %q is not in the grammar registry", filetype) + continue + } + if definition.ID != id { + t.Errorf("filetype %q: grammar id %q, want %q", filetype, definition.ID, id) + } + } + if _, ok := definitionForFiletype("ruby"); ok { + t.Error("ruby unexpectedly mapped to a linked grammar") + } + if got := len(definitionForFiletypeForTest(t, "typescript").Queries); got != 2 { + t.Errorf("typescript query layers: got %d, want 2", got) + } + if got := len(definitionForFiletypeForTest(t, "c++").Queries); got != 2 { + t.Errorf("c++ query layers: got %d, want 2", got) + } +} + +func definitionForFiletypeForTest(t *testing.T, filetype string) Definition { + t.Helper() + definition, ok := definitionForFiletype(filetype) + if !ok { + t.Fatalf("filetype %q is not in the grammar registry", filetype) + } + return definition +} + +// TestLanguageSpans checks that every linked grammar parses a small snippet +// and emits the expected micro groups. Each case also exercises the query +// layers: TypeScript over JavaScript and C++ over C. +func TestLanguageSpans(t *testing.T) { + type want struct { + text string + group string + } + cases := []struct { + filetype string + text string + wants []want + }{ + {"go", "package main\n\nfunc main() { s := \"alpha\" }\n", []want{{"alpha", "constant.string"}, {"func", "statement"}}}, + {"python", "def greet(name):\n return \"alpha\" # note\n", []want{{"alpha", "constant.string"}, {"note", "comment"}, {"def", "statement"}}}, + {"javascript", "const n = 42; // note\n", []want{{"42", "constant.number"}, {"note", "comment"}}}, + {"typescript", "interface User {\n id: number;\n}\n\n// note\nconst user: User = { id: 1 };\n", []want{{"interface", "statement"}, {"note", "comment"}}}, + {"rust", "fn main() {\n let s = \"alpha\"; // note\n}\n", []want{{"alpha", "constant.string"}, {"note", "comment"}}}, + {"c", "int main(void) { return 42; } // note\n", []want{{"42", "constant.number"}, {"note", "comment"}}}, + {"c++", "const char *s = \"alpha\"; // note\n", []want{{"alpha", "constant.string"}, {"note", "comment"}}}, + {"shell", "echo \"alpha\" # note\n", []want{{"alpha", "constant.string"}, {"note", "comment"}}}, + {"json", "{\"key\": 42, \"flag\": true}\n", []want{{"key", "constant.string"}, {"true", "constant.bool"}}}, + {"toml", "[server]\nport = 8080 # note\nname = \"alpha\"\n", []want{{"note", "comment"}, {"alpha", "constant.string"}}}, + {"html", "
hi
\n", []want{{"div", "symbol.tag"}, {"alpha", "constant.string"}}}, + {"css", "body { color: red; } /* note */\n", []want{{"body", "symbol.tag"}, {"note", "comment"}}}, + {"lua", "local s = \"alpha\" -- note\n", []want{{"alpha", "constant.string"}, {"note", "comment"}}}, + {"yaml", "name: alpha\nenabled: true\n", []want{{"alpha", "constant.string"}, {"true", "constant.bool"}}}, + } + + engine := New("../../queries") + defer engine.Close() + for _, tc := range cases { + t.Run(tc.filetype, func(t *testing.T) { + spans, err := engine.Spans(int64(len(tc.filetype)), tc.filetype, []byte(tc.text)) + if err != nil { + t.Fatalf("%s: %v", tc.filetype, err) + } + if len(spans) == 0 { + t.Fatalf("%s: no spans", tc.filetype) + } + for _, want := range tc.wants { + line, column := positionOf(t, tc.text, want.text) + if got := groupAt(spans, line, column); got != want.group { + t.Errorf("%s: group at %q = %q, want %q; line spans %v", + tc.filetype, want.text, got, want.group, lineSpans(spans, line)) + } + } + lines := splitLines([]byte(tc.text)) + for _, span := range spans { + if span.Line < 0 || span.Line >= len(lines) { + t.Fatalf("%s: span line out of range: %+v", tc.filetype, span) + } + if span.Start < 0 || span.End > len(lines[span.Line]) || span.End <= span.Start { + t.Fatalf("%s: span out of bounds: %+v (line length %d)", tc.filetype, span, len(lines[span.Line])) + } + } + }) + } +} + func TestGoSampleSpans(t *testing.T) { engine := New("../../queries") defer engine.Close() @@ -48,15 +179,15 @@ func TestGoSampleSpans(t *testing.T) { } lines := splitLines(text) - // Go keyword refinement: declarations are preproc, control flow is - // special, and conditional keywords stay statement. - if got := groupAt(spans, 0, 0); got != "preproc" { + // The upstream Go query emits one @keyword capture for every keyword, so + // all of them map to the generic statement group. + if got := groupAt(spans, 0, 0); got != "statement" { t.Errorf("package keyword: got %q", got) } - if got := groupAt(spans, 11, 0); got != "preproc" { + if got := groupAt(spans, 11, 0); got != "statement" { t.Errorf("func keyword: got %q", got) } - if got := groupAt(spans, 16, 1); got != "special" { + if got := groupAt(spans, 16, 1); got != "statement" { t.Errorf("return keyword: got %q", got) } @@ -123,59 +254,57 @@ func TestGoSampleSpans(t *testing.T) { } } -func lineSpans(spans []Span, line int) []Span { - var out []Span - for _, span := range spans { - if span.Line == line { - out = append(out, span) - } +func TestIncrementalMatchesFullParse(t *testing.T) { + bases := map[string]string{ + "go": string(loadSample(t)), + "python": "def f(x):\n return [x, \"a\"]\n\nclass C:\n pass\n", } - return out -} + insertions := []string{ + "x", "\n", "# note\n", "\"", "é", "`", "/*", "\t", " ", + "日本語", "\r\n", ")", "}", " pass\n", "alpha", + } + for filetype, base := range bases { + t.Run(filetype, func(t *testing.T) { + engine := New("../../queries") + defer engine.Close() + text := append([]byte(nil), base...) -func TestIncrementalMatchesFullParse(t *testing.T) { - engine := New("../../queries") - defer engine.Close() - text := append([]byte(nil), loadSample(t)...) + random := rand.New(rand.NewSource(7)) + for step := 0; step < 25; step++ { + switch random.Intn(3) { + case 0: + at := random.Intn(len(text) + 1) + insert := insertions[random.Intn(len(insertions))] + text = append(text[:at:at], append([]byte(insert), text[at:]...)...) + case 1: + if len(text) > 1 { + at := random.Intn(len(text)) + n := 1 + random.Intn(min(5, len(text)-at)) + text = append(text[:at:at], text[at+n:]...) + } + default: + if len(text) > 1 { + at := random.Intn(len(text)) + n := 1 + random.Intn(min(5, len(text)-at)) + text = append(text[:at:at], append([]byte(insertions[random.Intn(len(insertions))]), text[at+n:]...)...) + } + } - insertions := []string{ - "x", "\n", "// note\n", "func f() {}\n", "é", "`", "/*", "\tvar _ = 1\n", - "日本語", "\r\n", "string", "}", "package p\n", - } - random := rand.New(rand.NewSource(7)) - for step := 0; step < 40; step++ { - switch random.Intn(3) { - case 0: - at := random.Intn(len(text) + 1) - insert := insertions[random.Intn(len(insertions))] - text = append(text[:at:at], append([]byte(insert), text[at:]...)...) - case 1: - if len(text) > 1 { - at := random.Intn(len(text)) - n := 1 + random.Intn(min(5, len(text)-at)) - text = append(text[:at:at], text[at+n:]...) - } - default: - if len(text) > 1 { - at := random.Intn(len(text)) - n := 1 + random.Intn(min(5, len(text)-at)) - text = append(text[:at:at], append([]byte(insertions[random.Intn(len(insertions))]), text[at+n:]...)...) + got, err := engine.Spans(1, filetype, text) + if err != nil { + t.Fatalf("step %d: incremental parse: %v", step, err) + } + fresh := New("../../queries") + want, err := fresh.Spans(1, filetype, text) + fresh.Close() + if err != nil { + t.Fatalf("step %d: full parse: %v", step, err) + } + if !equalSpans(got, want) { + t.Fatalf("step %d: incremental spans differ from full parse\ngot %v\nwant %v\ntext:\n%s", step, got, want, text) + } } - } - - got, err := engine.Spans(1, "go", text) - if err != nil { - t.Fatalf("step %d: incremental parse: %v", step, err) - } - fresh := New("../../queries") - want, err := fresh.Spans(1, "go", text) - fresh.Close() - if err != nil { - t.Fatalf("step %d: full parse: %v", step, err) - } - if !equalSpans(got, want) { - t.Fatalf("step %d: incremental spans differ from full parse\ngot %v\nwant %v\ntext:\n%s", step, got, want, text) - } + }) } } @@ -271,7 +400,7 @@ func TestSessionsAreBounded(t *testing.T) { func TestUnsupportedFiletype(t *testing.T) { engine := New("../../queries") defer engine.Close() - if _, err := engine.Spans(1, "python", []byte("print(1)\n")); err == nil { + if _, err := engine.Spans(1, "ruby", []byte("puts 1\n")); err == nil { t.Fatal("expected an error for an unlinked filetype") } } -- knotmirror 0.0.1