Type-Constrained Code Generation with Language Models

1 Introduction

• Syntax-only constrained decoding improves syntax [8] [41] [57] [66], yet syntax errors are 6% of TypeScript compilation errors, while failed type checks account for 94%.Because next-token sampling provides no formal guarantees, generated code can also contain logic flaws or security vulnerabilities.Source: §1 · pp. 1–2
• Type systems cannot generally be captured by context-free grammars, and checking partial expressions requires solving type inhabitation to determine whether they can be completed with a required type.This challenge prevents directly applying prior syntax-focused constrained-decoding methods.Source: §1 · p. 2
• A sound prefix automaton tracks abstract syntax trees and type-relevant context, while type search tests whether partial expressions can inhabit required types, retaining only completable well-typed prefixes.The approach is constructed for a generic, simply typed Turing-complete calculus.Source: §1 · p. 2
• Type-constrained decoding reduced compilation errors by more than half, relatively increased functional correctness by 3.5% to 5.5% in synthesis and translation, and improved correct repair by 37% on average.The evaluation used TypeScript versions of HumanEval and MBPP across LLMs with 2B–34B parameters.Source: §1 · p. 2
• The approach applies to languages derivable from the core calculus, code-generation tasks in those languages, and any LLM using next-token generation.The paper instantiates it on a non-trivial TypeScript subset.Source: §1 · p. 2

2.1 Background on LLM-based Code Generation and Constrained Decoding

• Constrained decoding accepts sampled tokens only when the partial program can still be completed in L, and ends only when the program is in L, guaranteeing validity.The sample-and-check loop is bounded by the vocabulary size and requires no additional LLM inference.Source: §2.1 · pp. 3–4
• Grammar and type checkers that judge only complete programs cannot serve as completion engines because they do not identify prefixes that can be completed.Source: §2.1 · p. 4

2.2 Overview of Our Type Constraining Approach

• Syntax accounts for only 3.5% of functional errors and 6% of compilation errors in evaluated LLM-generated code.Source: §2.2 · p. 4
• The approach more than halves compiler errors and increases functional correctness across synthesis, translation, and repair for a significant TypeScript subset.The paper evaluates the approach across various LLMs.Source: §2.2 · pp. 5–6
• Type-constrained decoding is challenging because type systems exceed context-free grammars and checking completable well-typed expressions also requires type inhabitation.Source: §2.2 · p. 5
• The proposed prefix automata parse syntax while tracking typing context, including initialized variables and enclosing functions, to admit only type-valid transitions.The approach enforces well-typedness in a simply typed language and extends to a core subset of TypeScript.Source: §2.2 · p. 5
• Type reachability searches a graph of types and well-typed operations to find completions such as num.toString() that turn a number expression into a string.The search is invoked for each partially parsed expression to determine valid automaton transitions.Source: §2.2 · p. 6

3.1 A Simply Typed Language

• L_B is a Turing-complete, simply typed language with TypeScript-like syntax, typing rules, statement-updated environments, and checks that function returns match their declared type on every path.Its expression rules draw on safeFTS, a type-safe subset of TypeScript, and its grammar can be adapted to other statically typed languages.Source: §3.1 · pp. 6–7

3.2 Prefix Automaton Definition

• A prefix automaton’s reachable strings are exactly prefixes of accepted strings, enabling a completion engine; parsing a subset of well-typed programs enforces validity but may omit valid programs.The final automaton covers a significant but incomplete subset because high-order types are excluded to ensure termination.Source: §3.2 · pp. 9–10

3.3 Prefix Automata for Identifiers, Literals, and Types

• Automata for identifiers, literals, and type annotations exactly match their language counterparts, with identifiers drawn from the current environment and function types defined recursively.The recursive type automaton preserves the prefix property.Source: §3.3 · pp. 10–11

3.4 Prefix Automaton for Expressions

• Type-constrained expression automata use a two-tier search over derivable and reachable types to prune transitions, accepting only well-typed expressions while potentially excluding some valid ones.The search is sound; to terminate, it prunes paths through high-complexity types, so it is incomplete.Source: §3.4 · pp. 11–13
• Expression automata accept only well-typed expressions in L_B, and reachability of derivable(q) for q in γ(I_e,s) ensures a suffix can complete s with type G.The type restriction applies recursively through A_e↓T and A_e.Source: §3.4 · p. 14

3.5 Prefix Automata for Statements

• The statement automaton is a prefix automaton whose accepted language is a subset of L_B, except for well-typed expressions rejected by A_e.It combines statement forms recursively, carries compatible identifier environments across sequences, discards local environments after blocks, conditionals, and functions, and adds function signatures and parameters within function bodies. Function parsing enforces matching return types and returns on every execution path, forcing further statements when needed.Source: §3.5 · p. 14

