c AST : Enhancing Code Retrieval-Augmented Generation with Structural Chunking via Abstract Syntax Tree (2025)

Abstract: Retrieval-Augmented Generation (RAG) has become essential for large-scale code generation, grounding predictions in external code corpora to improve factuality. 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 (cAST), 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.

## introduction large-scale code generation has emerged as a cornerstone of modern software engineering, powering tasks that range from automated bug fixing @xcite to full-fledged repository-level completion @xcite . retrieval-augmented generation (rag) pushes this frontier further by allowing language models to ground their predictions in a rich external corpus of data @xcite , effectively mitigating hallucinations and improving factual correctness @xcite . one crucial preprocessing step in retrieval-augmented generation (rag) is chunking @xcite -breaking large documents into manageable segments that can be efficiently indexed, ✅ 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 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. can we chunk documents more intelligently, preserving their original structure? in this work, we explore cast-chunking via abstract syntax trees. asts represent code as hierarchical trees with typed nodes corresponding to program units. by parsing source code into an ast, we apply a recursive, split-then-merge algorithm to convert tree structures into chunks that are better aligned with syntactic boundaries. extensive experiments show that cast improves performance across a range of code generation tasks. specifically, it offers three key advantages: (1) structure-preserving chunks: ast traversal yields more self-contained chunks, improving both retrieval and generation. for instance, starcoder2-7b sees an average of 5.5 points gain on repoeval @xcite . (2) crosslanguage consistency: the language-agnostic nature of cast enables better generalization across programming languages, achieving up to 4.3 points gain on crosscodeeval @xcite . (3) metadata retention: ast-based chunks more faithfully capture metadata at the file, class, and function levels, enhancing context matching in hybrid code+natural language tasks, e.g., up to 2.7 points gain on swe-bench @xcite , which focuses on resolving github issues. ## cast we focus on the first stage of the rag pipeline: chunking. in this step, source code is parsed into semantically meaningful units (such as functions or classes) while preserving the structure of the code. these units are then grouped into coherent chunks, which serve as the retrievable context that can be obtained by a subsequent retriever and used to prompt a language model. design goal. our design for cast pursues four aligned goals: (1) syntactic integrity-whenever possible, chunk boundaries should align with complete syntactic units instead of splitting them; (2) high information density-each chunk is packed up to, but not beyond, a fixed size budget to maximize content utility; (3) language invariance-the algorithm employs no language-specific heuristics so it works unchanged across diverse programming languages and code-related tasks; and (4) plugand-play compatibility-concatenating the chunks must reproduce the original file verbatim, enabling seamless drop-in replacement within existing rag pipelines. ## experiments we evaluate cast with various top retrieval and generation models in various code task settings. we present results of selected end-to-end racg pipelines (retriever + lm) in section 3.2 and full tables in the appendix (5, 6, 7, 8). ## ablations necessity of merging. the motivation for introducing merging in our algorithm is to maximize the information density of each chunk. under a ## conclusion and discussion in this work, we present cast as a simple and effective chunking strategy for retrieval-augmented code generation. through the structural awareness brought by ast, we are allowed to maintain syntactic integrity and high information density during chunking. extensive experiments on various retrievers, llm generators, and code generation tasks, validate the gain from cast over the commonly used fixed-size chunking strategy on both retrieval and rag tasks. by maintaining the original rag pipeline, for the code agent practitioner, cast could be used as a simple plug-and-play tool to provide informative and formatted chunks for later stage agent use. for code rag benchmark developers, cast could serve as additional resources and an effective alternative or complementary retrieval unit.