Typestate-Guided Context Retrieval

• Typestate-guided context retrieval is a formal method using typestate automata to track resource states and extract relevant, structured program context for downstream systems like LLMs.
• This approach traces error propagation paths by monitoring state transitions of memory objects across function calls and aliases, capturing essential spatial and temporal information.
• By providing concise, semantically rich traces instead of raw code, it enables LLMs to effectively repair complex, interprocedural errors like memory bugs within strict token constraints.

Typestate-guided context retrieval is a formal approach for extracting, structuring, and supplying relevant program context to downstream reasoning or repair systems, such as LLMs, based on an explicit tracking of resource states through typestate automata. In the setting of memory error repair for C codebases, this method leverages finite typestate automata to track both the sequence of stateful memory operations (alloc, free, use) and their propagation across interprocedural paths and aliases. This results in concise, semantically rich trace information that enables accurate detection and automated repair of severe memory errors, despite token-window limitations in LLMs and the distributed nature of C’s manual memory management.

1. Typestate Automata for Memory State Modeling

A typestate automaton is defined as a tuple: $$\mathcal{A}_{\textup{ET}} = \left<\Sigma ,\mathbb{T},T_u,\delta,T_\textup{ET} \right>$$ where:

• $\Sigma$ is the set of all relevant memory operations (e.g., alloc, free, use).
• $\mathbb{T}$ is a set of abstract typestates such as uninitialized ($T_u$), live ($T_l$), dead ($T_d$), and error ($T_\textup{ET}$).
• $T_u$ denotes the initial typestate.
• $\delta$ is the transition function, mapping $(\text{state}, \text{operation})$ pairs to new typestates.
• $T_\textup{ET}$ represents specific error states (e.g., UAF, double-free, memory leak).

For example, the UAF automaton allows: alloc $\rightarrow$ live (use), free (on live) $\rightarrow$ dead, use (on dead) $\rightarrow$ UAF error.

Each typestate automaton is tailored per error type, and these automata serve as executable specifications for detecting not only the presence but also the origin and propagation of memory mismanagement.

2. Typestate-Guided Context Extraction and Error-Propagation Paths

The central methodology is to trace the error-propagation path for each relevant memory object address, as it transitions through various typestates from allocation through the erroneous operation. The core extraction process is as follows:

• Instrument execution (e.g., via GDB) using a proof-of-concept input that reliably triggers the error.
• Monitor every step involving the memory address of interest. Each operation in $\Sigma$ is observed.
• For each such operation, evaluate $\delta$ to track typestate transitions for that address.
• Build an error-propagation path $\pi_e = (s_i)_{i=m}^n$ as an ordered series of statements from allocation to the error trigger, only including steps where a typestate transition occurs.

Snapshots are captured at each typestate transition point, recording the operation, code location (file:line), and current call stack. The set of these context-rich events across the error-propagation path is assembled into a map ($\tMap$), which summarizes the information needed for subsequent reasoning or repair.

3. Spatial and Temporal Context Capture

This methodology explicitly encodes both:

• Spatial dimension: The active concrete memory state (live, dead, error) for each tracked allocation, including all aliasing relationships (multiple pointers to the same address). Each memory state transition considers the alias set; thus, interventions or errors through any pointer variable are tracked for the underlying memory object.
• Temporal dimension: The execution history, reflected as the precise ordering of operations/events (alloc, use, free, error) from allocation through error manifestation. This path may span multiple functions and files, providing a cross-cutting trace across the global codebase.

A typestate-guided error-propagation path thereby encodes the minimal, semantically essential trace required to understand and repair an error.

4. LLM Integration under Token Constraints

In this system, LLMs are prompted not with raw, window-greedy code slices, but with crafted, context-rich prompts containing only the typestate-guided context. Specifically:

• Inputs include error type, error location, a sequential trace of state transitions (as code snippets), and associated call stacks.
• The context provided is typically 41× shorter than comparable LLM APR frameworks (such as SWE-Agent), avoiding the “lost in the middle” issue for critical error signals.
• Role-based and stepwise prompts further guide the LLM’s reasoning.

This approach allows LLMs to generate patches for complex, interprocedural errors, spanning multiple functions and call contexts, while minimizing unnecessary context and keeping within strict token budgets.

5. Addressing Repair Challenges: Interprocedural Reasoning and Window Limits

Traditional program repair systems and LLM-based agents are challenged by the need for repository-scale reasoning and the constraints of fixed LLM context windows. Typestate-guided context retrieval addresses:

• Interprocedural and aliasing complexity: By following object lifecycles and propagating typestate transitions across function boundaries and pointer aliases, the extracted traces naturally bridge gaps left by traditional local or syntactic approaches.
• Token budget and efficiency: Only transitions—moments of semantic state change—are included, not every piece of code or every statement in the execution path, leading to highly compressed but informative traces for repair.

6. Practical Outcomes and Impact

This methodology demonstrates significantly improved efficacy in automated memory error repair:

• In empirical evaluation, typestate-guided context retrieval enabled automated repair of 31 of 40 real-world memory errors in large C repositories—a clear improvement over both traditional APRs (e.g., SAVER, ProveNFix) and LLM-based greedy-context baselines (e.g., SWE-Agent).
• Successfully produced patches accepted as fixes for zero-day vulnerabilities.
• Resulted in lower rates of new error introduction, due to targeted, semantics-rich prompting.

Applications extend to real-world repair of large systems, complex compiler/interpreter bugs, and scalable analysis across entire repositories. The paradigm is extensible to other classes of temporal or protocol errors beyond memory management.

7. Summary Table: Paper’s Core Concepts

Aspect	Approach/Contribution	Key Data/Paper Section
Typestate Automaton	Formal FTA for resource lifecycles/errors; guides context extraction	Def 1, Table 1, Fig. 1
Error-Propagation Path	Dynamic tracing, context only at state transitions	Algorithm 1, Defs 4–7
Spatial/Temporal Context	Alias-aware, lifecycle-tracking, cross-function execution order	Motivating example, Fig. 2, 3
LLM Integration	Thin, transition-centric prompts, structured for patch synthesis	§Approach, §Evaluation, Figs. 2, 4
Challenges & Solutions	Scales to large codebases, solves context loss, no window issues	§Introduction, §Ablation
Impact	Robustness, high repair rates, accepted upstream, low error induction	§Evaluation, §Case Study, Fig. 4

Conclusion

Typestate-guided context retrieval for memory error repair leverages finite typestate automata to extract and condense the essential interprocedural traces and state transitions leading to errors. This precision enables LLMs to repair complex memory mistakes in C systems, overcoming both the semantic and technical obstacles previously faced by automated and LLM-based repair systems. The method’s ability to combine formal program analysis with tight, meaningful context provides a blueprint for scalable, semantically aware automated code repair for memory and, by extension, other temporal or protocol-driven bugs.