Learning to Guarantee Type Correctness in Code Generation through Type-Guided Program Synthesis

Language models have shown remarkable proficiency in code generation; nevertheless, ensuring type correctness remains a challenge. Although traditional methods, such as constrained decoding, alleviate this problem by externally rejecting untypable code, the model itself does not effectively learn type reasoning internally, which ultimately limits its overall performance. This paper introduces TyFlow, a novel system that internalizes type reasoning within code generation to guide the model to learn the type system. The core of our approach is a novel type-guided program synthesis system that maintains an isomorphism between type derivation trees and synthesis derivation trees, enabling a new code representation based on synthesis decision sequences rather than traditional text-based token sequences. By offloading the complexity of type system learning to the representation itself, models can redirect their computational resources toward higher-level program semantics. Our evaluation shows that TyFlow not only eliminates type errors but also significantly improves functional correctness, highlighting the importance of aligning LMs with type systems internally.

💡 Research Summary

The paper addresses a persistent problem in neural code generation: despite impressive performance of large language models (LLMs) in producing syntactically plausible code, a substantial fraction of generated programs fail due to type errors. Existing mitigation strategies such as constrained decoding simply filter out ill‑typed outputs after generation, offering no learning signal for the model to internalize the type system. Consequently, models continue to allocate capacity to low‑level type reasoning rather than higher‑level algorithmic reasoning.

TyFlow proposes a fundamentally different paradigm. The authors observe that, under constructive logic, an existential proof of well‑typedness (∃p. welltyped(p)) necessarily constructs a witness program p. By treating the construction of such a proof as a program synthesis process, they replace the traditional token‑by‑token generation with a sequence of synthesis decisions that are tightly coupled to type derivations. Each decision either applies a typing rule to decompose a typing goal into sub‑goals or instantiates a variable in the current context. Because the proof and the program are built simultaneously, the decision sequence encodes both the program text and its type‑correctness proof.

The key technical contributions are:

1. Isomorphic Representation – The authors formalize a bijection between type derivation trees and synthesis derivation trees. This guarantees that every well‑typed program corresponds to a unique decision sequence and vice‑versa, satisfying the “Data Usability” requirement.

2. Typed‑Guided Synthesis Rules – Typing rules of the target language are translated into constrained Horn clauses that serve as synthesis rules. The paper demonstrates this translation for the simply‑typed λ‑calculus (λ→) and for Java, showing how rules such as T‑Var, T‑Abs, and T‑App become actionable decisions.

3. Encoder‑Decoder Architecture – An encoder processes both the natural‑language specification and the current synthesis goal (including the local typing context). The decoder autoregressively emits decision tokens (e.g., “ApplyRule(T‑App)”, “InstantiateVar(x)”). Because each decision only requires a bounded local context, the model avoids costly global reasoning, fulfilling the “Context Locality” and “Derivation Vicinality” desiderata.

4. TyFlow System – TyFlow automates the entire pipeline: given a language definition (syntax, typing rules, parser, type‑checker) and a dataset of prompt‑program pairs, it extracts decision sequences from existing programs (using the type‑checker to obtain proofs) and trains the model. At inference time, the model generates a decision sequence, which TyFlow deterministically reconstructs into a well‑typed program together with its proof.

Empirical evaluation covers two domains. In Java, a baseline CodeT5 model exhibits a 24 % compilation‑error rate (largely type errors). TyFlow reduces type errors to 0 % and markedly improves unit‑test pass rates. In the λ→ setting, TyFlow similarly guarantees type correctness while achieving higher functional accuracy than a token‑based baseline. These results indicate that internalizing type reasoning frees model capacity for semantic reasoning, leading to better overall performance.

The paper’s broader significance lies in demonstrating that static‑analysis properties can be baked directly into the generation representation, rather than treated as an after‑the‑fact filter. While the current work focuses on type systems, the underlying synthesis framework is generic and could be extended to enforce other decidable safety properties such as memory safety, ownership, or even domain‑specific invariants. By aligning LLMs with the formal semantics of the target language, TyFlow opens a path toward more reliable, semantically aware code synthesis.