CodeMEM: AST-Guided Adaptive Memory for Repository-Level Iterative Code Generation

TL;DR

CodeMEM tackles repository-level iterative code generation by introducing two AST-guided memory components: Code Context Memory for dynamic, code-structured context updates and Code Session Memory for memory of iterative edits. The system uses AST-based selectors and detectors to preserve relevant code and detect forgetting, yielding state-of-the-art instruction-following and session-level performance on CodeIF-Bench and CoderEval while maintaining competitive latency and token efficiency. Key contributions include the AST-guided memory selector, the memory-based forgetting detector, and extensive ablations validating the importance of each component. The approach advances practical, scalable tool support for long-horizon, repository-aware code generation in real-world software development workflows.

Abstract

Large language models (LLMs) substantially enhance developer productivity in repository-level code generation through interactive collaboration. However, as interactions progress, repository context must be continuously preserved and updated to integrate newly validated information. Meanwhile, the expanding session history increases cognitive burden, often leading to forgetting and the reintroduction of previously resolved errors. Existing memory management approaches show promise but remain limited by natural language-centric representations. To overcome these limitations, we propose CodeMEM, an AST-guided dynamic memory management system tailored for repository-level iterative code generation. Specifically, CodeMEM introduces the Code Context Memory component that dynamically maintains and updates repository context through AST-guided LLM operations, along with the Code Session Memory that constructs a code-centric representation of interaction history and explicitly detects and mitigates forgetting through AST-based analysis. Experimental results on the instruction-following benchmark CodeIF-Bench and the code generation benchmark CoderEval demonstrate that CodeMEM achieves state-of-the-art performance, improving instruction following by 12.2% for the current turn and 11.5% for the session level, and reducing interaction rounds by 2-3, while maintaining competitive inference latency and token efficiency.

Figures (4)

• Figure 1: Examples of failure cases in iterative code generation. Left: As interactions progress, the LLM exhibits forgetting and overwrites previously correct modifications. Right: Static context handling propagates erroneous information and prevents the incorporation of relevant contexts, resulting in repeated incorrect code.
• Figure 2: The overall pipline of CodeMEM.
• Figure 3: Cumulative recall and precision per round for code contexts on CodeIF-Bench.
• Figure 4: Forgetting rate (IFR) result on CodeIF-Bench.