• The TypeScript extension implements a core subset of the language; a comprehensive list of supported and unsupported features is provided in Appendix B.Source: §4 · p. 14

Constant Variable Declarations

• For const declarations, the type environment tracks identifier mutability and rejects assignments to immutable identifiers.Source: §4 · p. 14

Arrays

• Array support covers annotations, expressions, and field access or assignment; elements must share a type, and reachability pruning handles added type nesting to ensure termination.Source: §4 · p. 15

Loops

• Loop constructs are implemented as statement-block parser variations, while for..of requires an array and uses generic type •[] to match arrays of any type.Source: §4 · p. 15

Additional Operators and Types; Operator Precedence

• Additional arithmetic and logical operators add type-search edges, while precedence-aware parsing and search determine expression grouping and omit invalid operator extensions.Source: §4 · p. 15

Global Identifiers and Imports

• Global identifiers are added through program type-environment initialization, with additional types for identifiers such as Math and modeled imports for the crypto library using require.Source: §4 · p. 15

Polymorphic Built-In Members

• For polymorphic built-in members such as array map, type search instantiates type parameters from argument patterns before continuing, and constrains generated anonymous callbacks to match those patterns.Source: §4 · p. 15

Type Annotations

• To aid generation and constrain compilation failures, the implementation requires function parameter and return annotations, initialized or annotated variables, and a type annotation on reduce callbacks’ first parameter.These requirements trade off practical correctness with theoretical language completeness.Source: §4 · pp. 15–16

5.1 Experimental Setup

• The evaluation covers synthesis, translation, and repair on TypeScript-translated HumanEval and MBPP tasks, using six instruction-tuned open-weight models and comparing Vanilla, idealized Syntax, and Types.HumanEval has 159 instances and is sampled four times per instance; MBPP has 384 instances and is sampled once. Repair uses 292 HumanEval and 248 MBPP non-compiling instances.Source: §5.1 · p. 16

5.2 Results on Compilation and Functional Correctness

• Type constraining reduced compilation errors by 75.3% on HumanEval and 52.1% on MBPP across synthesis and translation; in synthesis, Table 2 reports reductions of 74.8% and 56.0%, versus 9.0% and 4.8% for idealized syntax-only constraints.Models across sizes and families benefited, with minimum error reductions of 54.8% on HumanEval and 27.3% on MBPP.Source: Table 2 · p. 17, §5.2 · p. 17
• Type constraints increased average pass@1 by 3.5% in synthesis, 5.0% in translation, and 37.0% in repair, improving functional correctness across all three tasks.For synthesis, 17 of 26 resampled Gemma 2 27B outputs and 15 of 21 Qwen2.5 32B outputs were also functionally correct.Source: Table 3 · p. 18, §5.2 · p. 18
• Type-constrained repair resolved 53.7% more compilation errors on average than vanilla decoding, raising Gemma 2 2B repair rates from 33.5% to 56.4% on HumanEval and from 25.8% to 58.4% on MBPP.Repair inputs were non-compiling programs from unconstrained synthesis.Source: §5.1 · p. 16, §5.2 · p. 18

5.3 Runtime Analysis

• Type-constrained decoding increased median runtime by 39.1% on HumanEval and 52.1% on MBPP, while 99.4% of sample-and-check cases required only one iteration.The overhead was measured with an unoptimized Python implementation; the iteration result reflects the evaluated cases, with higher one-iteration rates for Gemma 2 9B and 27B.Source: §5.3 · p. 19, Figure 8 · p. 19

5.4 Case Study

• Case studies show type constraints correcting missing arguments, missing guaranteed returns, and incorrect type inference in generated TypeScript.For the reduce example, an annotation prevents TypeScript from inferring never[] for an empty-array accumulator.Source: §5.4 · pp. 19–20, Figure 9 · p. 20

Implementation Effort

• Adopting the method requires manual language-specific completion-engine work, though features from the L_B and TypeScript implementations may reduce that effort.The authors suggest language developers build compilers as incremental completion engines, enabling automatic adoption alongside conventional parsing and type checking.Source: §6 · pp. 20–21

Broader Application to More Complex Tasks and Stronger LLMs

• The approach may extend to complex tasks, stricter typing, and low-resource languages, where current LLMs still produce substantial compilation errors.Reported rates include 18% for OpenAI o1-mini, 39% for DeepSeek R1 [15], and 27% for Claude 3.5 Sonnet on difficult Rust synthesis; OCaml and Haskell rates reach 40%–60% [24]. GPT-4 translation errors were 44.3% [54], while C-to-Rust errors were around 25% for OpenAI o1 models [61].Source: §6 · p. 21

