# How Headroom's CodeCompressor Uses AST and Tree-Sitter for Code Compression

> Discover how Headroom CodeCompressor leverages AST and Tree-Sitter to analyze code structure, calculate importance scores, and compress function bodies for syntactically valid output.

- Repository: [Tejas Chopra/headroom](https://github.com/chopratejas/headroom)
- Tags: how-to-guide
- Published: 2026-06-13

---

**Headroom's CodeCompressor parses source code into an Abstract Syntax Tree (AST) using Tree-Sitter, analyzes structural elements to compute semantic importance scores, and selectively compresses function bodies while guaranteeing syntactically valid output.**

The `CodeAwareCompressor` class in the `chopratejas/headroom` repository implements a language-agnostic compression pipeline that treats code as structured data rather than plain text. By leveraging Tree-Sitter to build ASTs across eight major programming languages, Headroom's CodeCompressor preserves critical interface elements—imports, type definitions, and function signatures—while intelligently truncating implementation details. This approach ensures that compressed code remains syntactically valid and semantically useful for LLM context windows.

## Language Detection via Tree-Sitter Error Analysis

The compression pipeline begins with robust language detection in [`headroom/transforms/code_compressor.py`](https://github.com/chopratejas/headroom/blob/main/headroom/transforms/code_compressor.py). The `detect_language` method first applies a lightweight regex pre-filter to shortlist candidate languages, then parses the code with Tree-Sitter for each candidate.

It selects the language that yields the **fewest error nodes** in the resulting AST (see lines 63-89 and 97-104). This error-counting heuristic ensures that even with ambiguous file extensions or mixed-language snippets, the parser chooses the grammar that best fits the actual syntax.

## Thread-Local Parser Caching

Tree-Sitter parsers are **lazy-loaded** and stored in a thread-local cache (`_tree_sitter_local`) to avoid concurrency issues. In [`headroom/transforms/code_compressor.py`](https://github.com/chopratejas/headroom/blob/main/headroom/transforms/code_compressor.py), the helper `_get_parser` (lines 14-24 and 25-33) creates a `tree_sitter.Parser`, sets its language via `tree_sitter_language_pack.get_language`, and caches it per thread.

This design prevents the "unsendable" crash that occurs when a `Parser` is shared across threads, ensuring thread safety without the overhead of recreating parsers for every compression request.

## Data-Driven Language Configuration

Rather than hardcoding language-specific logic, the compressor uses a single data structure called `_LANG_CONFIGS` (lines 29-87 in [`headroom/transforms/code_compressor.py`](https://github.com/chopratejas/headroom/blob/main/headroom/transforms/code_compressor.py)). This configuration maps AST node types to semantic categories—imports, functions, classes, type definitions, and comment syntax—across all supported languages.

This declarative approach eliminates per-language extraction code and enables uniform processing of Python, JavaScript, TypeScript, Go, Rust, Java, C, and C++.

## Extracting Code Structure from the AST

The private method `_extract_structure` (lines 158-200 and 212-270) performs a depth-first traversal of the AST, matching node types against `_LANG_CONFIGS`. It builds a `CodeStructure` object containing:

- Import statements and package declarations
- Type definitions and class declarations
- Function signatures and their original bodies
- Decorators and export statements
- Top-level code and comments

Export statements and decorated definitions receive special handling to preserve their surrounding syntax and semantic relationships.

## Semantic Importance Analysis

After extraction, the compressor computes a **semantic importance score** for each symbol (functions, methods, classes). The analysis collects definition nodes, identifier usage patterns, call relationships, and reference counts, then normalizes these signals to a 0-1 scale (lines 311-370).

These scores determine how much of each function body survives compression, ensuring that heavily-referenced utility functions retain more implementation detail than rarely-used helpers.

## Allocating Line Budgets and Compressing Bodies

### Budget Distribution

The `_allocate_body_budget` method (lines 69-106) distributes a global line budget—derived from the target compression rate—across all functions proportionally to their importance scores. High-scoring symbols receive larger line allocations, while low-scoring functions may be reduced to a single placeholder.

### Function Body Compression

The `_compress_function_ast` method (lines 494-523) receives a function node, its language configuration, and the allotted line limit. It rebuilds the function signature unchanged to preserve the interface, then truncates or compresses the body to fit within the allocated budget. The compressed fragments are collected for final reassembly.

## Reassembly and Syntax Validation

### Code Reconstruction

The `_assemble_compressed` method (lines 445-461) concatenates the preserved structural elements: imports, type definitions, class definitions, function signatures (with their compressed bodies), and any remaining top-level code. This guarantees that the output maintains proper declaration order and scoping.

### Syntax Verification

To ensure reliability, the compressor runs a **syntax-verification pass** using `_verify_syntax` (lines 512-525). It re-parses the generated code with Tree-Sitter; if any error nodes remain, the method returns the original code unchanged. This safety mechanism guarantees that Headroom never serves syntactically broken source code.

## Fallback Mechanisms

When Tree-Sitter is unavailable or the AST path fails, the compressor optionally falls back to the **Kompress** model or a regex-only compressor. This behavior is controlled by the `fallback_to_kompress` setting in `CodeCompressorConfig` (lines 73-84).

## Usage Examples

### Basic Compression

```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
        processed = item.strip().lower()
        results.append(processed)
    return results
''')
print(result.compressed)

```

Output:

```python
import os
from typing import List

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

    pass

```

The `result.syntax_valid` property always returns `True` due to the verification pass.

### Language Hint

When you know the language in advance, provide it to bypass auto-detection:

```python
result = compressor.compress(js_code, language="javascript")

```

### Disabling Tree-Sitter

To force fallback behavior when Tree-Sitter is unavailable:

```python
from headroom.transforms import CodeCompressorConfig

cfg = CodeCompressorConfig(fallback_to_kompress=False)
compressor = CodeAwareCompressor(config=cfg)
result = compressor.compress(python_code)  # Returns original if tree-sitter missing

```

## Summary

- **Tree-Sitter Integration**: Headroom's CodeCompressor uses Tree-Sitter to parse source code into ASTs, enabling language-aware compression across eight programming languages.
- **Error-Based Detection**: The `detect_language` method selects the best grammar by minimizing AST error nodes, ensuring accurate parsing even for ambiguous code.
- **Thread-Safe Parsing**: Parsers are cached in thread-local storage via `_tree_sitter_local` to prevent concurrency crashes.
- **Semantic Scoring**: Importance scores derived from AST analysis (definition nodes, call relationships) drive intelligent line-budget allocation.
- **Guaranteed Validity**: The `_verify_syntax` method re-parses compressed output to ensure零 error nodes before returning results.
- **Graceful Degradation**: Configurable fallbacks to the Kompress model or regex-only compression when Tree-Sitter is unavailable.

## Frequently Asked Questions

### What is Tree-Sitter and why does Headroom use it for code compression?

Tree-Sitter is a parser generator tool that builds concrete syntax trees for source code with error recovery. Headroom uses it because AST-based compression preserves the structural relationships between code elements—imports, class definitions, and function signatures—while allowing safe removal of implementation details. This ensures compressed output remains syntactically valid and semantically coherent, unlike naive text-based truncation.

### How does Headroom ensure compressed code remains syntactically valid?

The compressor performs a two-pass validation: first, it builds an AST from the original code; second, after reassembly, it runs `_verify_syntax` to parse the compressed output. If Tree-Sitter detects any error nodes in the final result, the compressor discards the compressed version and returns the original source unchanged. This guarantees that broken syntax never reaches the output.

### What happens if Tree-Sitter is not installed?

If Tree-Sitter is unavailable, the `CodeAwareCompressor` checks the `fallback_to_kompress` configuration in `CodeCompressorConfig`. When enabled, it falls back to the Kompress model or a regex-based compressor. If disabled, it returns the original code unchanged. This ensures the system degrades gracefully without crashing when dependencies are missing.

### Which programming languages does Headroom's CodeCompressor support?

The compressor supports eight languages through the `_LANG_CONFIGS` data structure: **Python**, **JavaScript**, **TypeScript**, **Go**, **Rust**, **Java**, **C**, and **C++**. Each language is defined by its Tree-Sitter grammar and a mapping of AST node types to semantic categories (imports, functions, classes, etc.), allowing uniform processing across diverse syntaxes.