Query-Aware Code Compression for LLMs: 1.3M LOC to 1,500 Lines in 2.8s

The Problem

You want to ask an LLM a question about your codebase. You have three options:

1. Paste everything — Repomix dumps your entire repo. For Next.js (1.29M LOC), that's 80k+ tokens. Expensive, slow, and the model loses the answer somewhere in the middle.

2. Truncate — Cut at 20k tokens. Hope the relevant code is in the first 400 files alphabetically. It usually isn't.

3. RAG — Embed your code, retrieve chunks by similarity. You get 8k tokens of tangentially related fragments, split mid-function, missing the imports and types you need to understand them.

None of these are query-aware at the AST level. If you ask "where is authentication?", you want the authenticate() function, its callers, its types — not a random embedding-similar chunk that mentions "auth" in a comment.

CodeShrink

codeshrink "where is authentication?" --path ./next.js

Output: 1,509 lines from 1.29M LOC. 99.9% compression. 2.8 seconds. Only auth-related functions with their signatures, imports, and 5 lines of context.

The key idea: use Tree-sitter ASTs, not embeddings. Parse the code into symbols (functions, classes, types), match them against the query, rank them, and extract the top-K with context.

How It Works

Step 1: Scan

Tree-sitter parses every file into an AST in parallel (via rayon). We support 7 languages: TypeScript, TSX, JavaScript, Python, Rust, Go, Java.

From each AST, we extract symbol definitions: function name, start/end line, signature (parameter types, return type), and the file path.

For a 78k LOC repo (Fastify), this takes ~40ms.

Step 2: Match

Your query is split into terms. Each term gets expanded with a semantic group:

• "auth" → also matches "login", "token", "jwt", "session", "credential"
• "route" → also matches "handler", "endpoint", "path", "middleware"
• "database" → also matches "db", "query", "connection", "pool", "migration"

This is a static mapping, not an embedding model. It covers the common synonyms that matter in code.

Step 3: Rank

Each symbol gets a score from multiple signals:

• Name match: does the function/class name contain a query term? (strongest signal)
• Path match: is the file in a directory named after the query? (src/auth/ for "auth")
• Signature match: do parameter types or return types match?
• 1-hop expansion: if a high-scoring function calls another function, that callee gets a boost

Scores are combined with configurable weights. The default ranking produces good results without tuning.

Step 4: Extract

Top-K symbols (default 50) are extracted with N lines of context above and below (default 5). When two symbols overlap or are adjacent, their ranges merge into one block.

Step 5: Output

Formatted as Markdown (default), XML (for LLM system prompts), or plain text. Each block includes the file path, line numbers, and the symbol's score.

Benchmarks

All measured on Apple Silicon, release build, single run (no warm-up caching):

Repo	Query	Input LOC	Output LOC	Compression	Latency
Express (21k LOC)	"middleware"	21,475	628	97.1%	28ms
Express (21k LOC)	"routing"	21,475	655	96.9%	25ms
Fastify (78k LOC)	"route handler"	77,959	1,454	98.1%	93ms
Fastify (78k LOC)	"plugin"	77,959	721	99.1%	68ms
Next.js (1.29M LOC)	"server action"	1,294,421	1,509	99.9%	2,823ms
Next.js (1.29M LOC)	"middleware"	1,294,421	1,709	99.9%	2,253ms

The latency is dominated by Tree-sitter parsing on large repos. For typical project sizes (10-100k LOC), it's under 100ms.

vs. Alternatives

	CodeShrink	Repomix	CodeGraph	Truncation	RAG
Query-aware	Yes	No (full dump)	Yes (MCP)	No	Partial
Standalone CLI	Yes	Yes	No (MCP server)	N/A	No
Latency (78k LOC)	93ms	~500ms	~2s	0ms	~200ms
Output (78k LOC)	1.4k lines	78k lines	~2k lines	20k tokens	~8k tokens
Dependencies	0 (single binary)	Node.js	Python + SQLite	N/A	Embeddings model
npm package	Yes	Yes	No	N/A	Varies

Usage

Install:

# Rust CLI
cargo install codeshrink

# or npm
npm install codeshrink

CLI:

# Basic query
codeshrink "where is the database connection?" --path ./my-project

# Narrow context
codeshrink "error handling" -c 2

# XML output for LLM system prompts
codeshrink "routing" --format xml

# Pipe into clipboard
codeshrink "auth" --path ./app | pbcopy

As a Node.js library (napi-rs bindings):

const { shrink } = require('codeshrink');

const result = shrink('authentication', '/path/to/repo', {
    contextLines: 5,
    maxSymbols: 50,
    format: 'markdown',
});

console.log(result.stats);
// { filesScanned: 141, symbolsReturned: 50,
//   inputLines: 21475, outputLines: 628,
//   compressionRatio: 0.971 }

As a Rust library:

use codeshrink_core::{shrink, ShrinkOptions};
use std::path::Path;

let result = shrink(
    "where is authentication?",
    Path::new("./my-project"),
    &ShrinkOptions::default(),
)?;

println!("{}", result.compressed);

Limitations

• Semantic grouping is static. The synonym map covers common code terms but won't know that "billing" relates to "stripe" in your codebase. A project-specific config is on the roadmap.
• No cross-file data flow. 1-hop expansion follows function calls within the same file. Cross-file call graph analysis would improve recall but adds complexity.
• 7 languages. Tree-sitter grammars exist for 100+ languages, but each needs a custom symbol extractor. Adding a new language takes ~50 lines of Rust.

What's Next

• MCP server for direct LLM tool use
• Cross-file call graph analysis
• Project-specific synonym configs

MIT/Apache-2.0: github.com/TimurRakhmatullin86/codeshrink

What queries would you run on your codebase? What output format works best for your LLM workflow?

Comments

[Arsen Apostolov]

The "not a random embedding-similar chunk that mentions auth in a comment" point is the whole game. Semantic similarity and structural relevance are different things, and RAG pipelines quietly conflate them — you retrieve the comment that says "handles auth" instead of the function that does.

I hit the downstream version of this problem when I run code-heavy local models: the compression wins you get here are the exact wins that decide whether a 27B on a 2x3090 can actually answer "where is X handled" in a big repo. The context that survives compression is what has to fit in VRAM, and if you feed it the wrong 1,500 lines you've just moved the hallucination upstream. What I ended up doing is treating the compressed context as a benchmarkable input — same query, different compression strategy, measure whether the model still lands on the right function. homelab-monitor's benchmark tab is where I run those comparisons, because the whole point is that "smaller context" only helps if the answer quality holds, and you can't tell by eye.

[Tim]

Exactly — "semantic similarity" and "structural relevance" are different things, and most RAG pipelines treat them as the same. An embedding that scores high on "handles auth" could be a comment, a test name, or the actual function body. AST extraction guarantees you get complete, callable code.

The benchmarkable-input approach is interesting. Right now CodeShrink reports compression ratio and latency but doesn't measure answer quality downstream — whether the compressed context actually helps the model answer correctly. That's the metric that matters.

For local models with constrained VRAM, the format matters too. The XML output mode (--format xml) with explicit file paths and line numbers gives more structure for the model to anchor on than raw markdown. I've seen better results with Llama models when the context has clear section boundaries.

What context sizes are you typically working with on the 2x3090 setup? That would help me calibrate the default --max-symbols parameter.