API Access

• Constrained decoding requires next-token probability distributions that commercial black-box APIs do not expose, so backend integration is a possible solution.The authors point to backend support for JSON Schema adherence as an example.Source: §6 · p. 21

Remaining Compiler Errors

• Compilation errors remain when constrained generation enters loops and fails to terminate within token or time limits, despite guaranteeing valid results upon termination.The authors suggest adding constraints to stop such loops, for example by limiting expression complexity.Source: §6 · p. 21, Figure 10 · p. 21

7 Related Work

• LLMs support code synthesis, repair, and translation [35], but even state-of-the-art open-weight models with over 30 billion parameters frequently make code-generation errors.Such models are typically trained on billions to trillions of tokens and have billions of parameters [28] [46] [59] [64].Source: §7 · p. 21
• Accuracy methods include resource-intensive fine-tuning [65], retrieval augmentation with related context [57], and compiler or execution feedback requiring resampling after generation [34].Constrained decoding is orthogonal to these approaches; combining it with retrieval augmentation or compiler feedback additionally improves performance [57].Source: §7 · p. 22
• Syntax errors account for only 6% of compilation errors on average, while type-constrained decoding more than halves compilation errors.Source: §7 · p. 22
• Prior constrained decoding has covered context-free languages [7] [8] [57] and features such as Python indentation and Go scope markers [41] [66], but remains limited to syntax.Source: §7 · p. 22
• Prior type-guided synthesis targeted SQL column names [57], function-call completion, member accesses, or context rather than hard constraints [11]; specialized search methods were incompatible with LLM generation.These approaches were confined to specialized settings and could not constrain general, complex program generation.Source: §7 · p. 22

8 Conclusion

• Compilation errors were more than halved, and functional correctness consistently increased across a diverse set of models.These outcomes were observed in evaluations of code synthesis, translation, and repair.Source: §8 · p. 22
• The authors designed and implemented prefix automata to constrain generation for a foundational simply typed language and extended the approach to TypeScript.Source: §8 · p. 22
• The approach was evaluated for code synthesis, translation, and repair, with qualitative exploration of its positive impact on code generation.Source: §8 · p. 22
• The authors conclude that type constraining should be implemented for more programming languages.Source: §8 · p. 22
• The authors say type constraining has the potential to improve code generation across many domains.Source: §8 · p. 22

A Detailed Prefix Automaton Definitions

• The section provides detailed definitions and analysis of the various prefix automata for L_B.Source: A Detailed Prefix Automaton Definitions · p. 26

A.1 Base Automata

• Union combines the automata’s initial and accepting states and transitions, yielding the union of their languages while preserving the prefix-automaton property.The proof uses a lemma that reachable states from a state set equal the union of the states reachable from each member.Source: §A.1 · p. 26
• Concatenation preserves each automaton’s parsing behavior and adds transitions from accepting states of the first to initial states of the second, recognizing concatenated words.The result is a prefix automaton when both inputs are prefix automata and the second language is nonempty.Source: §A.1 · pp. 26–27
• Kleene-Star makes initial states accepting and adds transitions from final states back to initial states, accepting repetitions including the empty word.If the input is a prefix automaton, the resulting automaton is also a prefix automaton; nonemptiness of the input language is not required.Source: §A.1 · p. 27
• The terminal automaton uses suffix states and transitions that remove the consumed character, accepting exactly its specified terminal and forming a prefix automaton.The section modifies it to allow arbitrary whitespace before the terminal, making parsing syntactically indifferent to whitespace between terminals.Source: §A.1 · p. 27

A.2 Expressions

• Each state tracks its type environment, left-hand side, and any constrained type; accepting states record the parsed expression’s type, which is checked against the environment.State information is passed through transitions, and the current state of a parsed automaton is its final accepting state.Source: §A.2 · p. 27
• Arithmetic, function-call, and member-access automata use the left-hand side to determine valid extensions and expression types.For an extending expression, the left-hand side’s type and the extension’s type determine the resulting type.Source: §A.2 · pp. 27–28
• Function-call automata are valid only when the left-hand side is a function, and they accept only signatures compatible with the argument and result types.The argument is constrained by the function’s parameter type, while the call expression receives its return type.Source: §A.2 · p. 28
• Member-access automata are formed as a union of automata that parse the left-hand side’s attributes.The type-restricted versions are defined in §3.4 and are not separately listed here.Source: §A.2 · p. 28

A.3 Pruning the Type Search

