diff --git a/README.md b/README.md index e3a78a6..0787adf 100644 --- a/README.md +++ b/README.md @@ -102,6 +102,9 @@ Compress(msgs) │ SearchCrusher — group grep matches + context by file (path shown once); │ keep first/last + top-scored + errors; offload the rest, │ with a redundancy-adaptive global cap + │ CodeCrusher — Go (go/ast): keep package/imports/types/signatures, + │ elide function bodies; Python/TS/JS/Java/C/C++ via the + │ optional tree-sitter codelang module; else → text │ TextCrusher — line mode: dedup near-identical lines (numeric-insensitive) │ prose mode: extractive sentence selection (hybrid │ BM25 + matched-term boost), verbatim only @@ -218,10 +221,11 @@ requiring every skill/prompt author to document compression. | `ctxzip` | `Compress`, `Unzip`, `Message` / `Options` / `Result` | | `detect` | content-type detection (heuristic cascade, most-specific first) | | `router` | content type → crusher | -| `crush` | `JSONCrusher`, `LogCrusher`, `DiffCrusher`, `SearchCrusher`, `TextCrusher`, `Scorer` (hybrid BM25), error floor | +| `crush` | `JSONCrusher`, `LogCrusher`, `DiffCrusher`, `SearchCrusher`, `CodeCrusher`, `TextCrusher`, `Scorer` (hybrid BM25), error floor | | `ccr` | `Store` interface, `MemoryStore`, `BoltStore`, hashing, marker grammar | | `tokenize` | approximate token counting (CJK-aware; deliberately estimates high) | | `cmd/ctxzip-demo` | pipe-anything demo CLI | +| `codelang` | optional tree-sitter Tier B code crusher (Python/TS/JS/Java/C/C++) — separate module (CGO) | Design invariants, enforced by tests: non-empty input never compresses to empty output; the caller's message slice is never mutated; a failed or unprofitable @@ -256,6 +260,17 @@ Shipped: repetitive results aren't padded to the ceiling. Content-complete (every match shown or retrievable byte-for-byte), though grouping regroups line order by design +- [x] code crusher — keeps package clause, imports, type/const/var decls, and + every function signature, and elides function bodies (position-preserving, + byte-exact round-trip on expansion). The generic error floor is + intentionally not applied to code (Go bodies are saturated with "error"); + MustKeep and query terms still protect a body. + - **Tier A** (core, no third-party dep): Go via the stdlib `go/parser`+`go/ast`. + - **Tier B** (optional [`codelang`](codelang) module): Python / TypeScript / + JavaScript / Java / C / C++ via tree-sitter, wired through + `CodeCrusher.Fallback`. Isolated in its own module so the core stays + dependency-light (CGO + grammars only where wanted). Other languages + still fall back to extractive text, so nothing regresses - [x] line-mode text compression (grep/log layout preserved byte-faithfully) - [x] durable `BoltStore` (restart-safe originals) - [x] caller `MustKeep` vocabulary + extended k8s error floor @@ -265,7 +280,8 @@ Shipped: Planned: -- [ ] dedicated code crusher (AST, build-tagged tree-sitter) +- [ ] code crusher — more tree-sitter languages (Rust, Ruby, Kotlin, …) in the + `codelang` module; score-aware body budgeting - [ ] richer categorical drop summaries ("149 Running, 3 error-like") in markers - [ ] optional ML prose path behind the same `Compressor` interface - [ ] optional embedding relevance scorer (build-tagged) fused with BM25 diff --git a/codelang/README.md b/codelang/README.md new file mode 100644 index 0000000..45fb0a4 --- /dev/null +++ b/codelang/README.md @@ -0,0 +1,64 @@ +# ctxzip/codelang + +Tree-sitter-backed **Tier B** code crusher for ctxzip: signature-keeping, +body-eliding compression for **Python, TypeScript, JavaScript, Java, C, and +C++** — the languages the core crusher's stdlib `go/ast` path (Tier A, Go only) +doesn't cover. + +It's a **separate Go module** on purpose. tree-sitter pulls in CGO and the C +grammars; ctxzip's core stays dependency-light (bbolt only), and only hosts that +want multi-language code compression take on this dependency. + +## Use + +Wire it into the core crusher via the `Fallback` seam — Go goes through the +stdlib AST path, everything else through tree-sitter: + +```go +import ( + "github.com/initializ/ctxzip/crush" + "github.com/initializ/ctxzip/codelang" +) + +code := crush.NewCodeCrusher() // Go via go/ast (no CGO) +code.Fallback = codelang.NewCrusher() // Python/TS/JS/Java/C/C++ via tree-sitter +// ...use `code` as the SourceCode compressor in your router. +``` + +`codelang.Crusher` implements `crush.Compressor`, so it can also be used stand-alone. + +## How it works + +1. **Detect** the language from content signals; candidates are tried in + likelihood order (specific dialects — TS over JS, C++ over C — first). +2. **Parse** with the candidate grammar and **verify**: a parse whose + ERROR/MISSING nodes cover more than 15% of the source is rejected and the + next candidate tried. A wrong-but-clean parse still elides real body nodes, so + compression stays valid. +3. **Elide** each function/method body worth eliding (`> MinBodyLines`), + preserving imports, type declarations, class headers, and every signature. + Braces are kept; only the bytes between them are offloaded (Python's + brace-less block is offloaded whole after the `:`). + +Same invariants as every ctxzip crusher: + +- **Reversible & byte-exact** — each body's interior is offloaded and replaced in + place by a `<>` marker, so expanding every marker reconstructs the + source byte-for-byte (tested per language). +- **Deterministic**, **fail-open** (nil store / unsupported language / parse + trouble → extractive text fallback, so nothing is lost), never-empty. +- The generic **error-vocabulary floor is not applied to code** (bodies are full + of `error`/`Errorf`); `MustKeep` and query terms still protect a body. + +## Develop + +```bash +go test ./... # requires CGO + a C toolchain (builds the tree-sitter grammars) +``` + +The `replace github.com/initializ/ctxzip => ../` in `go.mod` builds this module +against the working tree; it's dropped once the core module is tagged. + +## License + +Apache-2.0. diff --git a/codelang/crusher.go b/codelang/crusher.go new file mode 100644 index 0000000..8cccbf7 --- /dev/null +++ b/codelang/crusher.go @@ -0,0 +1,221 @@ +// Package codelang is the tree-sitter-backed Tier B code crusher for ctxzip: +// signature-keeping, body-eliding compression for Python, TypeScript, +// JavaScript, Java, C, and C++ — the languages ctxzip's stdlib go/ast crusher +// (Tier A, Go only) does not cover. +// +// It lives in its own Go module so ctxzip's core stays dependency-light (bbolt +// only): tree-sitter pulls in CGO and the C grammars, which only hosts that want +// multi-language code compression need. Wire it into the core crusher via +// crush.CodeCrusher.Fallback: +// +// code := crush.NewCodeCrusher() // Go via go/ast +// code.Fallback = codelang.NewCrusher() // Python/TS/JS/Java/C/C++ via tree-sitter +// +// It implements crush.Compressor and honors the same invariants: reversible +// (each elided body's interior is offloaded and replaced in place by a marker, +// so expanding every marker reconstructs the source byte-for-byte), +// deterministic, and fail-open (any parse trouble or unsupported language falls +// back to the extractive text crusher, so nothing is ever lost). +package codelang + +import ( + "context" + "fmt" + "sort" + "strings" + + sitter "github.com/smacker/go-tree-sitter" + + "github.com/initializ/ctxzip/ccr" + "github.com/initializ/ctxzip/crush" +) + +// Crusher compresses source code by eliding function bodies. It handles the six +// tree-sitter languages and delegates everything else to an extractive text +// crusher. +type Crusher struct { + // MinBodyLines is the body line-span below which a function is left whole. + MinBodyLines int + // MinLines is the file size (in newlines) below which it is left alone. + MinLines int + // maxErrorRatio rejects a parse whose ERROR/MISSING nodes cover more than + // this fraction of the source — a sign the grammar guess was wrong. + maxErrorRatio float64 + text *crush.TextCrusher +} + +// NewCrusher returns a Crusher with sensible defaults. +func NewCrusher() *Crusher { + return &Crusher{MinBodyLines: 3, MinLines: 15, maxErrorRatio: 0.15, text: crush.NewTextCrusher()} +} + +// Name implements crush.Compressor. +func (c *Crusher) Name() string { return "code_treesitter" } + +// Compress implements crush.Compressor. +func (c *Crusher) Compress(req crush.Request) (crush.Result, error) { + if req.Store == nil || strings.Count(req.Content, "\n") < c.MinLines { + return crush.Result{Compressed: req.Content, Strategy: c.Name()}, nil + } + // Try each candidate grammar in likelihood order; the first that parses + // cleanly and finds bodies to elide wins. A wrong-but-clean parse still + // elides real body nodes, so compression stays valid and reversible. + for i, lang := range detectCandidates(req.Content) { + if i >= 3 { + break // bound the work + } + if res, ok := c.compress(req, lang); ok { + return res, nil + } + } + // Unsupported language or nothing to elide: extractive text, so no regression. + return c.text.Compress(req) +} + +// span is the interior byte range [start,end) of one body to elide. +type span struct { + start, end, lines int +} + +func (c *Crusher) compress(req crush.Request, lang language) (crush.Result, bool) { + src := []byte(req.Content) + parser := sitter.NewParser() + parser.SetLanguage(lang.grammar()) + tree, err := parser.ParseCtx(context.Background(), nil, src) + if err != nil || tree == nil { + return crush.Result{}, false + } + defer tree.Close() + root := tree.RootNode() + if errorByteRatio(root, len(src)) > c.maxErrorRatio { + return crush.Result{}, false // wrong grammar for this content + } + + mustKeep := crush.NormalizeMustKeep(req.MustKeep) + terms := queryWords(req.Query) + var spans []span + collectBodies(root, lang, src, mustKeep, terms, c.MinBodyLines, &spans) + if len(spans) == 0 { + return crush.Result{}, false + } + sort.Slice(spans, func(i, j int) bool { return spans[i].start < spans[j].start }) + + var b strings.Builder + var markers []string + prev := 0 + for _, sp := range spans { + b.Write(src[prev:sp.start]) + interior := src[sp.start:sp.end] + hash := ccr.Hash(interior) + if err := req.Store.Put(hash, interior, ccr.Meta{ + ToolName: req.ToolName, + Query: req.Query, + ItemCount: sp.lines, + OriginalKind: "code", + }); err != nil { + b.Write(interior) // fail-open: keep this body verbatim + prev = sp.end + continue + } + markers = append(markers, hash) + b.WriteString(ccr.Marker(hash, fmt.Sprintf("%d_lines_offloaded", sp.lines))) + prev = sp.end + } + b.Write(src[prev:]) + if len(markers) == 0 { + return crush.Result{}, false + } + return crush.Result{Compressed: b.String(), Strategy: c.Name() + ":" + lang.name, Markers: markers}, true +} + +// collectBodies walks the tree recording the interior span of each function body +// worth eliding. It does not descend into an elided body (its interior is +// offloaded whole), so recorded spans never overlap. +func collectBodies(n *sitter.Node, lang language, src []byte, mustKeep, terms []string, minLines int, out *[]span) { + for i := 0; i < int(n.NamedChildCount()); i++ { + child := n.NamedChild(i) + if lang.funcTypes[child.Type()] { + if body := child.ChildByFieldName("body"); body != nil && lang.bodyTypes[body.Type()] { + if sp, ok := bodySpan(body, lang, src, mustKeep, terms, minLines); ok { + *out = append(*out, sp) + continue // do not descend into an elided function + } + } + } + collectBodies(child, lang, src, mustKeep, terms, minLines, out) + } +} + +// bodySpan returns the elidable interior of a body node, or ok=false if it is +// too small or protected by MustKeep / a query term. +func bodySpan(body *sitter.Node, lang language, src []byte, mustKeep, terms []string, minLines int) (span, bool) { + lines := int(body.EndPoint().Row) - int(body.StartPoint().Row) + if lines <= minLines { + return span{}, false + } + start, end := int(body.StartByte()), int(body.EndByte()) + if lang.usesBraces { + // Keep the braces; elide only what's between them. + if end <= start+1 || src[start] != '{' || src[end-1] != '}' { + return span{}, false + } + start, end = start+1, end-1 + } + if start >= end { + return span{}, false + } + // Caller/query intent keeps a body whole. The generic error-vocabulary floor + // is intentionally not applied to code (bodies are full of "error"/"Errorf"); + // bodies are reversible anyway. Mirrors the Go crusher. + lower := strings.ToLower(string(src[start:end])) + if containsAny(lower, mustKeep) || (len(terms) > 0 && containsAny(lower, terms)) { + return span{}, false + } + return span{start, end, lines}, true +} + +// errorByteRatio is the fraction of source bytes covered by ERROR/MISSING nodes. +func errorByteRatio(root *sitter.Node, total int) float64 { + if total == 0 { + return 1 + } + var errBytes int + var walk func(n *sitter.Node) + walk = func(n *sitter.Node) { + if n.IsError() || n.IsMissing() { + errBytes += int(n.EndByte() - n.StartByte()) + return // don't double-count the subtree + } + for i := 0; i < int(n.ChildCount()); i++ { + walk(n.Child(i)) + } + } + walk(root) + return float64(errBytes) / float64(total) +} + +// queryWords splits a query into lowercase words of 3+ chars (symbol-ish tokens). +func queryWords(q string) []string { + if q == "" { + return nil + } + fields := strings.FieldsFunc(strings.ToLower(q), func(r rune) bool { + return !(r >= 'a' && r <= 'z' || r >= '0' && r <= '9' || r == '_') + }) + out := fields[:0] + for _, f := range fields { + if len(f) >= 3 { + out = append(out, f) + } + } + return out +} + +func containsAny(s string, terms []string) bool { + for _, t := range terms { + if t != "" && strings.Contains(s, t) { + return true + } + } + return false +} diff --git a/codelang/crusher_test.go b/codelang/crusher_test.go new file mode 100644 index 0000000..f512797 --- /dev/null +++ b/codelang/crusher_test.go @@ -0,0 +1,326 @@ +package codelang + +import ( + "regexp" + "strings" + "testing" + + "github.com/initializ/ctxzip/ccr" + "github.com/initializ/ctxzip/crush" +) + +var markerRe = regexp.MustCompile(`<]*)?>>`) + +func expand(t *testing.T, compressed string, store ccr.Store) string { + t.Helper() + return markerRe.ReplaceAllStringFunc(compressed, func(m string) string { + h := markerRe.FindStringSubmatch(m)[1] + e, ok := store.Get(h) + if !ok { + t.Fatalf("marker %s not retrievable", h) + } + return string(e.Original) + }) +} + +// Each sample is a realistically-sized file (> MinLines) with at least one +// function whose body is worth eliding (> MinBodyLines) and a distinctive +// innard, plus enough structure to detect the language. +var samples = map[string]struct { + lang string // expected strategy language suffix + src string + innard string // a body line that must be offloaded (gone from output) + keep string // a structural line that must survive +}{ + "python": {"python", `# widget builder module +import os +import sys + + +class Widget: + def build(self, name, size): + if not name: + raise KeyError("required") + total = size * 2 + label = name.upper() + parts = [label, str(total)] + joined = "-".join(parts) + return joined + + def rename(self, new): + old = self.name + self.name = new + self.dirty = True + return old +`, "total = size * 2", "def build(self, name, size):"}, + + "javascript": {"javascript", `import fs from 'fs'; +import path from 'path'; + + +function build(name, size) { + if (!name) { + return null; + } + const total = size * 2; + const label = name.toUpperCase(); + const parts = [label, total]; + return parts.join('-'); +} + +function rename(obj, next) { + const old = obj.name; + obj.name = next; + obj.dirty = true; + return old; +} +`, "const total = size * 2", "function build(name, size) {"}, + + "typescript": {"typescript", `interface Config { + name: string; + size: number; +} + +function build(cfg: Config): string { + if (!cfg.name) { + return ''; + } + const total: number = cfg.size * 2; + const label = cfg.name.toUpperCase(); + const parts = [label, String(total)]; + return parts.join('-'); +} + +function rename(cfg: Config, next: string): string { + const old = cfg.name; + cfg.name = next; + return old; +} +`, "const total: number", "function build(cfg: Config): string {"}, + + "java": {"java", `package widget; + +import java.util.List; + +public class Builder { + private String name; + + public String build(String name, int size) { + if (name.isEmpty()) { + return ""; + } + int total = size * 2; + String label = name.toUpperCase(); + String joined = label + "-" + total; + return joined; + } + + public String rename(String next) { + String old = this.name; + this.name = next; + return old; + } +} +`, "int total = size * 2", "public String build(String name, int size) {"}, + + "c": {"c", `#include +#include + +int build(const char* name, int size) { + if (name == 0) { + return -1; + } + int total = size * 2; + int scaled = total + 1; + printf("%s %d\n", name, scaled); + return scaled; +} + +int copy_name(char* dst, const char* src) { + int n = 0; + while (src[n]) { + dst[n] = src[n]; + n = n + 1; + } + return n; +} +`, "int total = size * 2", "int build(const char* name, int size) {"}, + + "cpp": {"cpp", `#include +#include + +class Builder { +public: + std::string build(const std::string& name, int size) { + if (name.empty()) { + return ""; + } + int total = size * 2; + std::vector parts; + parts.push_back(name); + return name + std::to_string(total); + } + + std::string rename(const std::string& next) { + std::string old = this->name_; + this->name_ = next; + return old; + } + +private: + std::string name_; +}; +`, "int total = size * 2", "std::string build(const std::string& name, int size) {"}, +} + +func TestCrusher_PerLanguage(t *testing.T) { + for name, tc := range samples { + t.Run(name, func(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + res, err := NewCrusher().Compress(crush.Request{Content: tc.src, Store: store}) + if err != nil { + t.Fatal(err) + } + if res.Compressed == tc.src { + t.Fatalf("expected compression, got passthrough:\n%s", tc.src) + } + if want := "code_treesitter:" + tc.lang; res.Strategy != want { + t.Fatalf("strategy = %q, want %q", res.Strategy, want) + } + if !strings.Contains(res.Compressed, tc.keep) { + t.Fatalf("signature dropped: %q not in\n%s", tc.keep, res.Compressed) + } + if strings.Contains(res.Compressed, tc.innard) { + t.Fatalf("body not elided: %q still present", tc.innard) + } + if len(res.Markers) == 0 { + t.Fatal("no markers emitted") + } + // Byte-exact reversibility. + if got := expand(t, res.Compressed, store); got != tc.src { + t.Fatalf("round-trip not byte-exact:\n--- got ---\n%s", got) + } + }) + } +} + +// TestCrusher_TSNotClaimedByJS is the regression for the detection nit: a TS +// file with MANY JavaScript signals (const/function/=>/===) and only a couple of +// TS-only ones (interface, : number) must still be detected as TypeScript, so +// its typed signatures aren't parsed by the JS grammar and missed. +func TestCrusher_TSNotClaimedByJS(t *testing.T) { + src := `interface Opts { + limit: number; +} + +const factor = 2; +const scale = (n) => n * factor; + +function process(opts: Opts): number { + const base = opts.limit; + const out = scale(base); + const check = out === base; + return check ? out : base; +} + +function reduceAll(xs) { + const acc = xs.reduce((a, b) => a + b, 0); + const avg = acc / xs.length; + return avg; +} +` + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + res, _ := NewCrusher().Compress(crush.Request{Content: src, Store: store}) + if res.Strategy != "code_treesitter:typescript" { + t.Fatalf("TS file misdetected: strategy = %q, want code_treesitter:typescript", res.Strategy) + } + if got := expand(t, res.Compressed, store); got != src { + t.Fatal("round-trip not byte-exact") + } +} + +func TestCrusher_MustKeepProtectsBody(t *testing.T) { + src := `# module +import math + + +def alpha(x): + a = x + 1 + b = a * 2 + c = b - 3 + d = c + 4 + return d + + +def beta(y): + KEEPTOKEN = y + d = y + 1 + e = d * 2 + f = e - 3 + return f +` + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + res, _ := NewCrusher().Compress(crush.Request{ + Content: src, Store: store, MustKeep: []string{"KEEPTOKEN"}, + }) + if !strings.Contains(res.Compressed, "KEEPTOKEN") { + t.Fatal("MustKeep body was elided") + } + if strings.Contains(res.Compressed, "a = x + 1") { + t.Fatal("non-protected body should have been elided") + } +} + +func TestCrusher_UnsupportedLanguageFallsBack(t *testing.T) { + // Ruby is not one of the six grammars; the text crusher should handle it. + var sb strings.Builder + sb.WriteString("def handler\n") + for i := 0; i < 60; i++ { + sb.WriteString(" puts 'processing the current request'\n") + } + sb.WriteString("end\n") + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + res, _ := NewCrusher().Compress(crush.Request{Content: sb.String(), Store: store}) + if strings.HasPrefix(res.Strategy, "code_treesitter:") { + t.Fatalf("Ruby should not be claimed by a grammar, got %s", res.Strategy) + } +} + +func TestCrusher_Deterministic(t *testing.T) { + src := samples["python"].src + r1, _ := NewCrusher().Compress(crush.Request{Content: src, Store: ccr.NewMemoryStore(ccr.MemoryConfig{})}) + r2, _ := NewCrusher().Compress(crush.Request{Content: src, Store: ccr.NewMemoryStore(ccr.MemoryConfig{})}) + if r1.Compressed != r2.Compressed || strings.Join(r1.Markers, ",") != strings.Join(r2.Markers, ",") { + t.Fatal("compression is not deterministic") + } +} + +func TestCrusher_NilStore_Passthrough(t *testing.T) { + src := samples["python"].src + res, _ := NewCrusher().Compress(crush.Request{Content: src, Store: nil}) + if res.Compressed != src { + t.Fatal("nil store must force lossless passthrough") + } +} + +// TestIntegration_WiredIntoCoreCodeCrusher proves the seam: with the tree-sitter +// crusher injected as the core CodeCrusher's Fallback, Python routes to +// tree-sitter while Go still uses the stdlib AST path. +func TestIntegration_WiredIntoCoreCodeCrusher(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + code := crush.NewCodeCrusher() + code.Fallback = NewCrusher() + + res, _ := code.Compress(crush.Request{Content: samples["python"].src, Store: store}) + if !strings.HasPrefix(res.Strategy, "code_treesitter:") { + t.Fatalf("Python should route to the tree-sitter fallback, got %s", res.Strategy) + } + + goSrc := "package p\n\nimport \"fmt\"\n\n" + + "func Build(name string) string {\n\tx := 1\n\ty := 2\n\tz := x + y\n\treturn fmt.Sprintf(\"%s-%d\", name, z)\n}\n\n" + + "func Rename(a, b string) string {\n\told := a\n\ta = b\n\tb = old\n\treturn a + b\n}\n" + resGo, _ := code.Compress(crush.Request{Content: goSrc, Store: store}) + if resGo.Strategy != "code_crusher" { + t.Fatalf("Go should use the stdlib AST path, got %s", resGo.Strategy) + } +} diff --git a/codelang/go.mod b/codelang/go.mod new file mode 100644 index 0000000..ea23784 --- /dev/null +++ b/codelang/go.mod @@ -0,0 +1,15 @@ +module github.com/initializ/ctxzip/codelang + +go 1.26 + +require ( + github.com/initializ/ctxzip v0.0.0 + github.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82 +) + +require ( + go.etcd.io/bbolt v1.5.0 // indirect + golang.org/x/sys v0.45.0 // indirect +) + +replace github.com/initializ/ctxzip => ../ diff --git a/codelang/go.sum b/codelang/go.sum new file mode 100644 index 0000000..caccdd9 --- /dev/null +++ b/codelang/go.sum @@ -0,0 +1,16 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82 h1:6C8qej6f1bStuePVkLSFxoU22XBS165D3klxlzRg8F4= +github.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82/go.mod h1:xe4pgH49k4SsmkQq5OT8abwhWmnzkhpgnXeekbx2efw= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +go.etcd.io/bbolt v1.5.0 h1:S7GAl7Fxv12yohbwFfIbQCGDWbQbtDGPET4P/bD4lxU= +go.etcd.io/bbolt v1.5.0/go.mod h1:mkltfYE5aUHQxUct9N9V+Kp7aSjFqjgrhcXIS70Lrdk= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= +golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/codelang/lang.go b/codelang/lang.go new file mode 100644 index 0000000..07f5df5 --- /dev/null +++ b/codelang/lang.go @@ -0,0 +1,130 @@ +package codelang + +import ( + "regexp" + + sitter "github.com/smacker/go-tree-sitter" + "github.com/smacker/go-tree-sitter/c" + "github.com/smacker/go-tree-sitter/cpp" + "github.com/smacker/go-tree-sitter/java" + "github.com/smacker/go-tree-sitter/javascript" + "github.com/smacker/go-tree-sitter/python" + tsx "github.com/smacker/go-tree-sitter/typescript/typescript" +) + +// language couples a tree-sitter grammar with the node types whose function +// bodies this crusher elides. funcTypes are the definition nodes; bodyTypes are +// the valid body node types found via the "body" field. usesBraces is false for +// Python (its body is an indented block with no braces), true for the rest. +type language struct { + name string + grammar func() *sitter.Language + funcTypes map[string]bool + bodyTypes map[string]bool + usesBraces bool + // signals decides candidacy: does this look like the language at all? + signals *regexp.Regexp + // strong matches EXCLUSIVE markers — syntax a superset dialect has but its + // base does not (TS type annotations, C++ scope resolution). A hit here + // ranks the dialect ahead of its base, so TS beats JS and C++ beats C only + // when the dialect's own syntax is actually present, not merely on shared + // tokens like `class`. nil for the base languages (JS, C). + strong *regexp.Regexp +} + +func set(xs ...string) map[string]bool { + m := make(map[string]bool, len(xs)) + for _, x := range xs { + m[x] = true + } + return m +} + +// languages is the elision/detection registry. +var languages = []language{ + { + name: "typescript", grammar: tsx.GetLanguage, usesBraces: true, + funcTypes: set("function_declaration", "method_definition", "generator_function_declaration", "method_signature"), + bodyTypes: set("statement_block"), + // Candidacy requires a TS-specific marker (not plain function/const, which + // are shared with JS) so a pure-JS file never becomes a TS candidate. + signals: regexp.MustCompile(`\binterface\s+\w|:\s*(string|number|boolean|void|any)\b|\btype\s+\w+\s*=|\bimplements\b|\benum\s+\w|\bexport\s+(type|interface|enum)\b|\breadonly\s|\bas\s+\w`), + strong: regexp.MustCompile(`\binterface\s+\w|:\s*(string|number|boolean|void|any)\b|\btype\s+\w+\s*=|\bimplements\b|\benum\s+\w|\bexport\s+(type|interface|enum)\b|\breadonly\s`), + }, + { + name: "javascript", grammar: javascript.GetLanguage, usesBraces: true, + funcTypes: set("function_declaration", "method_definition", "generator_function_declaration"), + bodyTypes: set("statement_block"), + signals: regexp.MustCompile(`(?m)\bfunction\s|\bconst\s|\blet\s|=>|\brequire\(|\bmodule\.exports\b|\bconsole\.log\b|===`), + // base language: no exclusive markers (TS is the specific dialect). + }, + { + name: "cpp", grammar: cpp.GetLanguage, usesBraces: true, + funcTypes: set("function_definition"), + bodyTypes: set("compound_statement"), + // No bare #include here (shared with C); C++ candidacy needs C++-ish syntax. + signals: regexp.MustCompile(`\bstd::|::\w|\btemplate\s*<|\bnamespace\s|\bnullptr\b|\bcout\b|\bclass\s+\w`), + strong: regexp.MustCompile(`\bstd::|\w+::\w|\btemplate\s*<|\bnullptr\b|\bcout\b|\bcin\b|\bnamespace\s+\w`), + }, + { + name: "c", grammar: c.GetLanguage, usesBraces: true, + funcTypes: set("function_definition"), + bodyTypes: set("compound_statement"), + signals: regexp.MustCompile(`(?m)#include|\bint\s+main\b|\bprintf\b|\bstruct\s+\w|\btypedef\b|\bmalloc\b|->`), + // base language: C++ is the specific dialect. + }, + { + name: "java", grammar: java.GetLanguage, usesBraces: true, + funcTypes: set("method_declaration", "constructor_declaration"), + bodyTypes: set("block", "constructor_body"), + signals: regexp.MustCompile(`\bpublic\s+class\b|\bpackage\s+\w|\bimport\s+java|@Override\b|\bSystem\.|\bpublic\s+static\s+void\s+main|\bclass\s+\w+\s*\{`), + strong: regexp.MustCompile(`\bpublic\s+class\b|\bpackage\s+[\w.]+\s*;|\bimport\s+java|@\w+|\bSystem\.(out|err)\b|\bpublic\s+static\s+void\s+main`), + }, + { + name: "python", grammar: python.GetLanguage, usesBraces: false, + funcTypes: set("function_definition"), + bodyTypes: set("block"), + signals: regexp.MustCompile(`(?m)^\s*def\s+\w+\s*\(|^\s*class\s+\w+.*:|^\s*(from\s+\w+\s+)?import\s+\w|\belif\b|\bself\b`), + strong: regexp.MustCompile(`(?m)^\s*def\s+\w+\s*\(.*\)\s*:|^\s*class\s+\w+.*:|\belif\b|\bself\b|\b__\w+__\b`), + }, +} + +// detectCandidates returns the languages whose signals fire in content, ordered +// by a score that ranks a specific dialect ahead of its base ONLY when the +// dialect's own exclusive syntax is present: score = strongHits*1000 + +// totalHits. So a TypeScript file (which also matches many JavaScript signals) +// is tried as TS before JS because its `interface`/type-annotation strong hits +// dominate, but a plain Java or C file — whose only cross-signal is a shared +// `class`/`#include` — is not hijacked by C++'s weight. The parse-verify step in +// the crusher then confirms the winner. +func detectCandidates(content string) []language { + type scored struct { + lang language + strong, all int + } + var hits []scored + for _, l := range languages { + all := len(l.signals.FindAllStringIndex(content, -1)) + if all == 0 { + continue + } + strong := 0 + if l.strong != nil { + strong = len(l.strong.FindAllStringIndex(content, -1)) + } + hits = append(hits, scored{l, strong, all}) + } + score := func(s scored) int { return s.strong*1000 + s.all } + for i := 0; i < len(hits); i++ { + for j := i + 1; j < len(hits); j++ { + if score(hits[j]) > score(hits[i]) { + hits[i], hits[j] = hits[j], hits[i] + } + } + } + out := make([]language, 0, len(hits)) + for _, h := range hits { + out = append(out, h.lang) + } + return out +} diff --git a/crush/code.go b/crush/code.go new file mode 100644 index 0000000..44a92f7 --- /dev/null +++ b/crush/code.go @@ -0,0 +1,157 @@ +package crush + +import ( + "fmt" + "go/ast" + "go/parser" + "go/token" + "strings" + + "github.com/initializ/ctxzip/ccr" +) + +// CodeCrusher compresses source code by eliding function bodies while keeping +// the file's shape — package clause, imports, type/const/var declarations, and +// every function signature. The body is the bulk; the signatures and structure +// are what a model usually needs to reason about a file. +// +// Tier A (this crusher) handles Go natively via the standard library's +// go/parser + go/ast — no third-party dependency, and AST-accurate brace +// matching so string/comment braces never fool it. Other languages go to +// Fallback: the tree-sitter-backed ctxzip/codelang module (Tier B — Python / +// TS / JS / Java / C / C++) handles them when injected, and it lives in its own +// module so the core keeps no CGO/tree-sitter dependency. With no Fallback set, +// they fall back to the extractive text crusher (unchanged from how SourceCode +// routed before), so nothing regresses either way. +// +// Reversibility is position-preserving and byte-exact: each elided body's +// interior (the bytes between its braces) is offloaded and replaced in place by +// a marker, so expanding every marker reconstructs the file byte-for-byte. The +// compressed form is for READING, not compiling. +// +// The generic error-vocabulary floor is deliberately NOT applied here. It exists +// to stop a log/tool-output crusher from dropping the error line a user is about +// to ask about; source code is different — Go bodies are saturated with "error" +// / "Errorf", so the floor would pin almost every function and defeat +// compression. Bodies are always reversible, so eliding one that mentions an +// error is not the catastrophic miss the floor guards against. Caller MustKeep +// vocabulary and query terms still protect a body from elision. +type CodeCrusher struct { + // MinBodyLines is the body line-span below which a function is left whole + // (a marker isn't worth it for a tiny body). + MinBodyLines int + // MinLines is the file size (in newlines) below which it is left alone. + MinLines int + // Fallback compresses source the Go AST path does not handle (non-Go, a + // package-less snippet, or unparseable input). When nil, an extractive text + // crusher is used — the same strategy SourceCode used before this crusher, + // so no language regresses. A tree-sitter-backed crusher (the ctxzip/codelang + // module) can be injected here to give Python/TS/JS/Java/C/C++ the same + // signature-keeping, body-eliding treatment Go gets. Must implement the same + // invariants (fail-open, reversible, deterministic). + Fallback Compressor + // text is the default fallback. + text *TextCrusher +} + +// NewCodeCrusher returns a CodeCrusher with sensible defaults. +func NewCodeCrusher() *CodeCrusher { + return &CodeCrusher{MinBodyLines: 3, MinLines: 15, text: NewTextCrusher()} +} + +// fallback returns the injected Fallback, or the default text crusher. +func (c *CodeCrusher) fallback() Compressor { + if c.Fallback != nil { + return c.Fallback + } + return c.text +} + +// Name implements Compressor. +func (c *CodeCrusher) Name() string { return "code_crusher" } + +// Compress implements Compressor. +func (c *CodeCrusher) Compress(req Request) (Result, error) { + if req.Store == nil || strings.Count(req.Content, "\n") < c.MinLines { + return passthrough(c.Name(), req.Content), nil + } + if res, ok := c.compressGo(req); ok { + return res, nil + } + // Non-Go, a bare snippet (no package clause), or unparseable: hand off to the + // fallback (a tree-sitter crusher if injected, else extractive text). + return c.fallback().Compress(req) +} + +// compressGo elides the bodies of top-level Go functions and methods. ok is +// false when the content is not a parseable Go file or nothing is worth eliding, +// so the caller can fall back. +func (c *CodeCrusher) compressGo(req Request) (Result, bool) { + src := req.Content + fset := token.NewFileSet() + f, err := parser.ParseFile(fset, "", src, parser.SkipObjectResolution) + if err != nil { + return Result{}, false + } + + terms := queryTerms(req.Query) + // interior byte range [start,end) of each body to elide, in source order. + type span struct { + start, end, lines int + } + var spans []span + for _, decl := range f.Decls { + fn, ok := decl.(*ast.FuncDecl) + if !ok || fn.Body == nil { + continue // imports/types/consts, or a body-less signature + } + lb := fset.Position(fn.Body.Lbrace) + rb := fset.Position(fn.Body.Rbrace) + lines := rb.Line - lb.Line + if lines <= c.MinBodyLines { + continue + } + start, end := lb.Offset+1, rb.Offset // exclusive of the braces + if start >= end { + continue + } + // Caller/query intent keeps a body whole (see the type doc on why the + // generic error floor is not applied to code). + body := strings.ToLower(src[start:end]) + if mustKeep(body, req.MustKeep) || (len(terms) > 0 && matchesAny(body, terms)) { + continue + } + spans = append(spans, span{start, end, lines}) + } + if len(spans) == 0 { + return Result{}, false + } + + var b strings.Builder + var markers []string + prev := 0 + for _, sp := range spans { + b.WriteString(src[prev:sp.start]) // through the opening brace + interior := src[sp.start:sp.end] + hash := ccr.Hash([]byte(interior)) + if err := req.Store.Put(hash, []byte(interior), ccr.Meta{ + ToolName: req.ToolName, + Query: req.Query, + ItemCount: sp.lines, + OriginalKind: "code", + }); err != nil { + b.WriteString(interior) // fail-open: keep this body verbatim + prev = sp.end + continue + } + markers = append(markers, hash) + b.WriteString(ccr.Marker(hash, fmt.Sprintf("%d_lines_offloaded", sp.lines))) + prev = sp.end + } + b.WriteString(src[prev:]) + + if len(markers) == 0 { + return Result{}, false + } + return Result{Compressed: b.String(), Strategy: c.Name(), Markers: markers}, true +} diff --git a/crush/code_test.go b/crush/code_test.go new file mode 100644 index 0000000..88fa628 --- /dev/null +++ b/crush/code_test.go @@ -0,0 +1,226 @@ +package crush + +import ( + "regexp" + "strings" + "testing" + + "github.com/initializ/ctxzip/ccr" +) + +// goSource is a small but realistic Go file: imports, a type, a const, a +// body-less signature is not valid at top level so we use a small func, and two +// functions with multi-line bodies worth eliding. +const goSource = `package widget + +import ( + "fmt" + "strings" +) + +// Config tunes a Widget. +type Config struct { + Name string + Size int +} + +const DefaultSize = 42 + +// Build assembles a widget from cfg and returns its label. +func Build(cfg Config) (string, error) { + if cfg.Name == "" { + return "", fmt.Errorf("name required") + } + size := cfg.Size + if size == 0 { + size = DefaultSize + } + label := strings.Repeat(cfg.Name, 1) + return fmt.Sprintf("%s[%d]", label, size), nil +} + +// tiny has a short body that is not worth eliding. +func tiny() int { + return DefaultSize +} + +func (c Config) Describe() string { + parts := []string{c.Name} + parts = append(parts, "sized") + parts = append(parts, strings.Repeat("x", c.Size)) + joined := strings.Join(parts, "-") + return joined +} +` + +var testMarkerRe = regexp.MustCompile(`<]*)?>>`) + +// expandCode replaces every marker with its stored original bytes. +func expandCode(t *testing.T, compressed string, store ccr.Store) string { + t.Helper() + return testMarkerRe.ReplaceAllStringFunc(compressed, func(m string) string { + h := testMarkerRe.FindStringSubmatch(m)[1] + e, ok := store.Get(h) + if !ok { + t.Fatalf("marker %s not retrievable", h) + } + return string(e.Original) + }) +} + +func TestCodeCrusher_ElidesBodiesKeepsShape(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + c := NewCodeCrusher() + + res, err := c.Compress(Request{Content: goSource, Store: store}) + if err != nil { + t.Fatal(err) + } + if res.Strategy != "code_crusher" { + t.Fatalf("expected code_crusher, got %s", res.Strategy) + } + if res.Compressed == goSource { + t.Fatal("expected body elision, got passthrough") + } + // Structure is kept verbatim. + for _, want := range []string{ + "package widget", + `import (`, + "type Config struct {", + "const DefaultSize = 42", + "func Build(cfg Config) (string, error) {", + "func (c Config) Describe() string {", + } { + if !strings.Contains(res.Compressed, want) { + t.Fatalf("structural line dropped: %q", want) + } + } + // The big bodies' innards are gone from the visible output... + if strings.Contains(res.Compressed, `fmt.Errorf("name required")`) { + t.Fatal("Build body not elided") + } + // ...but the tiny body is kept whole (<= MinBodyLines). + if !strings.Contains(res.Compressed, "return DefaultSize") { + t.Fatal("tiny body should not be elided") + } + // Two multi-line bodies elided -> two markers, both retrievable. + if len(res.Markers) != 2 { + t.Fatalf("want 2 markers, got %d", len(res.Markers)) + } + for _, h := range res.Markers { + if _, ok := store.Get(h); !ok { + t.Fatalf("body %s not retrievable", h) + } + } +} + +func TestCodeCrusher_RoundTripByteExact(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + res, _ := NewCodeCrusher().Compress(Request{Content: goSource, Store: store}) + if got := expandCode(t, res.Compressed, store); got != goSource { + t.Fatalf("round-trip not byte-exact:\n--- got ---\n%s", got) + } +} + +func TestCodeCrusher_MustKeepProtectsBody(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + c := NewCodeCrusher() + res, _ := c.Compress(Request{Content: goSource, Store: store, MustKeep: NormalizeMustKeep([]string{"DefaultSize"})}) + // Build's body references DefaultSize, so it must stay whole. + if !strings.Contains(res.Compressed, "size = DefaultSize") { + t.Fatal("MustKeep did not protect the body referencing the term") + } +} + +func TestCodeCrusher_NonGoFallsBackToText(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + c := NewCodeCrusher() + // Repetitive non-Go (Python-ish) source: the text fallback should dedup it. + var sb strings.Builder + sb.WriteString("def handler(req):\n") + for i := 0; i < 60; i++ { + sb.WriteString(" log.info('processing request')\n") + } + res, _ := c.Compress(Request{Content: sb.String(), Store: store}) + if res.Strategy == "code_crusher" { + t.Fatal("non-Go should not use the Go AST path") + } + if res.Compressed == sb.String() { + t.Fatal("expected the text fallback to compress repetitive non-Go source") + } +} + +func TestCodeCrusher_SnippetNoPackageFallsBack(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + c := NewCodeCrusher() + // A bare function snippet (no package clause) is not a parseable Go file. + snippet := "func orphan() {\n" + strings.Repeat("\tx := compute()\n", 40) + "}\n" + res, _ := c.Compress(Request{Content: snippet, Store: store}) + if res.Strategy == "code_crusher" { + t.Fatal("a package-less snippet should fall back, not use the AST path") + } +} + +// fakeFallback records whether it was invoked, standing in for the injectable +// tree-sitter crusher. +type fakeFallback struct{ called bool } + +func (f *fakeFallback) Name() string { return "fake_fallback" } +func (f *fakeFallback) Compress(req Request) (Result, error) { + f.called = true + return Result{Compressed: "FALLBACK", Strategy: f.Name()}, nil +} + +func TestCodeCrusher_FallbackSeam(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + fb := &fakeFallback{} + c := NewCodeCrusher() + c.Fallback = fb + + // Non-Go input routes to the injected fallback... + nonGo := "def f():\n" + strings.Repeat(" x = 1\n", 20) + res, _ := c.Compress(Request{Content: nonGo, Store: store}) + if !fb.called || res.Strategy != "fake_fallback" { + t.Fatal("non-Go input did not route to the injected Fallback") + } + // ...but Go still uses the native AST path, not the fallback. + fb.called = false + res, _ = c.Compress(Request{Content: goSource, Store: store}) + if fb.called || res.Strategy != "code_crusher" { + t.Fatal("Go input should use the AST path, not the fallback") + } +} + +func TestCodeCrusher_Deterministic(t *testing.T) { + r1, _ := NewCodeCrusher().Compress(Request{Content: goSource, Store: ccr.NewMemoryStore(ccr.MemoryConfig{})}) + r2, _ := NewCodeCrusher().Compress(Request{Content: goSource, Store: ccr.NewMemoryStore(ccr.MemoryConfig{})}) + if r1.Compressed != r2.Compressed || strings.Join(r1.Markers, ",") != strings.Join(r2.Markers, ",") { + t.Fatal("compression is not deterministic") + } +} + +func TestCodeCrusher_SmallFile_Passthrough(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + c := NewCodeCrusher() + in := "package p\n\nfunc f() int { return 1 }\n" + res, _ := c.Compress(Request{Content: in, Store: store}) + if res.Compressed != in { + t.Fatal("small file should pass through") + } +} + +func TestCodeCrusher_NilStore_Passthrough(t *testing.T) { + c := NewCodeCrusher() + res, _ := c.Compress(Request{Content: goSource, Store: nil}) + if res.Compressed != goSource { + t.Fatal("nil store must force lossless passthrough") + } +} + +func TestCodeCrusher_NeverEmptyOutput(t *testing.T) { + store := ccr.NewMemoryStore(ccr.MemoryConfig{}) + res, _ := NewCodeCrusher().Compress(Request{Content: goSource, Store: store}) + if strings.TrimSpace(res.Compressed) == "" { + t.Fatal("non-empty input compressed to empty output") + } +} diff --git a/router/router.go b/router/router.go index 73d0f74..2c9e012 100644 --- a/router/router.go +++ b/router/router.go @@ -16,6 +16,7 @@ type Router struct { yaml *crush.YAMLCrusher diff *crush.DiffCrusher search *crush.SearchCrusher + code *crush.CodeCrusher } // New builds a Router with the default compressors. @@ -27,6 +28,7 @@ func New() *Router { yaml: crush.NewYAMLCrusher(), diff: crush.NewDiffCrusher(), search: crush.NewSearchCrusher(), + code: crush.NewCodeCrusher(), } } @@ -47,9 +49,7 @@ func (r *Router) For(ct detect.ContentType) crush.Compressor { case detect.GitDiff: return r.diff case detect.SourceCode: - // TODO: dedicated AST crusher. Route to extractive text for now — - // it only drops near-duplicate lines without a query, so it is safe. - return r.text + return r.code default: return r.text } diff --git a/router/router_test.go b/router/router_test.go index 284d25a..0a4a893 100644 --- a/router/router_test.go +++ b/router/router_test.go @@ -15,7 +15,7 @@ func TestRouter_For(t *testing.T) { detect.PlainText: "text_extractive", detect.SearchResults: "search_crusher", detect.GitDiff: "diff_crusher", - detect.SourceCode: "text_extractive", + detect.SourceCode: "code_crusher", } for ct, want := range cases { if got := r.For(ct).Name(); got != want {