# How CodeCompressor Uses AST Parsing for Different Programming Languages in Headroom

> Discover how Headroom's CodeCompressor employs AST parsing with tree-sitter for multiple programming languages, ensuring syntactic validity and efficient code compression. Learn how it works.

- Repository: [Tejas Chopra/headroom](https://github.com/chopratejas/headroom)
- Tags: deep-dive
- Published: 2026-06-18

---

**Headroom’s CodeAwareCompressor leverages tree-sitter parsers with thread-local caching and a data-driven `LangConfig` table to detect languages and preserve syntactic validity while compressing code across Python, JavaScript, Go, Rust, and other supported languages.**

The `CodeAwareCompressor` class in [`headroom/transforms/code_compressor.py`](https://github.com/chopratejas/headroom/blob/main/headroom/transforms/code_compressor.py) intelligently reduces source code size without breaking syntax by walking abstract syntax trees (ASTs). Unlike naive truncation methods, it uses language-specific tree-sitter grammars to identify imports, function signatures, and class definitions before selectively trimming function bodies based on symbol importance.

## Language Detection Pipeline

Before parsing, the compressor must identify which programming language the source code represents. The `detect_language` function (lines 22‑108) implements a three-stage detection pipeline that balances speed and accuracy.

### Regex Prefilter Stage

The compressor first runs a fast regex pre-filter to narrow down candidate languages. It iterates over `_LANGUAGE_PREFILTER` patterns for each supported language and assigns initial scores based on characteristic syntax markers:

```python
for lang, patterns in _LANGUAGE_PREFILTER.items():
    # Score based on regex matches for keywords like "def ", "import ", "func "

```

This stage quickly eliminates obvious mismatches before invoking the heavier tree-sitter parser.

### Tree-Sitter Disambiguation

When multiple candidates remain (for example, TypeScript versus JavaScript, or C++ versus C), the compressor loads a tree-sitter parser for each candidate via `_get_parser(lang.value)` and parses a slice of the code. The language with the fewest `ERROR` or `MISSING` nodes in its AST wins. The confidence score derives from the error rate, ensuring the parser selects the grammar that most accurately represents the source structure.

If tree-sitter is not installed, the function falls back to regex-only scoring.

## Thread-Local Parser Management

Tree-sitter parsers are unsendable objects that must remain on the thread that created them. Headroom solves this concurrency constraint using thread-local storage:

```python
_tree_sitter_local = threading.local()

```

The `_get_parser` function (lines 76‑136) checks this thread-local cache before creating a new parser instance. When creation is necessary, it loads the grammar via `tree_sitter_language_pack.get_language()` and stores the configured parser in the thread-local dictionary. This approach allows safe reuse of heavy grammar files across worker threads without cross-thread serialization issues.

## Data-Driven Language Configuration

Rather than hardcoding per-language extraction logic, the compressor uses a single `LangConfig` dataclass (lines 203‑226) to define AST node types for each supported language. The `_LANG_CONFIGS` dictionary (lines 229‑319) maps each `CodeLanguage` enum to its configuration:

```python
CodeLanguage.PYTHON: LangConfig(
    import_nodes=frozenset({"import_statement", "import_from_statement"}),
    function_nodes=frozenset({"function_definition"}),
    class_nodes=frozenset({"class_definition"}),
    type_nodes=frozenset({"type_alias_statement"}),
    body_node_types=frozenset({"block"}),
    decorator_node="decorated_definition",
    comment_prefix="#",
    uses_colon_after_signature=True,
    detection_hints=("def ", "import ", "from ", "class ", "async def"),
)

```

All other languages—including JavaScript, TypeScript, Go, Rust, Java, C, and C++—follow the same pattern. This data-driven approach allows the generic AST visitor to handle any supported language by simply looking up the appropriate node types in the configuration table.

## AST Walking and Structure Extraction

Once the language is detected and the parser loaded, the compressor extracts code structure using a recursive AST visitor.

### Structure Extraction

The `_extract_structure` function (lines 161‑277) walks the tree, comparing each `node.type` against the language-specific sets in `LangConfig` (`import_nodes`, `function_nodes`, `class_nodes`, etc.). It captures source text via `_get_node_text(node, code)` and dispatches to `_compress_function_ast` or `_compress_class_ast` for functions and classes. Because the visitor relies exclusively on the `LangConfig` tables, the same extraction logic works uniformly across all languages.

### Symbol Importance Analysis

Before trimming bodies, the compressor analyzes symbol importance via `_analyze_symbol_importance` (lines 108‑158). This function collects:

- **Definition occurrences** where symbols are declared
- **Identifier frequencies** showing how often names are used
- **Call relationships** mapping function invocations
- **Visibility modifiers** distinguishing public from private APIs

It calculates a raw importance score per symbol and normalizes values to a 0‑1 range. These scores feed into `_allocate_body_budget` (lines 68‑115), which determines how many lines of each function body the compressor may retain.

### Re-assembly and Syntactic Validity

The `_assemble_compressed` function reconstructs the final output by concatenating preserved imports, signatures, decorators, type definitions, and allocated body lines. The output guarantees syntactic validity because:

- All import statements and function signatures remain intact
- Class definitions and type annotations are preserved in full
- Only whole function bodies are trimmed according to the computed budget

## Content Router Integration

The `ContentRouter` class in [`headroom/transforms/content_router.py`](https://github.com/chopratejas/headroom/blob/main/headroom/transforms/content_router.py) lazy-loads the compressor only when processing source-code content:

```python
if self._code_compressor is None:
    from .code_compressor import CodeAwareCompressor, _check_tree_sitter_available
    self._code_compressor = CodeAwareCompressor()

```

This lazy loading (lines 1683‑1690) ensures that the heavy tree-sitter dependency is imported only when needed, keeping the rest of the transformation pipeline lightweight.

## Practical Usage Example

The following example demonstrates automatic language detection and compression for Python code:

```python
>>> from headroom.transforms import CodeAwareCompressor
>>> compressor = CodeAwareCompressor()
>>> result = compressor.compress("""\
... import os
... from typing import List
...
... def process_data(items: List[str]) -> List[str]:
...     \"\"\"Process a list of items.\"\"\"
...     results = []
...     for item in items:
...         if not item:
...             continue
...         results.append(item.strip().lower())
...     return results
... """)
>>> print(result.compressed)
import os
from typing import List

def process_data(items: List[str]) -> List[str]:
    """Process a list of items."""
    # ... (body compressed: 8 lines → 2 lines)

    pass
>>> print(result.syntax_valid)   # always True

True

```

The same `compress()` call works for JavaScript, Go, Rust, and other supported languages; the compressor automatically selects the correct parser based on the detection step.

## Summary

- **Language detection** combines regex pre-filtering with tree-sitter parsing to select the grammar with the fewest errors.
- **Thread-local storage** in `_get_parser` ensures parsers remain on their creating threads while enabling reuse across compressions.
- **`LangConfig` tables** provide a data-driven approach to AST node identification, eliminating the need for per-language extraction code.
- **Symbol importance analysis** allocates line budgets to function bodies based on usage frequency and visibility.
- **Syntactic validity** is guaranteed by keeping imports, signatures, and decorators intact while trimming only function internals.

## Frequently Asked Questions

### What programming languages does CodeCompressor support?

CodeCompressor supports Python, JavaScript, TypeScript, Go, Rust, Java, C, and C++. Each language is defined in the `_LANG_CONFIGS` dictionary in [`headroom/transforms/code_compressor.py`](https://github.com/chopratejas/headroom/blob/main/headroom/transforms/code_compressor.py), which maps AST node types to semantic categories like imports, functions, and classes.

### How does CodeCompressor maintain thread safety with tree-sitter parsers?

The compressor stores parser instances in a `threading.local()` cache. The `_get_parser` function checks this thread-local storage before creating new parsers, ensuring that unsendable tree-sitter objects never migrate between threads while still allowing efficient reuse within a single worker thread.

### Why does CodeCompressor use tree-sitter instead of regular expressions?

Tree-sitter provides accurate AST representations that distinguish between comments, string literals, and actual code structure. This allows the compressor to identify function boundaries, import statements, and class definitions reliably—tasks that regex cannot perform safely across nested structures and language-specific syntax variations.

### Is the compressed output guaranteed to be syntactically valid?

Yes. The compressor preserves all import statements, function signatures, class definitions, decorators, and type annotations in their original form. It only truncates function bodies according to calculated line budgets, ensuring the resulting code remains parseable by the original language's compiler or interpreter.