[Blog](https://supermemory.ai/blog/) · [Engineering](https://supermemory.ai/blog/tag/engineering/)

# Building code-chunk: AST Aware Code Chunking

By Shoubhit December 29, 2025 · 6 min read 

![Supermemory header reading "Code chunking follows the structure of your code" — Selected Lego gateway with a stepped top, doorway and connected side walls](https://supermemory.ai/_astro/cover.BX_Y5eWe_ZhuJP9.webp)

At [Supermemory](https://supermemory.ai/), we're building context engineering infrastructure for AI. A huge part of that is dealing with code: ingesting repos, understanding structure, and making it searchable. The problem is that most code chunking solutions are terrible.

We built [code-chunk](https://github.com/supermemoryai/code-chunk) to fix this. It's now the best AST-based code chunking library out there, and I want to walk through how we built it.

## The Problem with Text Splitters

Here's what most RAG pipelines do with code:

```
Source code → Split by character count → Generate embeddings → Store in vector DB
```

This is fine for prose. But code isn't prose. When you split Python at 500 characters, you end up with chunks like:

```
def calculate_total(items):
    """Calculate the total price of all items with tax."""
    subtotal = 0
    for item in items:
        subtotal += item.price * item.qu
```

Congratulations, you just cut a function in half. The embedding for this chunk has no idea what `qu` is. It doesn't know the function returns anything. If someone searches for "calculate total with tax", they might get this useless fragment. Text splitters don't understand code structure.

## Enter cAST

I found [this paper](https://arxiv.org/abs/2506.15655v2) from CMU called "cAST: Enhancing Code Retrieval-Augmented Generation with Structural Chunking via Abstract Syntax Tree". The core idea is elegant:

1. Parse code into an AST (Abstract Syntax Tree)
2. Use the tree structure to find natural split points (functions, classes, methods)
3. Chunk at semantic boundaries instead of arbitrary character limits

The paper shows solid improvements: 4.3 points on Recall@5 for RepoEval, 2.67 points on Pass@1 for SWE-bench. But reading through the paper, I realized their implementation was missing a lot of what you actually need in production.

So we took the core idea and built something better.

## How code-chunk Works

### Step 1: Parse with tree-sitter

We use [tree-sitter](https://tree-sitter.github.io/tree-sitter/) to parse source code into an AST. tree-sitter is battle-tested (it powers syntax highlighting in editors like Neovim, Helix, and Zed) and supports basically every language.

```
import { chunk } from 'code-chunk'

const chunks = await chunk('src/user.ts', sourceCode)
```

From a filepath, we detect the language and load the appropriate tree-sitter grammar. The result is a structured tree where every node has a type (function\_declaration, class\_body, import\_statement, etc.) and byte ranges.

### Step 2: Extract Entities

Here's where we diverge from the paper. We don't just traverse the tree. We extract semantic entities:

```
interface ExtractedEntity {
  type: 'function' | 'method' | 'class' | 'interface' | 'type' | 'enum' | 'import'
  name: string
  signature: string              // e.g., "async getUser(id: string): Promise<User>"
  docstring: string | null       // JSDoc, docstrings, etc.
  byteRange: { start: number; end: number }
  lineRange: { start: number; end: number }
  parent: string | null          // Which class/function contains this
}
```

For each entity, we extract:

* The full signature (not just the name)
* Any documentation comments
* Parent relationships (a method knows it belongs to a class)
* Import sources and symbols

This metadata becomes critical later.

### Step 3: Build the Scope Tree

Entities are organized into a hierarchical scope tree. A method inside a class knows its parent. A nested function knows its containing function.

```
interface ScopeNode {
  entity: ExtractedEntity
  children: ScopeNode[]
  parent: ScopeNode | null
}
```

We build this by sorting entities by byte range and using DFS to find the deepest container for each entity:

```
const findParentNode = (roots: ScopeNode[], entity: ExtractedEntity): ScopeNode | null => {
  const findInNode = (node: ScopeNode): ScopeNode | null => {
    // Does this node contain the entity?
    if (!rangeContains(node.entity.byteRange, entity.byteRange)) {
      return null
    }
    // Check children for deeper match
    for (const child of node.children) {
      const deeperMatch = findInNode(child)
      if (deeperMatch) return deeperMatch
    }
    // No child contains it, so this is the deepest
    return node
  }
  
  for (const root of roots) {
    const found = findInNode(root)
    if (found) return found
  }
  return null
}
```

This scope tree lets us answer questions like "where does this code live?" with context like `UserService > getUser`.

### Step 4: The Chunking Algorithm

The paper calls it "recursive split-then-merge". The idea is to pack as many complete syntactic units as possible into each chunk without exceeding a size limit. We do this with greedy window assignment:

```
function* greedyAssignWindows(nodes: SyntaxNode[], code: string, maxSize: number) {
  let currentWindow = { nodes: [], size: 0 }
  
  for (const node of nodes) {
    const nodeSize = getNwsCount(node) // Non-whitespace character count
    
    if (currentWindow.size + nodeSize <= maxSize) {
      // Node fits - add it
      currentWindow.nodes.push(node)
      currentWindow.size += nodeSize
    } else if (nodeSize > maxSize) {
      // Node is too big - recurse into children
      yield currentWindow
      yield* greedyAssignWindows(node.children, code, maxSize)
      currentWindow = { nodes: [], size: 0 }
    } else {
      // Node doesn't fit but isn't oversized - start new window
      yield currentWindow
      currentWindow = { nodes: [node], size: nodeSize }
    }
  }
  
  if (currentWindow.nodes.length > 0) {
    yield currentWindow
  }
}
```

We iterate through AST nodes and keep adding them to the current window until we hit the size limit. If a node is too big to fit anywhere, we recurse into its children and chunk those instead. If a node just doesn't fit in the current window but isn't oversized, we start a new window.

The paper uses non-whitespace character count instead of lines, and we do the same. Two chunks with the same line count can have wildly different amounts of actual code. A file full of blank lines and comments shouldn't count the same as dense logic.

We precompute a cumulative sum array for O(1) range queries:

```
const preprocessNwsCumsum = (code: string): Uint32Array => {
  const cumsum = new Uint32Array(code.length + 1)
  let count = 0
  for (let i = 0; i < code.length; i++) {
    if (code.charCodeAt(i) > 32) count++ // Not whitespace
    cumsum[i + 1] = count
  }
  return cumsum
}

// O(1) range query
const getNwsCount = (start: number, end: number) => cumsum[end] - cumsum[start]
```

After generating windows, we merge adjacent small ones to reduce fragmentation:

```
function* mergeAdjacentWindows(windows: Generator<ASTWindow>, maxSize: number) {
  let current: ASTWindow | null = null
  
  for (const window of windows) {
    if (!current) {
      current = window
    } else if (current.size + window.size <= maxSize) {
      current = mergeWindows(current, window)
    } else {
      yield current
      current = window
    }
  }
  
  if (current) yield current
}
```

### Step 5: Contextualized Text

Here's where code-chunk really differentiates itself. Raw chunk text isn't enough for good embeddings. You need context.

Each chunk includes a `contextualizedText` field:

```
# src/services/user.ts
# Scope: UserService
# Defines: async getUser(id: string): Promise<User>
# Uses: Database
# After: constructor

  async getUser(id: string): Promise<User> {
    return this.db.query('SELECT * FROM users WHERE id = ?', [id])
  }
```

This prepends semantic context to the raw code:

1. **File path**: where this code lives
2. **Scope chain**: what class/function contains it
3. **Entity signatures**: what's defined here
4. **Imports used**: dependencies
5. **Sibling context**: what comes before/after

Why does this matter? Embedding models are trained on natural language. When you embed `async getUser(id: string)`, the model doesn't inherently know this is inside a UserService class or that it uses a Database. By prepending this context, the embedding captures semantic relationships that pure code misses.

The implementation:

```
function formatChunkWithContext(text: string, context: ChunkContext): string {
  const parts: string[] = []
  
  if (context.filepath) {
    parts.push(`# ${context.filepath.split('/').slice(-3).join('/')}`)
  }
  
  if (context.scope.length > 0) {
    parts.push(`# Scope: ${context.scope.map(s => s.name).reverse().join(' > ')}`)
  }
  
  const signatures = context.entities
    .filter(e => e.signature && e.type !== 'import')
    .map(e => e.signature)
  if (signatures.length > 0) {
    parts.push(`# Defines: ${signatures.join(', ')}`)
  }
  
  if (context.imports.length > 0) {
    parts.push(`# Uses: ${context.imports.slice(0, 10).map(i => i.name).join(', ')}`)
  }
  
  parts.push('')
  parts.push(text)
  return parts.join('\n')
}
```

## What We Added Beyond the Paper

The cAST paper is a great foundation, but it's research code. Here's what we built on top:

**1\. Rich context extraction.** The paper just chunks. We extract full entity metadata, build scope trees, and format context for embeddings.

**2\. Overlap support.** Chunks can include the last N lines from the previous chunk. This helps with queries that target code at chunk boundaries.

```
const chunks = await chunk('file.ts', code, { overlapLines: 10 })
```

**3\. Streaming.** Process large files without loading everything into memory:

```
for await (const chunk of chunkStream('large-file.ts', code)) {
  await process(chunk)
}
```

**4\. Batch processing.** Chunk entire codebases with controlled concurrency:

```
const results = await chunkBatch(files, { 
  concurrency: 10,
  onProgress: (done, total) => console.log(`${done}/${total}`)
})
```

**5\. Effect integration.** First-class support for the [Effect](https://effect.website/) library:

```
const program = Stream.runForEach(
  chunkStreamEffect('file.ts', code),
  (chunk) => Effect.log(chunk.text)
)
```

**6\. WASM support.** Works in Cloudflare Workers and other edge runtimes:

```
import { createChunkerFromWasm } from 'code-chunk/wasm'
import treeSitterWasm from 'web-tree-sitter/tree-sitter.wasm'
import typescriptWasm from 'tree-sitter-typescript/tree-sitter-tsx.wasm'

const chunker = await createChunkerFromWasm({
  treeSitter: treeSitterWasm,
  languages: { typescript: typescriptWasm }
})
```

## Evaluating It

We built an eval harness using RepoEval data to measure how well this actually works. First run: 100% recall across all chunkers. That's not a good result. That's a broken benchmark.

### The Benchmark Was Too Easy

The original setup had a fundamental flaw. The query was a code prefix (say, lines 0-19 of a file), and the target was the same file containing that exact prefix. Embedding models trivially found the target because the query text was literally inside the target file.

This is like asking someone to find a book when you've already told them the exact title and they're standing in front of it.

### Fixing the Eval

First, we added hard negatives: 500 distractor files from the same repository. Same-repo files share vocabulary and coding style, making them genuinely hard to distinguish from the target. Before we had 108 files from the target repo. After, we had the target file plus 500 same-repo distractors.

Second, we added an Intersection over Union (IoU) threshold of 0.3\. A chunk is only "relevant" if it actually overlaps with the ground truth lines, not just somewhere in the same file.

```
Ground truth: lines 0-19 (20 lines)

Chunk A (lines 0-61):  IoU = 20/62  = 0.32 ✓ Relevant
Chunk B (lines 0-100): IoU = 20/101 = 0.20 ✗ Too bloated
Chunk C (lines 0-25):  IoU = 20/26  = 0.77 ✓ Excellent
```

This penalizes chunkers that create oversized chunks. Getting "somewhere in the file" isn't good enough. You need to actually hit the relevant code.

We also started tracking IoU directly as a quality metric. Recall is binary: did you find it or not. IoU is continuous: how well did you find it. We track `iou_at_5` and `iou_at_10`, the average best IoU score in top-k retrieved chunks.

### Results

After fixing the benchmark:

Chunker

Recall@5

IoU@5

code-chunk (AST)

70.1%

0.43

chonkie-code

49.0%

0.38

Fixed-size baseline

42.4%

0.34

The AST-based chunker creates chunks that align better with semantic boundaries. Higher IoU means the retrieved chunks actually contain the relevant code, not just code from the same file.

We also ran an agent-based eval on SWE-bench Lite comparing ops-only (Read/Grep/Glob) vs ops+semantic-search:

Duration

Tokens

Cost

Tool Calls

Without search

2.0m

4.3k

$0.25

19

With semantic search

1.2m

2.4k

$0.20

12

The semantic search agent finds relevant files faster, especially for queries that need code understanding rather than just text matching.

## Using code-chunk

Installation:

```
bun add code-chunk
# or npm install code-chunk
```

Basic usage:

```
import { chunk } from 'code-chunk'

const chunks = await chunk('src/user.ts', sourceCode)

for (const c of chunks) {
  console.log(c.text)                // Raw code
  console.log(c.contextualizedText)  // Code with context for embeddings
  console.log(c.context.scope)       // [{ name: 'UserService', type: 'class' }]
  console.log(c.context.entities)    // [{ name: 'getUser', type: 'method', ... }]
}
```

For RAG, use `contextualizedText` when generating embeddings:

```
for (const c of chunks) {
  const embedding = await embed(c.contextualizedText)
  await vectorDB.upsert({
    id: `${filepath}:${c.index}`,
    embedding,
    metadata: { filepath, lines: c.lineRange }
  })
}
```

Supported languages: TypeScript, JavaScript, Python, Rust, Go, Java.

## Wrapping Up

Chunking code well means more than respecting size limits. You need to preserve and enrich semantic context, and that's what code-chunk does.

If you're building code RAG and using naive text splitters, you're leaving performance on the table. Give code-chunk a try.

* **GitHub**: [supermemoryai/code-chunk](https://github.com/supermemoryai/code-chunk)
* **npm**: [code-chunk](https://www.npmjs.com/package/code-chunk)
* **Paper**: [cAST: Enhancing Code Retrieval-Augmented Generation](https://arxiv.org/html/2506.15655v2)

Questions? I'm on [Twitter/X](https://x.com/nexxeln).

## Other posts.

1. [How AI memory actually works. A beginner's guide How AI apps remember things between chats, why a bigger context window doesn't solve it, and the parts every memory system shares. EngineeringOct 9, 2026 ](https://supermemory.ai/blog/how-ai-memory-works/)
2. [We're open sourcing the company brain. Here's how we designed the multiplayer harness Company Brain is now open source. A walkthrough of the multiplayer harness behind its Slack experience, from proactivity and memory boundaries to approvals and recovery. NewsSep 25, 2026 ](https://supermemory.ai/blog/open-sourcing-company-brain/)
3. [Jev changes a lot in memory & context engineering. Here's exactly how. We tested Jev across reranking, chunking, observation, and harness decisions. Here is where fast decision models help memory systems, where they cost more, and where they still fall short. EngineeringSep 24, 2026 ](https://supermemory.ai/blog/jev-memory-context-engineering/)
4. [I reverse-engineered Instinct's memory. Here's exactly how it works Instinct keeps its memory as git-tracked markdown files, found with grep rather than vectors. Here is the whole system as far as black-box probing can reconstruct it, and how to rebuild it on supermemory in about 60 lines. EngineeringSep 20, 2026 ](https://supermemory.ai/blog/reverse-engineering-instinct-memory/)
5. [An update to supermemory We've discontinued the supermemory company brain and Nova. Everyone who was charged has been refunded, our MCP and plugins continue to run, and we're going all in on the memory engine. NewsSep 10, 2026 ](https://supermemory.ai/blog/an-update-to-supermemory/)
6. [Scaling Conversations: How Adapta Grew Usage Without Losing Context Adapta added Supermemory as a persistent memory layer so every conversation keeps its context — letting the team scale usage without losing the thread. Case StudyJun 12, 2026 ](https://supermemory.ai/blog/adapta-scaling-conversations/)
7. [How Chatarmin Ditched RAG and Went Memory-Only with Supermemory Chatarmin replaced a heavy RAG pipeline with Supermemory's memory layer — cutting average AI response time from 40s to 12s and token usage by 40–50%. Case StudyJun 10, 2026 ](https://supermemory.ai/blog/chatarmin-memory-only/)
8. [SMFS: making agentic retrieval 55% cheaper AND more accurate We launched SMFS.ai (Supermemory Filesystem) a few weeks ago, with a simple bet: We can redesign the filesystem specifically for agents, with special files, structures, and commands that it can use for it's tasks. Today, SMFS is used by hundreds of companies to power their agents. EngineeringMay 28, 2026 ](https://supermemory.ai/blog/smfs-making-agentic-retrieval-55-cheaper-and-more-accurate/)
9. [Introducing Dynamic Dreaming: supermemory now connects the dots, for you. Dreaming is magical. TLDR: We're launching Dynamic Dreaming in supermemory today, which automatically works if you're using supermemory in any way - API, OpenClaw, Hermes agent, etc. EngineeringMay 25, 2026 ](https://supermemory.ai/blog/introducing-dynamic-dreaming-supermemory-now-connects-the-dots-for-you/)
10. [Dear reader, we just made supermemory insanely cheap... the Context Cloud When I first started building supermemory, I had one goal: To build the best memory system for AI. I would talk to customers, and find out that memory was not the only thing they needed - They were all setting up 7-8 different vendors at the same time. EngineeringMay 18, 2026 ](https://supermemory.ai/blog/dear-reader-we-just-made-supermemory-insanely-cheap-the-context-cloud/)
11. [Introducing @supermemory/tools v2.0.0 Today we're releasing v2.0.0\. This release unifies the API across all agents sdk integrations from AI SDK to Mastra, makes conversation identity a first-class concept, and ships with memory saving on by default. EngineeringApr 27, 2026 ](https://supermemory.ai/blog/introducing-supermemory-tools-v2-0-0/)
12. [Solving the Precision-Recall Tradeoff: Search Result Aggregation When you're building memory for AI, search is your foundational layer. The way search generally works is straightforward: the user defines a query, and then sets a limit (top-K) on how many search results they want returned. Usually, this is set to 10 or 20. EngineeringApr 5, 2026 ](https://supermemory.ai/blog/solving-the-precision-recall-tradeoff-search-result-aggregation/)
13. [OpenClaw Memory Problems: Why It Forgets and How to Fix It (2026) TLDR: Today, we are releasing a new version of our openclaw plugin - https://github.com/supermemoryai/openclaw-supermemory. This post is going to be a bit technical, so bear with me (or bookmark for later!) In this post, I will talk about what we do about OpenClaw memory, and how we fix it. EngineeringFeb 19, 2026 ](https://supermemory.ai/blog/why-everyone-is-complaining-about-openclaws-memory-it-sucks-and-why-supermemory-fixes-it/)
14. [Stateful Coding Agents with Memory: Build Long-Running Agents (2026) We built a plugin for Claude Code and OpenCode that gives your coding agent persistent memory. It remembers your preferences, learns your codebase, and never loses context mid-conversation. The result is an agent you can run for months without starting over. EngineeringFeb 18, 2026 ](https://supermemory.ai/blog/infinitely-running-stateful-coding-agents/)
15. [Clawd / Molt bot's memory SUCKS. We gave it supermemory. I'm the founder of supermemory. Clawd/Molt bot is blowing up right now, with many, many use cases. I set it up, too, and have been using it through telegram. TLDR: just go to https://supermemory.ai/docs/integrations/clawdbot to set up supermemory for your clawd bot. EngineeringJan 28, 2026 ](https://supermemory.ai/blog/clawd-molt-bots-memory-sucks-we-gave-it-supermemory/)
16. [Catch up with our UNFORGETTABLE Launch Week Over the last year, one belief has guided almost everything we’ve built at Supermemory AI becomes meaningfully useful only when it remembers. Memory shouldn’t be something developers rebuild from scratch. It shouldn’t be fragile, expensive, or trapped inside a single tool. NewsJan 4, 2026 ](https://supermemory.ai/blog/catch-up-with-our-unforgettable-launch-week/)
17. [Empowering the Next Generation of Founders: Supermemory Startup Program If there’s one thing we’ve learned while building Supermemory, it’s that most startups don’t fail because they didn't build features; they fail when infrastructure slows them down, or they built too slow. NewsDec 31, 2025 ](https://supermemory.ai/blog/empowering-the-next-generation-of-founders-supermemory-startup-program/)
18. [Supermemory raises $3 million with the best memory engine for LLMs Today, I am excited to announce our first funding round to accelerate our mission of building an interoperable, scalable and reliable memory for LLMs and agents. Memory is one of the hardest challenges in AI right now. NewsOct 6, 2025 ](https://supermemory.ai/blog/supermemory-raises-3-million-and-building-the-best-memory-engine-for-llms/)
19. [Mem0 vs Supermemory: Why Scira Switched Scira AI moved its production memory layer from Mem0 to Supermemory. This is what failed, what improved, and how the team evaluated the two systems. Case StudyOct 2, 2025 ](https://supermemory.ai/blog/why-scira-ai-switched/)
20. [Never Record Again: How Montra Uses Supermemory to Rethink Video Creation Campbell Baron, the founder of Montra, has been making videos since he was twelve. By thirteen, he was already doing brand work. Today, he’s betting on a very different future for creators: a world where recording is the exception, and most videos are generated from scratch. Case StudyAug 21, 2025 ](https://supermemory.ai/blog/never-record-again-how-montra-uses-supermemory-to-rethink-video-creation/)
21. [Unified Memory That Works Where You Work: Your Second Brain With Supermemory Hi everyone, I’m Dhravya, the founder of Supermemory. I want to start with a little story behind why this product means so much to me. You can also skip straight to what it is and how it works below. EngineeringJul 25, 2025 ](https://supermemory.ai/blog/unified-memory-that-works-where-you-work-your-second-brain-with-supermemory/)
22. [Supermemory just got faster on PlanetScale What is Supermemory? Supermemory completes the missing part of the LLM puzzle: memory. Just as memory is crucial for human intelligence, it's essential for truly intelligent AI systems. EngineeringJul 18, 2025 ](https://supermemory.ai/blog/supermemory-just-got-faster-on-planetscale/)
23. [Faster, smarter, reliable infinite chat: Supermemory IS context engineering. People are obsessed with prompts and prompt engineering. Sure, what you say is important, but what the model knows when you say it is the difference between a stateless text generator and an intelligent AI system. In short, context is the most crucial component. NewsJul 9, 2025 ](https://supermemory.ai/blog/faster-smarter-reliable-infinite-chat-supermemory-is-context-engineering/)
24. [We solved AI API interoperability One API to rule them all, One spec to find them, One library to bring them all and in the TypeScript, bind them. When we were building the Infinite Chat API, initially, we only supported the OpenAI format. This was fine, until a lot of our customers started asking for more. EngineeringJul 7, 2025 ](https://supermemory.ai/blog/we-solved-ai-api-interoperability/)
25. [The Wow Factor of Memory - How Flow Used Supermemory To Build Smarter, Stickier Products Overview: Flow is a note-taking app built around a bold vision: to create a more personal, context-aware writing experience powered by AI. At the heart of this mission is memory. Case StudyJun 14, 2025 ](https://supermemory.ai/blog/the-wow-factor-of-memory-how-flow-used-supermemory-to-build-smarter-stickier-products/)
26. [The UX and technicalities of awesome MCPs Last month, we launched the Supermemory MCP, mostly to test our own infrastructure and get some initial traction. It blew up. To my absolute surprise, the initial launch itself got half a million impressions (!!!). Then, we launched and got #2 on ProductHunt too. EngineeringJun 8, 2025 ](https://supermemory.ai/blog/the-ux-and-technicalities-of-awesome-mcps/)
27. [Architecting a memory engine inspired by the human brain Language is at the heart of intelligence, but what truly powers meaningful interaction is memory — the ability to accumulate, recall, and contextualize information over time. Large Language Models (LLMs) have mastered language, but memory remains their Achilles’ heel. EngineeringJun 5, 2025 ](https://supermemory.ai/blog/memory-engine/)

## Start building with supermemory.

Memory and continual learning for any model, any harness. Available through our API, plugins, and MCP.

[Build with supermemory ](https://console.supermemory.ai/)