Source-linked AI summary
Type-Constrained Code Generation with Language Models
Niels Mündler, Jingxuan He, Hao Wang, Koushik Sen, Dawn Song, Martin Vechev
TL;DR
LLMs often generate uncompilable code because next-token generation does not formally model typing, while syntax-only constrained decoding cannot address this gap. The paper introduces type-constrained decoding with prefix automata and inhabitable-type search, and reports more than halved compilation errors with consistently higher functional correctness across coding tasks.
Problem
LLMs frequently generate typing errors, while syntax-only constraints cannot capture the formal rules needed to ensure well-typed code.
Method
The approach combines prefix automata with sound type search to determine whether partial programs can be completed into well-typed programs.
Results
The approach more than halves compilation errors and consistently increases functional correctness across synthesis, translation, and repair tasks and diverse models.
Takeaways & Limitations
Type-system constraints can guide LLM code generation beyond syntax across a simply typed language and a significant TypeScript subset.
Takeaways & Limitations
The evaluated TypeScript subset can still encounter compiler inference limitations, including empty-array inference to never[].
Abstract
from arXiv · showhide
Large language models (LLMs) have achieved notable success in code generation. However, they still frequently produce uncompilable output because their next-token inference procedure does not model formal aspects of code. Although constrained decoding is a promising approach to alleviate this issue, it has only been applied to handle either domain-specific languages or syntactic features of general-purpose programming languages. However, LLMs frequently generate code with typing errors, which are beyond the domain of syntax and generally hard to adequately constrain. To address this challenge, we introduce a type-constrained decoding approach that leverages type systems to guide code generation. For this purpose, we develop novel prefix automata and a search over inhabitable types, forming a sound approach to enforce well-typedness on LLM-generated code. We formalize our approach on a foundational simply-typed language and extend it to TypeScript to demonstrate practicality. Our evaluation on the HumanEval and MBPP datasets shows that our approach reduces compilation errors by more than half and significantly increases functional correctness in code synthesis, translation, and repair tasks across LLMs of various sizes and model families, including state-of-the-art open-weight models with more than 30B parameters. The results demonstrate the generality and effectiveness of our approach in constraining LLM code generation with formal rules of type systems.
1 Introduction
LLMs can synthesize, translate, and repair code, but probabilistic next-token generation provides no formal guarantees and often produces compilation errors. This work introduces type-constrained decoding, extending constrained generation beyond syntax to enforce well-typed code and evaluating it across languages, models, and tasks.
- LLMs support code synthesis, translation, and repair, yet generated code often contains compilation errors, logic flaws, or security vulnerabilities.
- 94% of compilation errors arise from failed type checks on average, while syntactic errors account for 6% in generated TypeScript code.
- Type-constrained decoding uses a prefix automaton and sound type search to determine whether partial programs can become well-typed.
- The approach is instantiated for TypeScript and evaluated on HumanEval and MBPP across synthesis, translation, and repair tasks.
- The approach is designed to apply broadly to languages derived from the core calculus, next-token LLMs, and additional production languages or closed-weight models.
2 Background and Overview
Constrained decoding filters sampled tokens using a completion engine, but syntax-only constraints cannot reject many type-invalid programs. The proposed approach combines prefix automata with type-reachability search to guide valid completions in TypeScript and related tasks.
- Background on constrained decoding: Constrained decoding repeatedly samples tokens and rejects those whose prefixes cannot be completed into well-formed programs.
- Background on constrained decoding: Token-level prefix checks provide inductive guarantees that the final returned program is valid in the constrained language.
- Limitations of syntax-only constraining: Syntax-only constraints accept some syntactically valid completions that still cause compilation errors, including undeclared identifiers, invalid calls, and wrong argument types.
- Type-constrained decoding: The type-constrained approach uses typing information in a prefix automaton to reject invalid completions and guide generation toward the valid candidate.
- Type-constrained decoding: Type reachability searches operator sequences over a type graph to determine whether partial expressions can attain required types.
- Evaluation: The implementation covers a significant TypeScript subset and evaluates synthesis, translation, and repair, reducing compilation errors and increasing functional correctness.
3 Our Type Constraining Approach
The paper defines a simply typed language and constructs prefix automata that parse fragments whose partial expressions can be completed into well-typed programs. A type-reachability search and soundness results support constrained decoding while termination heuristics make the automata incomplete.
- 3.1 A Simply Typed Language: The simply typed language 𝐿𝐵 combines a grammar with typing rules for expressions, statements, environments, and function returns.Its design resembles statically typed languages such as TypeScript, Java, and Go.
- 3.5 Prefix Automaton for Statements: Statement parsing tracks compatible type environments and enforces declared return types on every execution path.Function automata reject mismatched returns and require additional statements when not all paths return.
- 3.2 Prefix Automaton Definition: The final prefix automaton is intentionally incomplete because termination requires excluding high-order or highly complex types.The paper reports that this limitation primarily avoids types considered less likely in practical use.
- 3.4 Prefix Automaton for Expressions: Prefix automata parse well-typed expressions by tracking syntax and type-relevant context, including constrained target types.The constrained automaton accepts a partial expression only when its completion can inhabit the required type.
- 3.4 Prefix Automaton for Expressions: The two-tiered algorithm first derives types for a partial expression, then searches reachable types obtainable through further operator extensions.This addresses the possibility that repeated extensions change the expression’s type and may continue indefinitely.
- 3.4 Prefix Automaton for Expressions: The type-search algorithm is sound, and the resulting expression automata parse subsets of well-typed expressions in 𝐿𝐵.The automata preserve the prefix property while rejecting transitions that cannot lead to the required type.
4 Extension to TypeScript
The approach is extended to a core subset of TypeScript by supporting language features through type-environment tracking and adapted type search. The implementation adds practical constraints, including required annotations, that trade theoretical completeness for correctness.
- 4 Extension to TypeScript: The TypeScript extension supports a core subset of the language and documents supported and unsupported features separately.The completion engine from the foundational language is extended rather than replaced.
- Constant Variable Declarations: The implementation tracks immutable const identifiers and rejects assignments to them.Mutability is stored in the type environment alongside identifier types.
- Arrays: Array support enforces uniform element types and adapts type-reachability pruning to nested array types.The additional nesting dimension requires pruning changes to ensure termination.
- Loops: Loop support handles several loop constructs, using a generic array type •[] for for..of right-hand sides.The generic type matches arrays such as number[] and string[].
- Polymorphic Built-In Members: Polymorphic built-in members are handled by tracking type patterns and instantiating type parameters before continuing the type search.For x.map(f), the instantiated callback result determines the returned array type.
- Type Annotations: The implementation requires annotations for function parameters and returns, while variable declarations need either annotations or initializers.These restrictions provide more type information but trade practical correctness against theoretical language completeness.
5 Experimental Evaluation
The evaluation studies type-constrained decoding across TypeScript synthesis, translation, and repair tasks using multiple open-weight LLMs. It finds substantial reductions in compilation errors and improvements in functional correctness, with measurable runtime overhead and illustrative corrections of typing failures.
- Experimental Setup: The evaluation covers synthesis, translation, and repair on TypeScript-translated HumanEval and MBPP tasks across six open-weight LLMs.The models include Gemma 2 at 2B, 9B, and 27B parameters, DeepSeek Coder 33B, CodeLlama 34B, and Qwen2.5 32B.
- Compilation Errors: Type-constrained decoding reduces compilation errors by 75.3% on HumanEval and 52.1% on MBPP, with minimum reductions of 54.8% and 27.3% across models.Syntax errors account for only 9.0% and 4.9% of non-compiling synthesis and translation instances in the unconstrained settings.
- Compilation Errors: Type-constrained decoding substantially improves compilation-error repair, resolving 53.7% more errors on average than vanilla decoding.For Gemma 2 2B, repair rates increase from 33.5% to 56.4% on HumanEval and from 25.8% to 58.4% on MBPP.
- Functional Correctness: Type constraining increases average pass@1 by 3.5% for synthesis, 5.0% for translation, and 37.0% for repair.The largest gain occurs in repair because vanilla models generally struggle to generate functionally correct code.
- Case Studies: Type-constrained sampling corrects distinct TypeScript failures, including missing arguments, missing return paths, and inadequate callback type annotations.The case studies show the completion engine steering generation toward a required argument, fallback return statement, or explicit callback annotation.
6 Discussion
The discussion identifies practical boundaries for type-constrained decoding: implementation requires language-specific completion-engine work, stronger models may not solve harder settings, and constrained generation can still fail to terminate or require token distributions unavailable through black-box APIs.
- Implementation Effort: Adapting the method to another language currently requires manual development of a completion engine.The authors expect features from their simply typed language and TypeScript implementations to transfer, reducing future effort.
- Implementation Effort: Compiler integration as an incremental completion engine could automate adoption of constrained code generation alongside grammar parsing and type checking.
- Broader Application to More Complex Tasks and Stronger LLMs: Stronger models may make fewer typing errors, but difficult tasks, stricter type systems, and low-resource languages remain challenging.Reported compilation-error rates reach 40%–60% for OCaml and Haskell across models, while Rust tasks also retain substantial error rates.
- Access to Model Probabilities: Constrained decoding requires next-token probability distributions, which commercially available black-box APIs generally do not expose.The authors suggest backend integration by model providers as one possible solution.
- Remaining Compiler Errors: Even with valid results guaranteed upon termination, generation loops can leave compilation errors when constrained decoding prevents recovery within token or time limits.The discussion attributes these failures to nontermination after constraints amend generation.
7 Related Work
Related work covers code-generation models, accuracy-improvement techniques, syntax-focused constrained decoding, and earlier uses of type systems that do not generalize to complex program generation.
- Code Language Models: Code language models support synthesis, repair, and translation but are also known to make frequent mistakes.These models are trained on datasets ranging from billions to trillions of tokens and commonly contain billions of parameters.
- Improving Language Model Accuracy: Fine-tuning, retrieval augmentation, and compiler or execution feedback are three established approaches for improving code-generation accuracy.Fine-tuning is described as highly resource intensive, while retrieval augmentation supplies additional contextual information.
- Constrained Decoding: Prior constrained-decoding methods primarily enforce syntactic features, although some address simple context-sensitive properties.The paper notes that syntax errors account for only a small fraction of compilation errors in its evaluation.
- Type Systems for Code Synthesis: Earlier type-system methods targeted specialized settings and could not constrain general, complex program generation.Examples include SQL generation, function-call completion, and object-member access guidance.
8 Conclusion
The paper uses type systems to guide language-model decoding through prefix automata, evaluates the approach across code-generation tasks, and reports broad improvements in compilation and functional correctness.
- Conclusion: The approach designs prefix automata for a foundational simply typed language and extends them to TypeScript.
- Conclusion: The evaluation covers code synthesis, translation, and repair across a diverse set of models.
- Conclusion: The method more than halves compilation errors and consistently increases functional correctness.
- Conclusion: The paper develops detailed definitions and analyses for the automata used in its simply typed foundation.
A.1 Base Automata
The base automata are assembled from union, concatenation, Kleene-star, terminal, and whitespace-aware constructions, with proofs establishing their accepted languages and prefix properties.
- Union: Union combines initial and accepting states and merges transition functions to accept either component language.The construction preserves the prefix-automaton property when both component automata have it.
- Correctness Properties: The automata constructions rely on reachable-state decomposition to establish their accepted-language and prefix-automaton properties.
- Concatenation: Concatenation preserves the first automaton’s parsing behavior while adding transitions from its accepting states into the second automaton.Its accepted language is the concatenation of the component languages, and the construction remains a prefix automaton when the second language is nonempty.
- Kleene-Star: Kleene-star makes initial states accepting and links final states back to initial states to parse indefinite repetitions.The empty word is accepted because zero repetitions are allowed.
- Terminals: The terminal automaton parses exactly one specified terminal using states indexed by terminal suffixes.Its accepting state represents the empty suffix after the terminal has been consumed.
- Terminals: A whitespace-aware terminal construction permits arbitrary whitespace before parsing a terminal.
A.2 Expressions
The expression automata propagate type environments and expression metadata while restricting operators, calls, and member access to type-valid continuations.
- Automaton structure: Recursive expression automatons parse expressions and pass state information through transitions and composite automata.
- State information: Expression states carry the current type environment, left-hand-side expression, and type of the last coherent accepting expression.The typ attribute is defined only for accepting states, with lhs helping determine expression types.
- Operator and access constraints: Arithmetic operators are constrained to valid operators based on the left-hand side expression.
- Operator and access constraints: Function-call automata accept calls only when the left-hand side is a function with a valid signature.
- Operator and access constraints: Member-access automata parse attributes of the left-hand-side expression by combining the corresponding attribute automata.
A.3 Pruning the Type Search
The pruning heuristic bounds type-reachability search using type depth and root types, while preserving exploration of higher-order types that can expose novel reachable types.
- Heuristic rationale: The heuristic prunes candidate types according to their complexity and novelty during recursive type-reachability search.
- Heuristic rationale: From a type, extensions reach the original type, arithmetic results, function return types, and member types; higher-order types are usually avoided unless they add relevant reachability.
- Depth and root types: Type depth measures function order, while root returns the minimal-depth types composing a higher-order type.
- Pruning rules: The search stops exploring a type when its depth exceeds the maximum depth of the goal and current types.
- Pruning rules: Higher-depth functions remain eligible when they return an unexplored type, whereas types whose roots are already explored are pruned.
- Scope: The heuristic robustly seeks relevant inhabitable types, but completeness is not guaranteed because the lookup function adds complexity.
A.4 Implementation of derivable
The implementation makes derivable-type computation finite and more efficient by restricting exploration to the target-bounded reachability graph and integrating nested searches.
- Derivable-type search: Computing derivable types directly can be intractable because function types may have arbitrarily high order and therefore infinitely many candidates.
- Derivable-type search: The implementation bounds candidate types using the finite reachability graph produced by pruning heuristics for the target type.
- Search integration: Nested reachability calls are integrated so discovered intermediate types transition directly to function types, prioritizing paths to the goal.
- Return tracking: Function bodies propagate expected return types and track whether blocks must return or have returned on all branches.
- Statement automata: The statement automaton recursively combines declaration, expression, return, block, function, and conditional automata.
- Return tracking: A block requiring a return cannot accept until another statement is generated, while literals and return statements preserve the prefix property.
B Details about Experimental Evaluation
The evaluation implementation uses heuristic output extraction, standard software dependencies, fixed experimental settings, and exclusions for six MBPP instances with overly broad annotations.
- Implementation: The decoding implementation extracts executable code from model responses and applies a throughput-oriented heuristic modification.
- Implementation: Regular-expression literal automata use the regex library, while LLM inference uses transformers; supported and unsupported TypeScript features are catalogued.
- Experimental settings: Experiments use A100 GPUs, CUDA 12.4, temperature 1, fixed seeds, 1000-token completion limits, and 300-second timeouts.
- Experimental settings: Syntactic correctness is measured with the Oxidation toolchain because the official TypeScript compiler does not clearly separate syntactic and semantic errors.
- Dataset handling: 6 MBPP instances are excluded because their TypeScript translations contain overly broad any or array of any annotations.
- Prompting and extraction: Prompts use pre-filled function signatures for unified unit testing, and repair prompts include the non-compilable model output.
- Prompting and extraction: Code extraction selects TypeScript code blocks and truncates after the last balanced closing brace.
C Case Study Full Outputs
The case studies compare unconstrained and type-constrained TypeScript generations across synthesis, translation, and repair examples. Type constraints address distinct typing failures, including invalid method calls, missing return paths, and incorrect inference in array reductions.
- Translation: In translation, unconstrained Gemma 2 2B calls string.split without parameters, which is legal in Python but invalid in TypeScript.
- Synthesis: In synthesis, unconstrained DeepSeek Coder 33B Instruct can return undefined when a loop does not execute, violating the declared number return type.The example remains functionally correct for positive inputs but fails the type-system requirement for negative inputs.
- Synthesis: In synthesis, unconstrained Gemma 2 9B infers a reduce accumulator as never[], so pushing a number causes a type violation.Type-constrained decoding forces callback annotations, enabling the correct accumulator type and avoiding this issue.