• The heuristic exploits the restriction that extensions reach only the current type, operator results, call return types, and member types; higher-order types do not unlock types beyond their return type.Thus, a higher-order candidate is useful when it enables a new type or is needed to reach a higher-order goal.Source: §A.3 · p. 28
• The heuristic stops exploring a candidate when its depth exceeds the greater of the current and goal type depths and its root types are already explored.Depth measures function order; root types are the constituent types of minimal depth.Source: §A.3 · pp. 28–29
• Higher-order candidates are retained when they expose novel root types, while candidates with no unexplored root types are pruned.The explored-root set is updated with the root types encountered during search.Source: §A.3 · p. 29
• The depth and root-type restrictions draw on analysis of succinct-type search by Gvero et al., but lookup complexity prevents a completeness guarantee.The authors instead cite their evaluation as empirical evidence of high search coverage.Source: §A.3 · p. 29

A.4 Implementation of derivable

• The final derivable-type search uses two heuristics to make exploration finite and prioritize paths toward the target type.Without bounding the search, derivable may enumerate infinitely many reachable higher-order types.Source: §A.4 · pp. 29–30
• Bounding derivable types by the pruned reachability graph for target G is sound but potentially incomplete, while retaining practically relevant types.The bound limits type complexity to types present in the finite graph reachable for G.Source: §A.4 · p. 29
• Integrating derivable-type discovery with reachability lets the search transition directly from a discovered body type to a function type, prioritizing depth toward G.This enables efficient discovery of a path to G when one exists.Source: §A.4 · p. 30

A.5 Statements

• The recursive statement automaton combines declarations, expressions, returns, blocks, functions, and conditionals; declarations reject existing variable names.If-Then-Else uses concatenation, while expression statements and blocks reuse their respective automata.Source: §A.5 · p. 30
• Function-body automata propagate the expected return type and track whether statements and branches have returned, with return statements enabled only inside functions.The expected return type is set from the function signature; nested blocks and conditionals update return-state information.Source: §A.5 · p. 30
• When a block must return but has not yet returned, its automaton forces another statement rather than accepting the incomplete block.A return statement can always be generated to fulfill the requirement because the requested type can be expressed through literals.Source: §A.5 · p. 30
• Forcing completion of required returns preserves the prefix-automaton property.The construction relies on the ability to generate a return statement and express the requested type through literals.Source: §A.5 · p. 30

Implementation Details

• The implementation uses the regex library for regular-expression literal automata and transformers for LLM inference, and reports supported and unsupported TypeScript features in Tables 5 and 6.Source: B Details about Experimental Evaluation · p. 31

Hyperparameters

• Runs used A100 GPUs with 80 GB VRAM, temperature 1, seeds 0–4 for HumanEval and 0 for MBPP, a 1000-token limit, and a 300-second timeout.Syntactic correctness was measured with Oxidation because the official TypeScript compiler does not clearly distinguish syntax from semantic errors.Source: B Details about Experimental Evaluation · p. 31

Excluded MBPP Instances

• 6 MBPP instances were excluded because their TypeScript translations used overly broad any annotations: mbpp_405_check_tuplex, mbpp_563_extract_values, mbpp_580_extract_even, mbpp_612_merge, mbpp_725_extract_quotation, and mbpp_791_remove_nested.Earlier invalid nested-tuple translations were resolved in the latest MBPP version and included in the evaluation.Source: B Details about Experimental Evaluation · p. 31

Complete Prompts

• Task prompts are templates, and repair prompts include the non-compilable output followed by compiler feedback with line numbers for localization.Completions start from pre-filled function signatures; Qwen2.5 32B’s prompt additionally instructs it not to include test cases.Source: B Details about Experimental Evaluation · p. 31

Extracting Output Code

• The extraction heuristic selects a TypeScript code block, then cuts off after the last balanced closing brace followed by a newline or semicolon.If the code block is unclosed or no qualifying brace is found, the output is not pruned.Source: B Details about Experimental Evaluation · p. 31

C Case Study Full Outputs

• In HumanEval select_words, unconstrained Gemma 2 2B [64] calls split without parameters, which is legal in Python but not TypeScript; Figure 16 compares it with type-constrained generation.The task selects words containing a specified number of consonants.Source: C Case Study Full Outputs · p. 36, Figure 16 · p. 36
• For HumanEval largest_divisor, unconstrained DeepSeek Coder 33B Instruct [28] can return undefined on negative inputs, violating type-system rules; Figure 17 compares it with type-constrained generation.The unconstrained function loops over numbers below n to find the largest divisor and is functionally correct for positive inputs.Source: C Case Study Full Outputs · p. 37, Figure 17 · p. 37
• Figure 18 compares unconstrained Gemma 2 9B [73] with type-constrained generation for HumanEval intersperse, which inserts a delimiter number between consecutive list elements.Source: Figure 18 · p. 38