Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

cAST: Enhancing Code Retrieval-Augmented Generation with Structural Chunking via Abstract Syntax Tree

Yilin Zhang, Xinran Zhao, Zora Zhiruo Wang, Chenyang Yang, Jiayi Wei, Tongshuang Wu

TL;DR

cAST tackles the chunking bottleneck in retrieval-augmented code generation by using Abstract Syntax Tree structures to produce semantically coherent chunks. It introduces a split-then-merge AST-based chunking algorithm, with chunk sizing defined by non-whitespace character count, and demonstrates consistent retrieval and generation gains across RepoEval, CrossCodeEval, and SWE-bench with diverse retrievers and language models. The results show that structure-aware chunking improves retrieval precision and recall and yields better code generation, especially when paired with code-specific retrievers, highlighting the importance of syntactic-aware preprocessing for scalable code intelligence. The work provides a practical, drop-in improvement for RAG pipelines and outlines avenues for future enhancement, including multi-view inputs and execution-aware chunking.

Abstract

Retrieval-Augmented Generation (RAG) has become essential for large-scale code generation, grounding predictions in external code corpora to improve actuality. However, a critical yet underexplored aspect of RAG pipelines is chunking -- the process of dividing documents into retrievable units. Existing line-based chunking heuristics often break semantic structures, splitting functions or merging unrelated code, which can degrade generation quality. We propose chunking via Abstract Syntax Trees (\ourwork), a structure-aware method that recursively breaks large AST nodes into smaller chunks and merges sibling nodes while respecting size limits. This approach generates self-contained, semantically coherent units across programming languages and tasks, improving performance on diverse code generation tasks, e.g., boosting Recall@5 by 4.3 points on RepoEval retrieval and Pass@1 by 2.67 points on SWE-bench generation. Our work highlights the importance of structure-aware chunking for scaling retrieval-enhanced code intelligence.

cAST: Enhancing Code Retrieval-Augmented Generation with Structural Chunking via Abstract Syntax Tree

TL;DR

cAST tackles the chunking bottleneck in retrieval-augmented code generation by using Abstract Syntax Tree structures to produce semantically coherent chunks. It introduces a split-then-merge AST-based chunking algorithm, with chunk sizing defined by non-whitespace character count, and demonstrates consistent retrieval and generation gains across RepoEval, CrossCodeEval, and SWE-bench with diverse retrievers and language models. The results show that structure-aware chunking improves retrieval precision and recall and yields better code generation, especially when paired with code-specific retrievers, highlighting the importance of syntactic-aware preprocessing for scalable code intelligence. The work provides a practical, drop-in improvement for RAG pipelines and outlines avenues for future enhancement, including multi-view inputs and execution-aware chunking.

Abstract

Retrieval-Augmented Generation (RAG) has become essential for large-scale code generation, grounding predictions in external code corpora to improve actuality. However, a critical yet underexplored aspect of RAG pipelines is chunking -- the process of dividing documents into retrievable units. Existing line-based chunking heuristics often break semantic structures, splitting functions or merging unrelated code, which can degrade generation quality. We propose chunking via Abstract Syntax Trees (\ourwork), a structure-aware method that recursively breaks large AST nodes into smaller chunks and merges sibling nodes while respecting size limits. This approach generates self-contained, semantically coherent units across programming languages and tasks, improving performance on diverse code generation tasks, e.g., boosting Recall@5 by 4.3 points on RepoEval retrieval and Pass@1 by 2.67 points on SWE-bench generation. Our work highlights the importance of structure-aware chunking for scaling retrieval-enhanced code intelligence.

Paper Structure

This paper contains 34 sections, 2 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. cAST
  3. Design Goal.
  4. AST Parsing.
  5. AST-based Recursive Chunking.
  6. Chunk size metric.
  7. Experiments
  8. Experiment Settings
  9. Datasets.
  10. Metrics.
  11. Retrieval and Generation Models.
  12. cAST Results and Analysis
  13. Retrieval.
  14. Generation.
  15. Correlation between retrieval and generation performance.
  16. ...and 19 more sections

Figures (2)

  • Figure 1: Syntax-agnostic chunking often omits crucial information needed to generate functional code. In this example, fixed-size chunking breaks the structure of the $\texttt{compute\_stats}$ method, causing the model to lose context regarding its return value. As a result, the model generates incorrect code based on a mistaken assumption of what is returned. In contrast, when given syntax-aware chunks, the model accurately identifies the return values and integrates them correctly within the existing codebase.
  • Figure 2: Comparison of fixed-size chunking vs. cAST. For cAST, we first parse the document into a tree of AST nodes. Then, starting from the first level, we greedily merge AST nodes into chunks. If adding a node would exceed the chunk size limit, we recursively break it into smaller nodes. The output of cAST is a list of chunks where each chunk contains a list of AST nodes.