Typechecker Zoo

This is a pet project of mine I’ve been working on for a while. We’re going to create minimal implementations of the most successful static type systems of the last 30 years. This will involve making toy implementations of programming languages and the core typechecking algorithms. These obviously have evolved a lot over the years, so we’ll start with the simple ones and proceed all the way up to modern dependent types.

We’re going to implement them all in Rust. For no particularly reason, other than it has a decent parser ecosystem and is easy to install. And I like the ironic synthesis of building pure functional languages in a language which is decidedly non-functional. It’s a bit of a heaven-and-hell thing going on, and I’ll leave it up to you to decide on the chirality of that metaphor.

This is going to be a more a fun weekend side project, than a formal introduction to these systems. If you want resources read TAPL and ATTAPL for the theory behind these systems. We’ll going for implementation, fun, and rock and roll. But I’ll link to the primary resources if you want the gory details.

The four little critters we’re going to build are:

Algorithm W (775 lines of code)

Robin Milner’s classic Hindley-Milner type inference algorithm from A Theory of Type Polymorphism in Programming. A toy polymorphic lambda calculus.

System F (1090 lines of code)

Second-order lambda calculus with parametric polymorphism using bidirectional type checking. A Mini-OCaml

An implementation of DK algorithm from Complete and Easy Bidirectional Typechecking for Higher-rank Polymorphism by Dunfield and Krishnaswami.

System F-ω (3196 lines of code)

Complete implementation of System F-ω with higher-kinded types, DK bidirectional type checking, existential type variables, polymorphic constructor applications, pattern matching, and datatypes. A Haskell-lite.

Uses the method of A Mechanical Formalization of Higher-Ranked Polymorphic Type Inference by Zhao et al.

Calculus of Constructions (6000 lines of code)

The Calculus of Constructions with a countable hierarchy of non-cumulative universes and inductive types. A teeny Lean-inspired dependent type checker.

Uses a bidirectional dependent typechecker outlined in A Universe Polymorphic Type System by Vladimir Voevodsky.

Type Systems Overview

So why do we build type systems? The answer is obviously because they are awesome and intellectually interesting. But second to that, because they are genuinely useful.

Our exploration starts with the lambda calculus, a formal system developed by Alonzo Church in the 1930s to express computation. It is the minimal, universal programming language. Its syntax consists of just three elements: variables, function abstractions (a way to define an anonymous function), and function application (the act of calling a function). For instance, the identity function, which simply returns its input, is written as \( \lambda x. x \). Despite this simplicity, any computable problem can be expressed and solved within the lambda calculus, making it Turing complete.

Next, we introduce type systems. A type system is a set of rules that assigns a property, known as a type, to the constructs of a program, such as variables, expressions, and functions. The primary purpose is to reduce bugs by preventing operations that don’t make sense, like dividing a number by a string. This process of verifying that a program obeys its language’s type rules is called type checking. Type checking can be performed at compile-time, known as static typing, or during program execution, known as dynamic typing. By enforcing these rules, type systems help ensure that different parts of a program connect in a consistent and error-free way.

This leads to the principle of well-typed programs. A program is called well-typed if it conforms to the rules of its type system. This can be expressed formally with a typing judgment, \( E \vdash M : A \), which asserts that in a given context \( E \), the program expression \( M \) has the type \( A \). The foundational promise of this approach was articulated by Robin Milner in his 1978 paper, A Theory of Type Polymorphism in Programming, with the phrase “Well-typed programs cannot go wrong”. This means a program that passes the type checker is guaranteed to be free of a certain class of runtime errors. Later extensions of this idea, such as “Well-typed programs don’t get stuck” and “Well-typed programs can’t be blamed,” further underscore the safety and reliability that robust type systems provide to software development.

A System of Types

At the heart of any type system is a set of formal rules for making logical deductions about a program. These deductions are called judgments. A judgment is an assertion that a piece of code has a certain property. The most common kind of judgment you will encounter is a typing judgment, which asserts that a given expression has a particular type within a specific context. We write this formally using a “turnstile” symbol \( \vdash \). The general form of a typing judgment looks like this:

\[ \Gamma \vdash e : T \]

This statement is read as, “In the context \( \Gamma \), the expression \( e \) has the type \( T \).” The context, represented by the Greek letter Gamma (\( \Gamma \)), is crucial. It acts as an environment that keeps track of the types of all the variables that are currently in scope. Our goal when type checking a program is to construct a valid derivation that proves this judgment holds for the entire program. Think of it as a giant dictionary that maps variable names to their types, that’s basically all it is.

The context \( \Gamma \) is essentially a map from variable names to their types. For example, \( x: \text{Int}, f: \text{Bool} \to \text{Int} \) is a context where the variable \( x \) has type \( \text{Int} \) and \( f \) has a function type from \( \text{Bool} \) to \( \text{Int} \). As we enter deeper scopes in a program, like inside a function body, we extend the context with new variable bindings. This is often written as \( \Gamma, x:T \), which means “the context \( \Gamma \) extended with a new binding stating that variable \( x \) has type \( T \).”

Typing judgments are derived using inference rules. An inference rule states that if you can prove a set of judgments, called the premises, then you can conclude another judgment, called the conclusion. Rules are written with the premises on top of a horizontal line and the conclusion below it. You should read them from top to bottom: “If all premises above the line are true, then the conclusion below the line is also true.” If a rule has no premises, it is an axiom, a self-evident truth that requires no prior proof to hold. These rules are the fundamental building blocks we use to reason about the types in a program.

When working with formal type systems, inference rules follow a consistent structural pattern that makes them easier to read and understand. Every rule has the form:

\[ \frac{\text{premises}}{\text{conclusion}} \text{(Rule-Name)} \]

This notation should be read as: “If all the premises above the line are true, then the conclusion below the line is also true.” The premises represent the conditions that must be satisfied, while the conclusion represents what we can deduce when those conditions hold. The rule name provides a convenient label for referencing the rule in discussions and proofs.

For example, an axiom rule that establishes the type of the literal zero might look like:

\[ \frac{}{\Gamma \vdash 0 : \text{Nat}} \text{(T-Zero)} \]

This rule has no premises above the line, making it an axiom. It simply states that in any context \( \Gamma \), the literal \( 0 \) has type \( \text{Nat} \). This is a fundamental fact that requires no further proof.

These axioms typically handle simple cases like variable lookups or literal values (sometimes called ground types). Rules with multiple premises, separated by spacing or explicit conjunction symbols, require all conditions to be satisfied simultaneously before the conclusion can be drawn.

A foundational rule in nearly every type system is the variable lookup rule, which lets us find the type of a variable from the context:

\[ \frac{x:T \in \Gamma}{\Gamma \vdash x : T} \]

This rule is an axiom because it has no premises above the line. It reads: “If the type binding \( x:T \) is present in the context \( \Gamma \), then we can conclude that in context \( \Gamma \), the expression \( x \) has type \( T \).” It formally defines the action of looking up a variable’s type in the current environment. For example the axiom that 1 has type \( \text{Int} \):

\[ \frac{}{\Gamma \vdash 1 : \text{Int}} \]

By defining a collection of these inference rules, we create a complete type system. Each rule defines how to determine the type of a specific kind of expression, like a function call, a literal value, or an if-then-else block. For instance, a rule for function application would require as its premises that we first prove the function itself has a function type \( T \to U \) and that its argument has the corresponding input type \( T \). If we can prove those premises, the rule allows us to conclude that the entire function application expression has the output type \( U \). By repeatedly applying these rules, we can build a derivation tree that starts from axioms about variables and literals and culminates in a single judgment about the type of our entire program, thereby proving it is well-typed.

Judgements

A common inference rule with multiple premises is the one for function application, which determines the type of a function call. In lambda calculus, this is simply an expression \( e_1 \ e_2 \), where \( e_1 \) is the function and \( e_2 \) is the argument. To assign a type to this expression, we must first determine the types of both the function and its argument. The rule, often called “application” or “elimination for functions” (\( \to E \)), is written as follows:

\[ \frac{\Gamma \vdash e_1 : T_1 \to T_2 \quad \Gamma \vdash e_2 : T_1}{\Gamma \vdash e_1 \ e_2 : T_2} \]

This rule has two premises above the line. It states that to conclude that the application \( e_1 \ e_2 \) has type \( T_2 \), we must first prove two things in the same context \( \Gamma \). First, we must show that \( e_1 \) has a function type, written as \( T_1 \to T_2 \), which means it takes an input of type \( T_1 \) and produces an output of type \( T_2 \). Second, we must show that the argument \( e_2 \) has the correct input type \( T_1 \). If both of these premises hold, the rule allows us to deduce that the entire expression \( e_1 \ e_2 \) results in the function’s output type, \( T_2 \).

Now, let’s look at an example that chains three judgments together to type check the expression \( f \ x \). We’ll work within a context \( \Gamma \) that contains bindings for both \( f \) and \( x \), specifically \( \Gamma = f:\text{Int} \to \text{Bool}, x:\text{Int} \). Our goal is to prove that \( \Gamma \vdash f \ x : \text{Bool} \).

Our derivation begins with two simple variable lookups, which are our axioms. These will form the premises of our application rule:

1. First Judgment (Variable Lookup for \( f \)): We use the variable rule to find the type of \( f \). Since \( f:\text{Int} \to \text{Bool} \) is in our context \( \Gamma \), we can conclude: \[ \frac{f:\text{Int} \to \text{Bool} \in \Gamma}{\Gamma \vdash f : \text{Int} \to \text{Bool}} \]

2. Second Judgment (Variable Lookup for \( x \)): Similarly, we look up the type of \( x \). The binding \( x:\text{Int} \) is in \( \Gamma \), so we can conclude: \[ \frac{x:\text{Int} \in \Gamma}{\Gamma \vdash x : \text{Int}} \]

3. Third Judgment (Function Application): Now we have the necessary premises to use the function application rule. We substitute \( e_1 \) with \( f \), \( e_2 \) with \( x \), \( T_1 \) with \( \text{Int} \), and \( T_2 \) with \( \text{Bool} \). Since our first two judgments successfully proved the premises, we can now form the final conclusion: \[ \frac{\Gamma \vdash f : \text{Int} \to \text{Bool} \quad \Gamma \vdash x : \text{Int}}{\Gamma \vdash f \ x : \text{Bool}} \]

Putting all of this together, we can represent the full derivation as a tree of nested inference rules, showing how the final judgment is built from the axioms for variable lookup:

\[ \frac{ \frac{f:\text{Int} \to \text{Bool} \in \Gamma}{\Gamma \vdash f : \text{Int} \to \text{Bool}} \quad \frac{x:\text{Int} \in \Gamma}{\Gamma \vdash x : \text{Int}} }{ \Gamma \vdash f \ x : \text{Bool} } \]

This nested structure is called a derivation tree or inference tree. Each node in the tree corresponds to an application of an inference rule, and the leaves are axioms (variable lookups from the context). The tree visually demonstrates how the type checker starts from basic facts (the types of variables in the context) and applies rules step by step to reach a conclusion about the type of a complex expression. In this example, the root of the tree is the judgment \( \Gamma \vdash f \ x : \text{Bool} \), and its two children are the judgments for \( f \) and \( x \), each justified by their respective axioms. This process generalizes to larger programs, where the derivation tree grows to reflect the structure of the program and the logical flow of type information.

Terms and Types

Historically, many programming languages enforced a strict separation between different layers of abstraction, a concept known as stratification. In this model, you had a “term language” and a “type language” which were syntactically and conceptually distinct. The term language consists of the expressions that actually compute values and run at runtime, like \( 5 + 2 \) or \( \text{if } x \text{ then } y \text{ else } z \). The type language, on the other hand, consists of expressions that describe the terms, like \( \text{Int} \) or \( \text{Bool} \to \text{Bool} \). These two worlds were kept separate; a type could not appear where a term was expected, and vice versa.

This stratification could be extended further. To bring order to the type language itself, a third layer called the “kind language” was often introduced. A kind can be thought of as the “type of a type.” For example, a concrete type like \( \text{Int} \) has the simplest kind, \( * \) (often pronounced “type”). But a type constructor like \( \text{List} \) isn’t a type on its own; it’s something that takes a type and produces a new one. \( \text{List} \) takes \( \text{Int} \) to produce \( \text{List} \ \text{Int} \). Therefore, its kind is \( * \to * \). This creates a rigid hierarchy: terms are classified by types, and types are classified by kinds. This clear separation makes type checking simpler and more predictable.

However, a major trend in the design of modern, powerful type systems has been to move away from this strict stratification and instead unify the term and type languages into a single, consistent syntactic framework. In these unified systems, the line between what is a “value” (a term) and what is a “description” (a type) begins to blur. The language’s grammar allows expressions that can be interpreted at different “levels” or “universes.” For instance, you might have a universe \( \text{Type}_0 \) which contains simple types like \( \text{Bool} \). Then you would have a higher universe, \( \text{Type}_1 \), whose only member is \( \text{Type}_0 \). This allows you to write functions that operate on types themselves, a key feature of dependently typed languages. More on this later.

To make this concrete, let’s consider a simple, stratified language for arithmetic. We can define its term language (\( e \)) and its corresponding type language (\( \tau \)) separately. The terms are the expressions we can compute, and the types are the static labels we can assign to them. Their definitions might look like this:

\[ \begin{align*} \text{terms} \quad e &::= n \mid e_1 + e_2 \mid \text{iszero}(e) \mid \text{true} \mid \text{false} \mid \text{if } e_1 \text{ then } e_2 \text{ else } e_3 \\ \text{types} \quad \tau &::= \text{Nat} \ \mid \ \text{Bool} \end{align*} \]

Here, the term language \( e \) defines natural numbers (\( n \)), addition, a function to check for zero, boolean constants, and a conditional. The type language \( \tau \) is much simpler; it only contains the types \( \text{Nat} \) and \( \text{Bool} \). In this stratified system, an expression like \( \text{Nat} + 5 \) would be a syntax error because \( \text{Nat} \) belongs to the type language and cannot be used in a term-level operation like addition. In a more modern, unified system, this rigid distinction would be relaxed.

Type Checking and Type Reconstruction

While the process of verifying that a program adheres to its type rules is called type checking, a related and historically significant challenge is type reconstruction, more commonly known as type inference. The goal of type inference is to have the compiler automatically deduce the types of expressions without the programmer needing to write explicit type annotations. For many years, developing algorithms that could perform full type inference for increasingly expressive languages was an active and vital area of research. The promise was seductive: achieve all the safety guarantees of a static type system without the verbose, manual effort of annotating every variable and function.

In more recent decades, the focus on achieving complete type inference has diminished. The primary reason is that as type systems have grown more powerful and complex, full inference has become computationally intractable or, in many cases, fundamentally undecidable. Modern languages often include features like higher-rank polymorphism (passing polymorphic functions as arguments), GADTs, and various forms of type-level programming where types themselves can involve computation. For these systems, a general algorithm that can always infer the single “best” type for any given expression simply does not exist. Attempting to do so would lead to impossibly complex algorithms and can result in inferred types that are enormous and incomprehensible to the programmer.

As a result, many modern statically-typed languages have converged on a practical and elegant middle ground: bidirectional type checking. Instead of having a single mode that always tries to infer types, a bidirectional checker operates in two distinct modes: a “checking” mode and a “synthesis” mode.

1. Checking Mode: In this mode, the algorithm verifies that an expression \( e \) conforms to a known, expected type \( \tau \). Information flows “down” from the context into the expression. We ask the question: “Can we prove that \( e \) has type \( \tau \)?”
2. Synthesis Mode: In this mode, the algorithm computes, or “synthesizes,” a type for an expression \( e \) without any prior expectation. Information flows “up” from the expression’s components. Here, we ask: “What is the type of \( e \)?”

This duality provides a powerful framework. The language designer can specify which syntactic constructs require annotations and which do not. For example, the arguments of a top-level function might require explicit annotations (putting the checker in checking mode for the function body), but the types of local variables within that body can be inferred (synthesized). This approach neatly sidesteps the difficulties of full inference by requiring the programmer to provide annotations only at key boundaries where ambiguity would otherwise arise. It offers a “best of both worlds” scenario: the convenience of local inference with the clarity and power of explicit annotations for complex, polymorphic, or ambiguous parts of the code, representing a theoretical sweet spot that balances expressiveness, usability, and implementability. More on this later.

The Frontiers

A deep insight has been unfolding since the late 1970s, a revelation that three immense, seemingly distinct realms of human thought are echoes of a single, deeper reality. This “computational trinitarianism” reveals a profound interplay between different disciplines:

• Computation
• Formal logic
• Category theory

Think of it as a three way harmony between Spaces, Logic and Computation.

These are not merely analogous; they are three perspectives on one underlying phenomenon. The key to this unification is a powerful act of abstraction. While mathematics seems to be about many different things like sets and shapes, it fundamentally relies on the notion of a map. Set theory is about functions, topology is about continuous maps, algebra is about homomorphisms, and logic is about implication.

With the right framework, we can abstract away the specific details. What matters are not the objects themselves, but the maps between them. When you redefine a concept in this universal language, something incredible happens: you can translate it into any other category. A “group” defined in the category of sets is just a normal group. But translate that same definition to the category of topological spaces, and you get a topological group. In the category of manifolds, you get a Lie group. This is the Rosetta Stone for the foundations of abstraction itself.

This is why this is an exciting time in software. These concepts elevate types from simple data labels into a rich language for encoding deep truths about a program’s behavior. These are no longer just academic curiosities; they are the architectural principles for the next generation of powerful and provably correct software.

This grand unification is an extension of the famous Curry-Howard correspondence, which revealed a stunning duality: a type is a logical proposition, and a program is a constructive proof of that proposition. Trinitarianism expands this into a three-fold revelation:

• A Proposition in logic corresponds to a Type in a programming language, which corresponds to an Object (like a space) in a category.
• A Proof of a proposition corresponds to a Program of a type, which corresponds to a Map (or morphism) in a category.
• The Simplification of a proof corresponds to the Computation (or evaluation) of a program.

Under this paradigm, programming becomes the act of discovering and articulating mathematical truth. Consider a simple example: a function to get the first element of a list. A naive type signature might be head(List<T>) -> T. This is a weak proposition: “Give me a list of things, and I’ll give you a thing.” This proposition is false, as the program will crash on an empty list.

An expressive type system lets us state a more honest proposition: safeHead(List<T>) -> Maybe<T>. The Maybe<T> type embodies the possibility of failure. This type is a stronger proposition: “Given any list, I can provide a context that maybe contains a thing.” A program satisfying this type is a proof that you have handled all possibilities, including the empty one. The compiler, by checking the type, verifies your proof and mathematically guarantees your program cannot crash from this error.

Now for a more profound example. Imagine a type for a provably sorted list: SortedList<T>. A function to merge two such lists has a type that is a powerful theorem: merge(SortedList<T>, SortedList<T>) -> SortedList<T>. This proposition states, “If you give me two proofs that lists are sorted, I will produce a new proof that the resulting merged list is also sorted.” The compiler will not accept a program for this type unless its logic correctly preserves the “sorted” property. The type system becomes a partner in reasoning, ensuring deep logical properties hold true.

This idea of structure-preserving transformations is a core engine of modern mathematics. In algebraic topology, a functor acts as a map between entire universes of mathematical objects. The homology functor, for example, translates topological spaces into algebraic groups, allowing geometric problems to be solved with algebra. The insight of trinitarianism is that a well-typed program is a functor: a concrete, computable map that translates a logical proof into a computational process while perfectly preserving the underlying structure.

This theoretical bedrock unleashes extraordinary practical power. When types are precise logical statements, composing programs is like composing theorems, allowing us to build massive systems with certainty. The ultimate ambition this enables is program synthesis. The programmer’s role shifts from mechanic to architect. They specify the what by writing a rich, descriptive type (a theorem), and the system discovers the how by searching for a valid proof, which is the program itself.

These frontiers are critical in the age of AI-generated code. As we task AI with generating vast software systems, the central challenge becomes one of trust. How do we know the code is correct? The answer is to demand more from the AI. We require it to produce not just code, but code whose very type is a formal proof of its correctness. The type checker then becomes our ultimate, infallible arbiter of truth.

This transforms programming from a manual craft into a discipline of rigorous specification, allowing us to build software whose correctness is not just hoped for, but mathematically assured. This is the awesome power of modern type systems: the journey from preventing simple errors to proving deep truths, turning the art of programming into a direct collaboration with logic itself.

And that’s why type systems are awesome, and hopefully you come to see how much fun they can be!

The Lambda Calculus

The lambda calculus is computation in its purest form. It is not an exaggeration to say that this elegant, tiny system is the theoretical bedrock of nearly every functional programming language, and its influence is felt across the entire landscape of computer science. Developed by Alonzo Church in the 1930s as a tool for studying the foundations of mathematics, it was later understood to be a universal model of computation. At its core, the lambda calculus is shockingly minimal; its syntax defines only three kinds of expressions, or “terms.”

Formally, we can define the syntax of the pure lambda calculus as follows:

\[ e ::= x \mid \lambda x . e \mid e_1 \ e_2 \]

Let’s break this down:

1. Variable (\(x\)): A name that acts as a placeholder for a value.
2. Abstraction (\(\lambda x . e\)): This is an anonymous function definition. The \(\lambda x\) is the function’s parameter, and the expression \(e\) is its body. The \(\lambda\) is read as “lambda.”
3. Application (\(e_1 \ e_2\)): This is the act of calling a function. The expression \(e_1\) is the function, and \(e_2\) is the argument being passed to it.

This formal definition maps almost directly onto a data structure in a language like Rust. We can represent these three core terms using an enum:

#![allow(unused)]
fn main() {
pub enum Expr {
    Var(String),
    Abs(String, Box<Expr>),
    App(Box<Expr>, Box<Expr>),
}
}

#![allow(unused)]
fn main() {
#[derive(Debug, Clone, PartialEq)]
pub enum Lit {
    Int(i64),
    Bool(bool),
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    Var(String),
    Arrow(Box<Type>, Box<Type>),
    Int,
    Bool,
    Tuple(Vec<Type>),
}
#[derive(Debug, Clone, PartialEq)]
pub struct Scheme {
    pub vars: Vec<String>, // Quantified type variables
    pub ty: Type,          // The type being quantified over
}
impl std::fmt::Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Expr::Var(name) => write!(f, "{}", name),
            Expr::Lit(Lit::Int(n)) => write!(f, "{}", n),
            Expr::Lit(Lit::Bool(b)) => write!(f, "{}", b),
            Expr::Abs(param, body) => write!(f, "λ{}.{}", param, body),
            Expr::App(func, arg) => match (func.as_ref(), arg.as_ref()) {
                (Expr::Abs(_, _), _) => write!(f, "({}) {}", func, arg),
                (_, Expr::App(_, _)) => write!(f, "{} ({})", func, arg),
                (_, Expr::Abs(_, _)) => write!(f, "{} ({})", func, arg),
                _ => write!(f, "{} {}", func, arg),
            },
            Expr::Let(var, value, body) => {
                write!(f, "let {} = {} in {}", var, value, body)
            }
            Expr::Tuple(exprs) => {
                write!(f, "(")?;
                for (i, expr) in exprs.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", expr)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl std::fmt::Display for Type {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Type::Var(name) => write!(f, "{}", name),
            Type::Int => write!(f, "Int"),
            Type::Bool => write!(f, "Bool"),
            Type::Arrow(t1, t2) => {
                // Add parentheses around left side if it's an arrow to avoid ambiguity
                match t1.as_ref() {
                    Type::Arrow(_, _) => write!(f, "({}) → {}", t1, t2),
                    _ => write!(f, "{} → {}", t1, t2),
                }
            }
            Type::Tuple(types) => {
                write!(f, "(")?;
                for (i, ty) in types.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", ty)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl std::fmt::Display for Scheme {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.vars.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "forall {}. {}", self.vars.join(" "), self.ty)
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    Var(String),
    App(Box<Expr>, Box<Expr>),
    Abs(String, Box<Expr>),
    Let(String, Box<Expr>, Box<Expr>),
    Lit(Lit),
    Tuple(Vec<Expr>),
}
}

Here, Var(name) corresponds to \(x\), Abs(param, body) corresponds to \(\lambda x . e\), and App(function, argument) corresponds to \(e_1 \ e_2\). The Box is a Rust detail that allows us to have recursive types of a known size.

The power of the lambda calculus emerges from a few fundamental concepts. The first is variable binding. In an abstraction like \(\lambda x . x + 1\), the variable \(x\) is said to be bound within the body of the lambda. Any variable that is not bound by an enclosing lambda is a free variable. This distinction is critical for understanding scope. This leads directly to the idea of alpha equivalence, which states that the name of a bound variable is irrelevant. The function \(\lambda x . x\) is semantically identical to \(\lambda y . y\); both are the identity function. An implementation must be able to recognize this equivalence to correctly handle variable naming.

The second core concept is beta reduction, which is the computational engine of the lambda calculus. It formally defines how function application works. When an abstraction is applied to an argument, we reduce the expression by substituting the argument for every free occurrence of the bound variable within the function’s body. For example, applying the function \(\lambda x . x + 1\) to the argument \(5\) is written as \((\lambda x . x + 1) \ 5\). The beta reduction rule tells us to replace \(x\) with \(5\) in the body \(x + 1\), yielding the result \(5 + 1\). This process of substitution is the fundamental mechanism of computation in this system. Implementations often avoid direct string substitution due to its complexity and risk of name collisions, instead using techniques like de Bruijn indices, which replace variable names with numbers representing their lexical distance to their binder.

While the pure lambda calculus is Turing complete, it is also notoriously difficult to use directly for practical programming. For example, representing the number \(3\) requires a complex expression like \(\lambda f . \lambda x . f (f (f x))\). To make programming more convenient, we almost always extend the core calculus with additional expression types. These can include let bindings for local variables, literals for concrete values like numbers and strings, and data structures like tuples.

#![allow(unused)]
fn main() {
#[derive(Debug, Clone, PartialEq)]
pub enum Lit {
    Int(i64),
    Bool(bool),
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    Var(String),
    Arrow(Box<Type>, Box<Type>),
    Int,
    Bool,
    Tuple(Vec<Type>),
}
#[derive(Debug, Clone, PartialEq)]
pub struct Scheme {
    pub vars: Vec<String>, // Quantified type variables
    pub ty: Type,          // The type being quantified over
}
impl std::fmt::Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Expr::Var(name) => write!(f, "{}", name),
            Expr::Lit(Lit::Int(n)) => write!(f, "{}", n),
            Expr::Lit(Lit::Bool(b)) => write!(f, "{}", b),
            Expr::Abs(param, body) => write!(f, "λ{}.{}", param, body),
            Expr::App(func, arg) => match (func.as_ref(), arg.as_ref()) {
                (Expr::Abs(_, _), _) => write!(f, "({}) {}", func, arg),
                (_, Expr::App(_, _)) => write!(f, "{} ({})", func, arg),
                (_, Expr::Abs(_, _)) => write!(f, "{} ({})", func, arg),
                _ => write!(f, "{} {}", func, arg),
            },
            Expr::Let(var, value, body) => {
                write!(f, "let {} = {} in {}", var, value, body)
            }
            Expr::Tuple(exprs) => {
                write!(f, "(")?;
                for (i, expr) in exprs.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", expr)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl std::fmt::Display for Type {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Type::Var(name) => write!(f, "{}", name),
            Type::Int => write!(f, "Int"),
            Type::Bool => write!(f, "Bool"),
            Type::Arrow(t1, t2) => {
                // Add parentheses around left side if it's an arrow to avoid ambiguity
                match t1.as_ref() {
                    Type::Arrow(_, _) => write!(f, "({}) → {}", t1, t2),
                    _ => write!(f, "{} → {}", t1, t2),
                }
            }
            Type::Tuple(types) => {
                write!(f, "(")?;
                for (i, ty) in types.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", ty)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl std::fmt::Display for Scheme {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.vars.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "forall {}. {}", self.vars.join(" "), self.ty)
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    Var(String),
    App(Box<Expr>, Box<Expr>),
    Abs(String, Box<Expr>),
    Let(String, Box<Expr>, Box<Expr>),
    Lit(Lit),
    Tuple(Vec<Expr>),
}
}

This raises a crucial design question: should these new constructs be “wired in” as primitives with their own semantic rules, or should they be defined in terms of the pure calculus? This represents a fundamental tradeoff in language design. For example, a let expression, \(\text{let } x = e_1 \text{ in } e_2\), can be treated as syntactic sugar that desugars directly into a pure lambda calculus expression: \((\lambda x . e_2) \ e_1\). This approach keeps the core language minimal and elegant. The alternative is to make let a primitive construct. This means the type checker and evaluator must have specific logic to handle it. This “wired-in” approach often leads to better performance and more precise error messages, as the compiler has more specific knowledge about the construct’s intent. However, it increases the complexity of the core system. Many languages strike a balance: fundamental and performance-critical features like numeric literals are often wired in, while higher-level patterns like let bindings might be treated as syntactic sugar, offering a convenient syntax that ultimately translates back to the foundational concepts of lambda, application, and variable.

Core Implementation

Algorithm W represents one of the most elegant solutions to the type inference problem in functional programming languages. Developed by Robin Milner in 1978, it provides a sound and complete method for inferring the most general types in the Hindley-Milner type system. This chapter explores our Rust implementation, examining how the mathematical foundations translate into practical code that can handle lambda abstractions, function applications, let-polymorphism, and complex unification scenarios.

The core insight of Algorithm W lies in its systematic approach to type inference through constraint generation and unification. Rather than attempting to determine types through local analysis, the algorithm builds a global picture by generating type variables, collecting constraints, and then solving these constraints through unification. This approach ensures that we always find the most general type possible, a property crucial for supporting polymorphism in functional languages.

You’ll often see this algorithm (or the type system, confusingly enough) referred to by many names:

• Hindley-Milner
• Hindley-Damas-Milner
• Damas-Milner
• HM
• Algorithm W

Typing Rules

Before diving into the implementation details, let’s establish the formal typing rules that govern the Hindley-Milner type system. We’ll be introducing mathematical symbols that capture the essence of type inference, but don’t worry, each symbol has a precise and intuitive meaning once you dive into the details.

• \( \Gamma \) (Gamma) - The type environment, which maps variables to their types. It’s like a dictionary that remembers what we know about each variable’s type.

• \( \vdash \) (Turnstile) - The “entails” or “proves” symbol. When we write \( \Gamma \vdash e : \tau \), we’re saying “in environment \( \Gamma \), expression \( e \) has type \( \tau \).”

• \( \tau \) (Tau) - Represents monomorphic types like \( \text{Int} \), \( \text{Bool} \), or \( \text{Int} \to \text{Bool} \). These are concrete, fully-determined types.

• \( \sigma \) (Sigma) - Represents polymorphic type schemes like \( \forall \alpha. \alpha \to \alpha \). These can be instantiated with different concrete types.

• \( \forall \alpha \) (Forall Alpha) - Universal quantification over type variables. It means “for any type \( \alpha \).” This is how we express polymorphism.

• \( \alpha, \beta, \gamma \) (Greek Letters) - Type variables that stand for unknown types during inference. Think of them as type-level unknowns that get solved.

• \( [\tau/\alpha]\sigma \) - Type substitution, replacing all occurrences of type variable \( \alpha \) with type \( \tau \) in scheme \( \sigma \). This is how we instantiate polymorphic types.

• \( S \) (Substitution) - A mapping from type variables to types, representing the solutions found by unification.

• \( \text{gen}(\Gamma, \tau) \) - Generalization, which turns a monotype into a polytype by quantifying over type variables not present in the environment.

• \( \text{inst}(\sigma) \) - Instantiation, which creates a fresh monotype from a polytype by replacing quantified variables with fresh type variables.

• \( \text{ftv}(\tau) \) - Free type variables, the set of unbound type variables appearing in type \( \tau \).

• \( \emptyset \) (Empty Set) - The empty substitution, representing no changes to types.

• \( [\alpha \mapsto \tau] \) - A substitution that maps type variable \( \alpha \) to type \( \tau \).

• \( S_1 \circ S_2 \) - Composition of substitutions, applying \( S_2 \) first, then \( S_1 \).

• \( \notin \) (Not In) - Set membership negation, used in the occurs check to prevent infinite types.

Now that we have our symbolic toolkit, let’s see how these pieces work together to create the elegant machinery of Algorithm W.

Core Typing Rules

The variable rule looks up types from the environment: \[ \frac{x : σ \in Γ \quad τ = \text{inst}(σ)}{Γ ⊢ x : τ} \text{(T-Var)} \]

Lambda abstraction introduces new variable bindings: \[ \frac{Γ, x : α ⊢ e : τ \quad α \text{ fresh}}{Γ ⊢ λx. e : α → τ} \text{(T-Lam)} \]

Function application combines types through unification: \[ \frac{Γ ⊢ e₁ : τ₁ \quad Γ ⊢ e₂ : τ₂ \quad α \text{ fresh} \quad S = \text{unify}(τ₁, τ₂ → α)}{Γ ⊢ e₁ \, e₂ : S(α)} \text{(T-App)} \]

Let-polymorphism allows generalization: \[ \frac{Γ ⊢ e₁ : τ₁ \quad σ = \text{gen}(Γ, τ₁) \quad Γ, x : σ ⊢ e₂ : τ₂}{Γ ⊢ \text{let } x = e₁ \text{ in } e₂ : τ₂} \text{(T-Let)} \]

Literals have their corresponding base types: \[ \frac{}{Γ ⊢ n : \text{Int}} \text{(T-LitInt)} \]

\[ \frac{}{Γ ⊢ b : \text{Bool}} \text{(T-LitBool)} \]

These rules capture the essence of the Hindley-Milner type system, where we infer the most general types while supporting true polymorphism through let-generalization.

Abstract Syntax Trees

Our implementation begins with a careful modeling of both expressions and types as algebraic data types. The expression language extends the pure lambda calculus with practical constructs while maintaining the theoretical foundation.

#![allow(unused)]
fn main() {
#[derive(Debug, Clone, PartialEq)]
pub enum Lit {
    Int(i64),
    Bool(bool),
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    Var(String),
    Arrow(Box<Type>, Box<Type>),
    Int,
    Bool,
    Tuple(Vec<Type>),
}
#[derive(Debug, Clone, PartialEq)]
pub struct Scheme {
    pub vars: Vec<String>, // Quantified type variables
    pub ty: Type,          // The type being quantified over
}
impl std::fmt::Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Expr::Var(name) => write!(f, "{}", name),
            Expr::Lit(Lit::Int(n)) => write!(f, "{}", n),
            Expr::Lit(Lit::Bool(b)) => write!(f, "{}", b),
            Expr::Abs(param, body) => write!(f, "λ{}.{}", param, body),
            Expr::App(func, arg) => match (func.as_ref(), arg.as_ref()) {
                (Expr::Abs(_, _), _) => write!(f, "({}) {}", func, arg),
                (_, Expr::App(_, _)) => write!(f, "{} ({})", func, arg),
                (_, Expr::Abs(_, _)) => write!(f, "{} ({})", func, arg),
                _ => write!(f, "{} {}", func, arg),
            },
            Expr::Let(var, value, body) => {
                write!(f, "let {} = {} in {}", var, value, body)
            }
            Expr::Tuple(exprs) => {
                write!(f, "(")?;
                for (i, expr) in exprs.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", expr)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl std::fmt::Display for Type {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Type::Var(name) => write!(f, "{}", name),
            Type::Int => write!(f, "Int"),
            Type::Bool => write!(f, "Bool"),
            Type::Arrow(t1, t2) => {
                // Add parentheses around left side if it's an arrow to avoid ambiguity
                match t1.as_ref() {
                    Type::Arrow(_, _) => write!(f, "({}) → {}", t1, t2),
                    _ => write!(f, "{} → {}", t1, t2),
                }
            }
            Type::Tuple(types) => {
                write!(f, "(")?;
                for (i, ty) in types.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", ty)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl std::fmt::Display for Scheme {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.vars.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "forall {}. {}", self.vars.join(" "), self.ty)
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    Var(String),
    App(Box<Expr>, Box<Expr>),
    Abs(String, Box<Expr>),
    Let(String, Box<Expr>, Box<Expr>),
    Lit(Lit),
    Tuple(Vec<Expr>),
}
}

The expression AST captures the essential constructs of our language. Variables (Var) and function abstractions (Abs) correspond directly to the lambda calculus. Function application (App) drives computation through beta reduction. The Let construct introduces local bindings with potential for polymorphic generalization, while literals (Lit) and tuples (Tuple) provide concrete data types that make the language practical for real programming tasks.

The type system mirrors this structure with its own AST that represents the types these expressions can have.

#![allow(unused)]
fn main() {
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    Var(String),
    App(Box<Expr>, Box<Expr>),
    Abs(String, Box<Expr>),
    Let(String, Box<Expr>, Box<Expr>),
    Lit(Lit),
    Tuple(Vec<Expr>),
}
#[derive(Debug, Clone, PartialEq)]
pub enum Lit {
    Int(i64),
    Bool(bool),
}
#[derive(Debug, Clone, PartialEq)]
pub struct Scheme {
    pub vars: Vec<String>, // Quantified type variables
    pub ty: Type,          // The type being quantified over
}
impl std::fmt::Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Expr::Var(name) => write!(f, "{}", name),
            Expr::Lit(Lit::Int(n)) => write!(f, "{}", n),
            Expr::Lit(Lit::Bool(b)) => write!(f, "{}", b),
            Expr::Abs(param, body) => write!(f, "λ{}.{}", param, body),
            Expr::App(func, arg) => match (func.as_ref(), arg.as_ref()) {
                (Expr::Abs(_, _), _) => write!(f, "({}) {}", func, arg),
                (_, Expr::App(_, _)) => write!(f, "{} ({})", func, arg),
                (_, Expr::Abs(_, _)) => write!(f, "{} ({})", func, arg),
                _ => write!(f, "{} {}", func, arg),
            },
            Expr::Let(var, value, body) => {
                write!(f, "let {} = {} in {}", var, value, body)
            }
            Expr::Tuple(exprs) => {
                write!(f, "(")?;
                for (i, expr) in exprs.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", expr)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl std::fmt::Display for Type {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Type::Var(name) => write!(f, "{}", name),
            Type::Int => write!(f, "Int"),
            Type::Bool => write!(f, "Bool"),
            Type::Arrow(t1, t2) => {
                // Add parentheses around left side if it's an arrow to avoid ambiguity
                match t1.as_ref() {
                    Type::Arrow(_, _) => write!(f, "({}) → {}", t1, t2),
                    _ => write!(f, "{} → {}", t1, t2),
                }
            }
            Type::Tuple(types) => {
                write!(f, "(")?;
                for (i, ty) in types.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", ty)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl std::fmt::Display for Scheme {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.vars.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "forall {}. {}", self.vars.join(" "), self.ty)
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    Var(String),
    Arrow(Box<Type>, Box<Type>),
    Int,
    Bool,
    Tuple(Vec<Type>),
}
}

Type variables (Type::Var) serve as placeholders during inference, eventually getting instantiated to concrete types through unification. Arrow types (Type::Arrow) represent function types, encoding both parameter and return types. Base types like Int and Bool provide the foundation, while tuple types (Type::Tuple) support structured data. The recursive nature of these types allows us to express arbitrarily complex type structures, from simple integers to higher-order functions that manipulate other functions.

Type Inference Algorithm

Algorithm W operates on several fundamental data structures that capture the essential concepts of type inference. These type aliases provide names for the core abstractions and make the algorithm’s implementation more readable.

The type variable abstraction represents unknown types that will be resolved during inference. Term variables represent program variables that appear in expressions. The type environment maps term variables to their types, while substitutions map type variables to concrete types.

#![allow(unused)]
fn main() {
pub type TyVar = String;
pub type TmVar = String;
pub type Env = BTreeMap<TmVar, Scheme>;  // Now stores schemes, not types
pub type Subst = HashMap<TyVar, Type>;
}

These aliases encapsulate the fundamental data flow in Algorithm W. Type variables like t0, t1, and t2 serve as placeholders that get unified with concrete types as inference progresses. Term variables represent the actual identifiers in source programs. The environment now tracks polymorphic type schemes rather than just types, enabling proper let-polymorphism, while substitutions record the solutions discovered by unification.

The choice of String for both type and term variables reflects the simplicity of our implementation. In a full implementation, systems often use more complex representations like de Bruijn indices for type variables or interned strings for performance, but strings provide clarity for understanding the fundamental algorithms.

The heart of our Algorithm W implementation lies in the TypeInference struct, which maintains the state necessary for sound type inference across an entire program.

#![allow(unused)]
fn main() {
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fmt;
use crate::ast::{Expr, Lit, Scheme, Type};
use crate::errors::{InferenceError, Result};
pub type TyVar = String;
pub type TmVar = String;
pub type Env = BTreeMap<TmVar, Scheme>;
pub type Subst = HashMap<TyVar, Type>;
#[derive(Debug)]
pub struct InferenceTree {
    pub rule: String,
    pub input: String,
    pub output: String,
    pub children: Vec<InferenceTree>,
}
impl InferenceTree {
    fn new(rule: &str, input: &str, output: &str, children: Vec<InferenceTree>) -> Self {
        Self {
            rule: rule.to_string(),
            input: input.to_string(),
            output: output.to_string(),
            children,
        }
    }
}
impl fmt::Display for InferenceTree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.display_with_indent(f, 0)
    }
}
impl InferenceTree {
    fn display_with_indent(&self, f: &mut fmt::Formatter, indent: usize) -> fmt::Result {
        let prefix = "  ".repeat(indent);
        writeln!(
            f,
            "{}{}: {} => {}",
            prefix, self.rule, self.input, self.output
        )?;
        for child in &self.children {
            child.display_with_indent(f, indent + 1)?;
        }
        Ok(())
    }
}
impl Default for TypeInference {
    fn default() -> Self {
        Self::new()
    }
}
#[allow(clippy::only_used_in_recursion)]
impl TypeInference {
    pub fn new() -> Self {
        Self { counter: 0 }
    }

    fn fresh_tyvar(&mut self) -> TyVar {
        let var = format!("t{}", self.counter);
        self.counter += 1;
        var
    }

    fn pretty_env(&self, env: &Env) -> String {
        if env.is_empty() {
            "{}".to_string()
        } else {
            let entries: Vec<String> = env.iter().map(|(k, v)| format!("{}: {}", k, v)).collect();
            format!("{{{}}}", entries.join(", "))
        }
    }

    fn pretty_subst(&self, subst: &Subst) -> String {
        if subst.is_empty() {
            "{}".to_string()
        } else {
            let entries: Vec<String> = subst.iter().map(|(k, v)| format!("{}/{}", v, k)).collect();
            format!("{{{}}}", entries.join(", "))
        }
    }

    fn apply_subst(&self, subst: &Subst, ty: &Type) -> Type {
        match ty {
            Type::Var(name) => subst.get(name).cloned().unwrap_or_else(|| ty.clone()),
            Type::Arrow(t1, t2) => Type::Arrow(
                Box::new(self.apply_subst(subst, t1)),
                Box::new(self.apply_subst(subst, t2)),
            ),
            Type::Tuple(types) => {
                Type::Tuple(types.iter().map(|t| self.apply_subst(subst, t)).collect())
            }
            Type::Int | Type::Bool => ty.clone(),
        }
    }

    fn apply_subst_scheme(&self, subst: &Subst, scheme: &Scheme) -> Scheme {
        // Remove bindings for quantified variables to avoid capture
        let mut filtered_subst = subst.clone();
        for var in &scheme.vars {
            filtered_subst.remove(var);
        }
        Scheme {
            vars: scheme.vars.clone(),
            ty: self.apply_subst(&filtered_subst, &scheme.ty),
        }
    }

    fn apply_subst_env(&self, subst: &Subst, env: &Env) -> Env {
        env.iter()
            .map(|(k, v)| (k.clone(), self.apply_subst_scheme(subst, v)))
            .collect()
    }

    fn compose_subst(&self, s1: &Subst, s2: &Subst) -> Subst {
        let mut result = s1.clone();
        for (k, v) in s2 {
            result.insert(k.clone(), self.apply_subst(s1, v));
        }
        result
    }

    fn free_type_vars(&self, ty: &Type) -> HashSet<TyVar> {
        match ty {
            Type::Var(name) => {
                let mut set = HashSet::new();
                set.insert(name.clone());
                set
            }
            Type::Arrow(t1, t2) => {
                let mut set = self.free_type_vars(t1);
                set.extend(self.free_type_vars(t2));
                set
            }
            Type::Tuple(types) => {
                let mut set = HashSet::new();
                for t in types {
                    set.extend(self.free_type_vars(t));
                }
                set
            }
            Type::Int | Type::Bool => HashSet::new(),
        }
    }

    fn free_type_vars_scheme(&self, scheme: &Scheme) -> HashSet<TyVar> {
        let mut set = self.free_type_vars(&scheme.ty);
        // Remove quantified variables
        for var in &scheme.vars {
            set.remove(var);
        }
        set
    }

    fn free_type_vars_env(&self, env: &Env) -> HashSet<TyVar> {
        let mut set = HashSet::new();
        for scheme in env.values() {
            set.extend(self.free_type_vars_scheme(scheme));
        }
        set
    }

    fn generalize(&self, env: &Env, ty: &Type) -> Scheme {
        let type_vars = self.free_type_vars(ty);
        let env_vars = self.free_type_vars_env(env);
        let mut free_vars: Vec<_> = type_vars.difference(&env_vars).cloned().collect();
        free_vars.sort(); // Sort for deterministic behavior

        Scheme {
            vars: free_vars,
            ty: ty.clone(),
        }
    }

    fn instantiate(&mut self, scheme: &Scheme) -> Type {
        // Create fresh type variables for each quantified variable
        let mut subst = HashMap::new();
        for var in &scheme.vars {
            let fresh = self.fresh_tyvar();
            subst.insert(var.clone(), Type::Var(fresh));
        }

        self.apply_subst(&subst, &scheme.ty)
    }

    fn occurs_check(&self, var: &TyVar, ty: &Type) -> bool {
        match ty {
            Type::Var(name) => name == var,
            Type::Arrow(t1, t2) => self.occurs_check(var, t1) || self.occurs_check(var, t2),
            Type::Tuple(types) => types.iter().any(|t| self.occurs_check(var, t)),
            Type::Int | Type::Bool => false,
        }
    }

    fn unify(&self, t1: &Type, t2: &Type) -> Result<(Subst, InferenceTree)> {
        let input = format!("{} ~ {}", t1, t2);

        match (t1, t2) {
            (Type::Int, Type::Int) | (Type::Bool, Type::Bool) => {
                let tree = InferenceTree::new("Unify-Base", &input, "{}", vec![]);
                Ok((HashMap::new(), tree))
            }
            (Type::Var(v), ty) | (ty, Type::Var(v)) => {
                if ty == &Type::Var(v.clone()) {
                    let tree = InferenceTree::new("Unify-Var-Same", &input, "{}", vec![]);
                    Ok((HashMap::new(), tree))
                } else if self.occurs_check(v, ty) {
                    Err(InferenceError::OccursCheck {
                        var: v.clone(),
                        ty: ty.clone(),
                    })
                } else {
                    let mut subst = HashMap::new();
                    subst.insert(v.clone(), ty.clone());
                    let output = format!("{{{}/{}}}", ty, v);
                    let tree = InferenceTree::new("Unify-Var", &input, &output, vec![]);
                    Ok((subst, tree))
                }
            }
            (Type::Arrow(a1, a2), Type::Arrow(b1, b2)) => {
                let (s1, tree1) = self.unify(a1, b1)?;
                let a2_subst = self.apply_subst(&s1, a2);
                let b2_subst = self.apply_subst(&s1, b2);
                let (s2, tree2) = self.unify(&a2_subst, &b2_subst)?;
                let final_subst = self.compose_subst(&s2, &s1);
                let output = self.pretty_subst(&final_subst);
                let tree = InferenceTree::new("Unify-Arrow", &input, &output, vec![tree1, tree2]);
                Ok((final_subst, tree))
            }
            (Type::Tuple(ts1), Type::Tuple(ts2)) => {
                if ts1.len() != ts2.len() {
                    return Err(InferenceError::TupleLengthMismatch {
                        left_len: ts1.len(),
                        right_len: ts2.len(),
                    });
                }

                let mut subst = HashMap::new();
                let mut trees = Vec::new();

                for (t1, t2) in ts1.iter().zip(ts2.iter()) {
                    let t1_subst = self.apply_subst(&subst, t1);
                    let t2_subst = self.apply_subst(&subst, t2);
                    let (s, tree) = self.unify(&t1_subst, &t2_subst)?;
                    subst = self.compose_subst(&s, &subst);
                    trees.push(tree);
                }

                let output = self.pretty_subst(&subst);
                let tree = InferenceTree::new("Unify-Tuple", &input, &output, trees);
                Ok((subst, tree))
            }
            _ => Err(InferenceError::UnificationFailure {
                expected: t1.clone(),
                actual: t2.clone(),
            }),
        }
    }

    pub fn infer(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
        match expr {
            Expr::Lit(Lit::Int(_)) => self.infer_lit_int(env, expr),
            Expr::Lit(Lit::Bool(_)) => self.infer_lit_bool(env, expr),
            Expr::Var(name) => self.infer_var(env, expr, name),
            Expr::Abs(param, body) => self.infer_abs(env, expr, param, body),
            Expr::App(func, arg) => self.infer_app(env, expr, func, arg),
            Expr::Let(var, value, body) => self.infer_let(env, expr, var, value, body),
            Expr::Tuple(exprs) => self.infer_tuple(env, expr, exprs),
        }
    }

    /// T-LitInt: ─────────────────
    ///           Γ ⊢ n : Int
    fn infer_lit_int(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);
        let tree = InferenceTree::new("T-Int", &input, "Int", vec![]);
        Ok((HashMap::new(), Type::Int, tree))
    }

    /// T-LitBool: ─────────────────
    ///            Γ ⊢ b : Bool
    fn infer_lit_bool(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);
        let tree = InferenceTree::new("T-Bool", &input, "Bool", vec![]);
        Ok((HashMap::new(), Type::Bool, tree))
    }

    /// T-Var: x : σ ∈ Γ    τ = inst(σ)
    ///        ─────────────────────────
    ///               Γ ⊢ x : τ
    fn infer_var(
        &mut self,
        env: &Env,
        expr: &Expr,
        name: &str,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        match env.get(name) {
            Some(scheme) => {
                let instantiated = self.instantiate(scheme);
                let output = format!("{}", instantiated);
                let tree = InferenceTree::new("T-Var", &input, &output, vec![]);
                Ok((HashMap::new(), instantiated, tree))
            }
            None => Err(InferenceError::UnboundVariable {
                name: name.to_string(),
            }),
        }
    }

    /// T-Lam: Γ, x : α ⊢ e : τ    α fresh
    ///        ─────────────────────────────
    ///           Γ ⊢ λx. e : α → τ
    fn infer_abs(
        &mut self,
        env: &Env,
        expr: &Expr,
        param: &str,
        body: &Expr,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let param_type = Type::Var(self.fresh_tyvar());
        let mut new_env = env.clone();
        // Insert a monomorphic scheme for the parameter
        let param_scheme = Scheme {
            vars: vec![],
            ty: param_type.clone(),
        };
        new_env.insert(param.to_string(), param_scheme);

        let (s1, body_type, tree1) = self.infer(&new_env, body)?;
        let param_type_subst = self.apply_subst(&s1, &param_type);
        let result_type = Type::Arrow(Box::new(param_type_subst), Box::new(body_type));

        let output = format!("{}", result_type);
        let tree = InferenceTree::new("T-Abs", &input, &output, vec![tree1]);
        Ok((s1, result_type, tree))
    }

    /// T-App: Γ ⊢ e₁ : τ₁    Γ ⊢ e₂ : τ₂    α fresh    S = unify(τ₁, τ₂ → α)
    ///        ──────────────────────────────────────────────────────────────
    ///                            Γ ⊢ e₁ e₂ : S(α)
    fn infer_app(
        &mut self,
        env: &Env,
        expr: &Expr,
        func: &Expr,
        arg: &Expr,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let result_type = Type::Var(self.fresh_tyvar());

        let (s1, func_type, tree1) = self.infer(env, func)?;
        let env_subst = self.apply_subst_env(&s1, env);
        let (s2, arg_type, tree2) = self.infer(&env_subst, arg)?;

        let func_type_subst = self.apply_subst(&s2, &func_type);
        let expected_func_type = Type::Arrow(Box::new(arg_type), Box::new(result_type.clone()));

        let (s3, tree3) = self.unify(&func_type_subst, &expected_func_type)?;

        let final_subst = self.compose_subst(&s3, &self.compose_subst(&s2, &s1));
        let final_type = self.apply_subst(&s3, &result_type);

        let output = format!("{}", final_type);
        let tree = InferenceTree::new("T-App", &input, &output, vec![tree1, tree2, tree3]);
        Ok((final_subst, final_type, tree))
    }

    /// T-Let: Γ ⊢ e₁ : τ₁    σ = gen(Γ, τ₁)    Γ, x : σ ⊢ e₂ : τ₂
    ///        ──────────────────────────────────────────────────────
    ///                     Γ ⊢ let x = e₁ in e₂ : τ₂
    fn infer_let(
        &mut self,
        env: &Env,
        expr: &Expr,
        var: &str,
        value: &Expr,
        body: &Expr,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let (s1, value_type, tree1) = self.infer(env, value)?;
        let env_subst = self.apply_subst_env(&s1, env);
        let generalized_type = self.generalize(&env_subst, &value_type);

        let mut new_env = env_subst;
        new_env.insert(var.to_string(), generalized_type);

        let (s2, body_type, tree2) = self.infer(&new_env, body)?;

        let final_subst = self.compose_subst(&s2, &s1);
        let output = format!("{}", body_type);
        let tree = InferenceTree::new("T-Let", &input, &output, vec![tree1, tree2]);
        Ok((final_subst, body_type, tree))
    }

    /// T-Tuple: Γ ⊢ e₁ : τ₁    ...    Γ ⊢ eₙ : τₙ
    ///          ─────────────────────────────────────
    ///              Γ ⊢ (e₁, ..., eₙ) : (τ₁, ..., τₙ)
    fn infer_tuple(
        &mut self,
        env: &Env,
        expr: &Expr,
        exprs: &[Expr],
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let mut subst = HashMap::new();
        let mut types = Vec::new();
        let mut trees = Vec::new();
        let mut current_env = env.clone();

        for expr in exprs {
            let (s, ty, tree) = self.infer(&current_env, expr)?;
            subst = self.compose_subst(&s, &subst);
            current_env = self.apply_subst_env(&s, &current_env);
            types.push(ty);
            trees.push(tree);
        }

        let result_type = Type::Tuple(types);
        let output = format!("{}", result_type);
        let tree = InferenceTree::new("T-Tuple", &input, &output, trees);
        Ok((subst, result_type, tree))
    }
}
pub fn run_inference(expr: &Expr) -> Result<InferenceTree> {
    let mut inference = TypeInference::new();
    let env = BTreeMap::new();
    let (_, _, tree) = inference.infer(&env, expr)?;
    Ok(tree)
}
pub fn infer_type_only(expr: &Expr) -> Result<Type> {
    let mut inference = TypeInference::new();
    let env = BTreeMap::new();
    let (_, ty, _) = inference.infer(&env, expr)?;
    Ok(ty)
}
pub struct TypeInference {
    counter: usize,
}
}

The inference engine’s primary responsibility is generating fresh type variables, a process that ensures each unknown type gets a unique identifier. This counter-based approach provides a simple but effective way to avoid naming collisions during the inference process.

#![allow(unused)]
fn main() {
fn fresh_tyvar(&mut self) -> TyVar {
    let var = format!("t{}", self.counter);
    self.counter += 1;
    var
}
}

Fresh variable generation forms the foundation for Algorithm W’s systematic approach to handling unknowns. Each time we encounter an expression whose type we don’t yet know, we assign it a fresh type variable. These variables later get unified with concrete types as we discover more information about the program’s structure.

Substitution and Unification

Type substitutions represent the core computational mechanism of Algorithm W. A substitution maps type variables to concrete types, effectively “solving” part of our type inference puzzle.

The application of substitutions must handle the recursive structure of types correctly, ensuring that substitutions propagate through compound types like arrows and tuples.

#![allow(unused)]
fn main() {
fn apply_subst(&self, subst: &Subst, ty: &Type) -> Type {
    match ty {
        Type::Var(name) => subst.get(name).cloned().unwrap_or_else(|| ty.clone()),
        Type::Arrow(t1, t2) => Type::Arrow(
            Box::new(self.apply_subst(subst, t1)),
            Box::new(self.apply_subst(subst, t2)),
        ),
        Type::Tuple(types) => {
            Type::Tuple(types.iter().map(|t| self.apply_subst(subst, t)).collect())
        }
        Type::Int | Type::Bool => ty.clone(),
    }
}
}

Substitution application demonstrates how type information flows through our system. When we apply a substitution to an arrow type, we must apply it recursively to both the parameter and return types. This ensures that type information discovered in one part of a program correctly influences other parts.

Composition of substitutions allows us to combine multiple partial solutions into a more complete understanding of our program’s types.

#![allow(unused)]
fn main() {
fn compose_subst(&self, s1: &Subst, s2: &Subst) -> Subst {
    let mut result = s1.clone();
    for (k, v) in s2 {
        result.insert(k.clone(), self.apply_subst(s1, v));
    }
    result
}
}

The composition operation ensures that when we have multiple substitutions from different parts of our inference process, we can combine them into a single, consistent substitution that represents our cumulative knowledge about the program’s types.

Substitutions must also be applied to entire type environments when we discover new type information. This operation updates all the types in the environment according to the current substitution.

#![allow(unused)]
fn main() {
fn apply_subst_env(&self, subst: &Subst, env: &Env) -> Env {
    env.iter()
        .map(|(k, v)| (k.clone(), self.apply_subst_scheme(subst, v)))
        .collect()
}
}

Environment substitution is crucial for maintaining consistency as inference progresses. When we discover that a type variable should be instantiated to a concrete type, we must update not just individual types but entire environments to reflect this new knowledge.

Unification

Unification is the heart of type inference, solving constraints between types. The unification algorithm produces substitutions that make two types equivalent:

Reflexivity - identical types unify trivially: \[ \frac{}{\text{unify}(τ, τ) = \emptyset} \text{(U-Refl)} \]

Variable unification with occurs check: \[ \frac{α \notin \text{ftv}(τ)}{\text{unify}(α, τ) = [α ↦ τ]} \text{(U-VarL)} \]

\[ \frac{α \notin \text{ftv}(τ)}{\text{unify}(τ, α) = [α ↦ τ]} \text{(U-VarR)} \]

Arrow type unification decomposes into domain and codomain: \[ \frac{S₁ = \text{unify}(τ₁, τ₃) \quad S₂ = \text{unify}(S₁(τ₂), S₁(τ₄))}{\text{unify}(τ₁ → τ₂, τ₃ → τ₄) = S₂ ∘ S₁} \text{(U-Arrow)} \]

Tuple unification requires component-wise unification: \[ \frac{S₁ = \text{unify}(τ₁, τ₃) \quad S₂ = \text{unify}(S₁(τ₂), S₁(τ₄))}{\text{unify}((τ₁, τ₂), (τ₃, τ₄)) = S₂ ∘ S₁} \text{(U-Tuple)} \]

Base type unification succeeds only for identical types: \[ \frac{}{\text{unify}(\text{Int}, \text{Int}) = \emptyset} \text{(U-Int)} \]

\[ \frac{}{\text{unify}(\text{Bool}, \text{Bool}) = \emptyset} \text{(U-Bool)} \]

These unification rules ensure that type constraints are solved systematically while maintaining soundness through the occurs check.

When we have two types that must be equal (such as the parameter type of a function and the type of an argument being passed to it), unification determines whether this constraint can be satisfied and, if so, what substitution makes them equal.

#![allow(unused)]
fn main() {
fn unify(&self, t1: &Type, t2: &Type) -> Result<(Subst, InferenceTree)> {
    let input = format!("{} ~ {}", t1, t2);

    match (t1, t2) {
        (Type::Int, Type::Int) | (Type::Bool, Type::Bool) => {
            let tree = InferenceTree::new("Unify-Base", &input, "{}", vec![]);
            Ok((HashMap::new(), tree))
        }
        (Type::Var(v), ty) | (ty, Type::Var(v)) => {
            if ty == &Type::Var(v.clone()) {
                let tree = InferenceTree::new("Unify-Var-Same", &input, "{}", vec![]);
                Ok((HashMap::new(), tree))
            } else if self.occurs_check(v, ty) {
                Err(InferenceError::OccursCheck {
                    var: v.clone(),
                    ty: ty.clone(),
                })
            } else {
                let mut subst = HashMap::new();
                subst.insert(v.clone(), ty.clone());
                let output = format!("{{{}/{}}}", ty, v);
                let tree = InferenceTree::new("Unify-Var", &input, &output, vec![]);
                Ok((subst, tree))
            }
        }
        (Type::Arrow(a1, a2), Type::Arrow(b1, b2)) => {
            let (s1, tree1) = self.unify(a1, b1)?;
            let a2_subst = self.apply_subst(&s1, a2);
            let b2_subst = self.apply_subst(&s1, b2);
            let (s2, tree2) = self.unify(&a2_subst, &b2_subst)?;
            let final_subst = self.compose_subst(&s2, &s1);
            let output = self.pretty_subst(&final_subst);
            let tree = InferenceTree::new("Unify-Arrow", &input, &output, vec![tree1, tree2]);
            Ok((final_subst, tree))
        }
        (Type::Tuple(ts1), Type::Tuple(ts2)) => {
            if ts1.len() != ts2.len() {
                return Err(InferenceError::TupleLengthMismatch {
                    left_len: ts1.len(),
                    right_len: ts2.len(),
                });
            }

            let mut subst = HashMap::new();
            let mut trees = Vec::new();

            for (t1, t2) in ts1.iter().zip(ts2.iter()) {
                let t1_subst = self.apply_subst(&subst, t1);
                let t2_subst = self.apply_subst(&subst, t2);
                let (s, tree) = self.unify(&t1_subst, &t2_subst)?;
                subst = self.compose_subst(&s, &subst);
                trees.push(tree);
            }

            let output = self.pretty_subst(&subst);
            let tree = InferenceTree::new("Unify-Tuple", &input, &output, trees);
            Ok((subst, tree))
        }
        _ => Err(InferenceError::UnificationFailure {
            expected: t1.clone(),
            actual: t2.clone(),
        }),
    }
}
}

The unification algorithm handles several distinct cases, each representing a different constraint-solving scenario. When unifying two identical base types like Int with Int, no substitution is needed. When unifying a type variable with any other type, we create a substitution that maps the variable to that type, provided the occurs check passes.

The occurs check prevents infinite types by ensuring that a type variable doesn’t appear within the type it’s being unified with.

#![allow(unused)]
fn main() {
fn occurs_check(&self, var: &TyVar, ty: &Type) -> bool {
    match ty {
        Type::Var(name) => name == var,
        Type::Arrow(t1, t2) => self.occurs_check(var, t1) || self.occurs_check(var, t2),
        Type::Tuple(types) => types.iter().any(|t| self.occurs_check(var, t)),
        Type::Int | Type::Bool => false,
    }
}
}

This check is essential for soundness. Without it, we might generate infinite types like t0 = t0 -> Int, which would break our type system’s decidability.

For compound types like arrows, unification becomes recursive. We must unify corresponding subcomponents and then compose the resulting substitutions. This process ensures that complex types maintain their structural relationships while allowing for flexible instantiation of type variables.

The Main Inference Algorithm

The central infer method implements Algorithm W proper, analyzing expressions to determine their types while accumulating the necessary substitutions. Our implementation uses a modular approach where each syntactic construct has its own helper method implementing the corresponding typing rule.

#![allow(unused)]
fn main() {
pub fn infer(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
    match expr {
        Expr::Lit(Lit::Int(_)) => self.infer_lit_int(env, expr),
        Expr::Lit(Lit::Bool(_)) => self.infer_lit_bool(env, expr),
        Expr::Var(name) => self.infer_var(env, expr, name),
        Expr::Abs(param, body) => self.infer_abs(env, expr, param, body),
        Expr::App(func, arg) => self.infer_app(env, expr, func, arg),
        Expr::Let(var, value, body) => self.infer_let(env, expr, var, value, body),
        Expr::Tuple(exprs) => self.infer_tuple(env, expr, exprs),
    }
}
}

Each helper method corresponds directly to a formal typing rule, making the relationship between theory and implementation explicit.

Variable Lookup

\[ \frac{x : σ \in Γ \quad τ = \text{inst}(σ)}{Γ ⊢ x : τ} \text{(T-Var)} \]

#![allow(unused)]
fn main() {
/// T-Var: x : σ ∈ Γ    τ = inst(σ)
///        ─────────────────────────
///               Γ ⊢ x : τ
fn infer_var(
    &mut self,
    env: &Env,
    expr: &Expr,
    name: &str,
) -> Result<(Subst, Type, InferenceTree)> {
    let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

    match env.get(name) {
        Some(scheme) => {
            let instantiated = self.instantiate(scheme);
            let output = format!("{}", instantiated);
            let tree = InferenceTree::new("T-Var", &input, &output, vec![]);
            Ok((HashMap::new(), instantiated, tree))
        }
        None => Err(InferenceError::UnboundVariable {
            name: name.to_string(),
        }),
    }
}
}

Variable lookup requires instantiation of polymorphic types. When we find a variable in the environment, it might have a polymorphic type scheme like ∀α. α → α. We create a fresh monomorphic instance by replacing quantified variables with fresh type variables.

Lambda Abstraction

\[ \frac{Γ, x : α ⊢ e : τ \quad α \text{ fresh}}{Γ ⊢ λx. e : α → τ} \text{(T-Lam)} \]

#![allow(unused)]
fn main() {
/// T-Lam: Γ, x : α ⊢ e : τ    α fresh
///        ─────────────────────────────
///           Γ ⊢ λx. e : α → τ
fn infer_abs(
    &mut self,
    env: &Env,
    expr: &Expr,
    param: &str,
    body: &Expr,
) -> Result<(Subst, Type, InferenceTree)> {
    let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

    let param_type = Type::Var(self.fresh_tyvar());
    let mut new_env = env.clone();
    // Insert a monomorphic scheme for the parameter
    let param_scheme = Scheme {
        vars: vec![],
        ty: param_type.clone(),
    };
    new_env.insert(param.to_string(), param_scheme);

    let (s1, body_type, tree1) = self.infer(&new_env, body)?;
    let param_type_subst = self.apply_subst(&s1, &param_type);
    let result_type = Type::Arrow(Box::new(param_type_subst), Box::new(body_type));

    let output = format!("{}", result_type);
    let tree = InferenceTree::new("T-Abs", &input, &output, vec![tree1]);
    Ok((s1, result_type, tree))
}
}

Lambda abstractions introduce new variable bindings. We assign a fresh type variable to the parameter, extend the environment, and infer the body’s type. Any constraints discovered during body inference get propagated back through substitution.

Function Application

\[ \frac{Γ ⊢ e₁ : τ₁ \quad Γ ⊢ e₂ : τ₂ \quad α \text{ fresh} \quad S = \text{unify}(τ₁, τ₂ → α)}{Γ ⊢ e₁ \, e₂ : S(α)} \text{(T-App)} \]

#![allow(unused)]
fn main() {
/// T-App: Γ ⊢ e₁ : τ₁    Γ ⊢ e₂ : τ₂    α fresh    S = unify(τ₁, τ₂ → α)
///        ──────────────────────────────────────────────────────────────
///                            Γ ⊢ e₁ e₂ : S(α)
fn infer_app(
    &mut self,
    env: &Env,
    expr: &Expr,
    func: &Expr,
    arg: &Expr,
) -> Result<(Subst, Type, InferenceTree)> {
    let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

    let result_type = Type::Var(self.fresh_tyvar());

    let (s1, func_type, tree1) = self.infer(env, func)?;
    let env_subst = self.apply_subst_env(&s1, env);
    let (s2, arg_type, tree2) = self.infer(&env_subst, arg)?;

    let func_type_subst = self.apply_subst(&s2, &func_type);
    let expected_func_type = Type::Arrow(Box::new(arg_type), Box::new(result_type.clone()));

    let (s3, tree3) = self.unify(&func_type_subst, &expected_func_type)?;

    let final_subst = self.compose_subst(&s3, &self.compose_subst(&s2, &s1));
    let final_type = self.apply_subst(&s3, &result_type);

    let output = format!("{}", final_type);
    let tree = InferenceTree::new("T-App", &input, &output, vec![tree1, tree2, tree3]);
    Ok((final_subst, final_type, tree))
}
}

Application drives constraint generation. We infer types for both function and argument, then unify the function type with an arrow type constructed from the argument type and a fresh result type variable.

Let-Polymorphism

\[ \frac{Γ ⊢ e₁ : τ₁ \quad σ = \text{gen}(Γ, τ₁) \quad Γ, x : σ ⊢ e₂ : τ₂}{Γ ⊢ \text{let } x = e₁ \text{ in } e₂ : τ₂} \text{(T-Let)} \]

#![allow(unused)]
fn main() {
/// T-Let: Γ ⊢ e₁ : τ₁    σ = gen(Γ, τ₁)    Γ, x : σ ⊢ e₂ : τ₂
///        ──────────────────────────────────────────────────────
///                     Γ ⊢ let x = e₁ in e₂ : τ₂
fn infer_let(
    &mut self,
    env: &Env,
    expr: &Expr,
    var: &str,
    value: &Expr,
    body: &Expr,
) -> Result<(Subst, Type, InferenceTree)> {
    let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

    let (s1, value_type, tree1) = self.infer(env, value)?;
    let env_subst = self.apply_subst_env(&s1, env);
    let generalized_type = self.generalize(&env_subst, &value_type);

    let mut new_env = env_subst;
    new_env.insert(var.to_string(), generalized_type);

    let (s2, body_type, tree2) = self.infer(&new_env, body)?;

    let final_subst = self.compose_subst(&s2, &s1);
    let output = format!("{}", body_type);
    let tree = InferenceTree::new("T-Let", &input, &output, vec![tree1, tree2]);
    Ok((final_subst, body_type, tree))
}
}

Let expressions enable polymorphism through generalization. After inferring the bound expression’s type, we generalize it by quantifying over type variables not constrained by the environment. This allows polymorphic usage in the let body.

Literal Types

\[ \frac{}{Γ ⊢ n : \text{Int}} \text{(T-LitInt)} \]

\[ \frac{}{Γ ⊢ b : \text{Bool}} \text{(T-LitBool)} \]

#![allow(unused)]
fn main() {
/// T-LitInt: ─────────────────
///           Γ ⊢ n : Int
fn infer_lit_int(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
    let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);
    let tree = InferenceTree::new("T-Int", &input, "Int", vec![]);
    Ok((HashMap::new(), Type::Int, tree))
}
}

#![allow(unused)]
fn main() {
/// T-LitBool: ─────────────────
///            Γ ⊢ b : Bool
fn infer_lit_bool(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
    let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);
    let tree = InferenceTree::new("T-Bool", &input, "Bool", vec![]);
    Ok((HashMap::new(), Type::Bool, tree))
}
}

Literals have known types and require no constraint generation.

Generalization and Instantiation

The generalization and instantiation mechanisms handle let-polymorphism, allowing variables bound in let expressions to be used with multiple different types.

Understanding Generalization with Examples

Generalization turns concrete types into polymorphic type schemes. Consider this simple example:

let id = \x -> x in (id 42, id true)

When we infer the type of \x -> x, we get something like t0 → t0 where t0 is a type variable. Since t0 doesn’t appear anywhere else in the environment, we can generalize it to ∀t0. t0 → t0, making id polymorphic.

This is what allows id to be used with both 42 (type Int) and true (type Bool) in the same expression. Without generalization, the first use would fix t0 to Int, making the second use fail.

#![allow(unused)]
fn main() {
fn generalize(&self, env: &Env, ty: &Type) -> Scheme {
    let type_vars = self.free_type_vars(ty);
    let env_vars = self.free_type_vars_env(env);
    let mut free_vars: Vec<_> = type_vars.difference(&env_vars).cloned().collect();
    free_vars.sort(); // Sort for deterministic behavior

    Scheme {
        vars: free_vars,
        ty: ty.clone(),
    }
}
}

Generalization identifies type variables that could be made polymorphic by checking which ones don’t appear free in the current environment. If a type variable isn’t constrained by anything else in scope, it’s safe to quantify over it.

Understanding Instantiation with Examples

Instantiation creates fresh monomorphic versions from polymorphic types. When we use a polymorphic function like our identity function, we need to create a fresh copy of its type for each use.

Consider this expression:

let id = \x -> x in id id

Here we’re applying the polymorphic identity function to itself. The first id gets instantiated to (α → α) → (α → α) while the second id gets instantiated to α → α. These different instantiations allow the application to type-check successfully.

#![allow(unused)]
fn main() {
fn instantiate(&mut self, scheme: &Scheme) -> Type {
    // Create fresh type variables for each quantified variable
    let mut subst = HashMap::new();
    for var in &scheme.vars {
        let fresh = self.fresh_tyvar();
        subst.insert(var.clone(), Type::Var(fresh));
    }

    self.apply_subst(&subst, &scheme.ty)
}
}

Instantiation replaces quantified type variables with fresh type variables. This ensures that each use of a polymorphic function gets its own independent type constraints, preventing interference between different call sites.

Free Type Variables

Generalization depends on computing the free type variables in both individual types and entire environments. These operations identify which type variables could potentially be generalized versus those that are already constrained.

#![allow(unused)]
fn main() {
fn free_type_vars(&self, ty: &Type) -> HashSet<TyVar> {
    match ty {
        Type::Var(name) => {
            let mut set = HashSet::new();
            set.insert(name.clone());
            set
        }
        Type::Arrow(t1, t2) => {
            let mut set = self.free_type_vars(t1);
            set.extend(self.free_type_vars(t2));
            set
        }
        Type::Tuple(types) => {
            let mut set = HashSet::new();
            for t in types {
                set.extend(self.free_type_vars(t));
            }
            set
        }
        Type::Int | Type::Bool => HashSet::new(),
    }
}
}

The free type variables computation traverses type structures recursively, collecting all type variables that appear unbound. For compound types like arrows and tuples, it must traverse all subcomponents to ensure no variables are missed.

#![allow(unused)]
fn main() {
fn free_type_vars_env(&self, env: &Env) -> HashSet<TyVar> {
    let mut set = HashSet::new();
    for scheme in env.values() {
        set.extend(self.free_type_vars_scheme(scheme));
    }
    set
}
}

Computing free variables across entire environments requires examining every type in the environment and taking the union of their free variables. This gives us the complete set of type variables that are constrained by the current context.

Our complete implementation fully supports polymorphic instantiation by generating fresh type variables for each quantified variable in a scheme when it is instantiated. This mechanism is what allows the identity function to work on integers in one context and booleans in another, as demonstrated by expressions like let id = \x -> x in (id, id) which produces the type (t1 -> t1, t2 -> t2) showing proper polymorphic instantiation.

Error Handling and Inference Trees

Our implementation provides detailed error reporting and generates inference trees that show the step-by-step reasoning process.

#![allow(unused)]
fn main() {
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fmt;
use crate::ast::{Expr, Lit, Scheme, Type};
use crate::errors::{InferenceError, Result};
pub type TyVar = String;
pub type TmVar = String;
pub type Env = BTreeMap<TmVar, Scheme>;
pub type Subst = HashMap<TyVar, Type>;
impl InferenceTree {
    fn new(rule: &str, input: &str, output: &str, children: Vec<InferenceTree>) -> Self {
        Self {
            rule: rule.to_string(),
            input: input.to_string(),
            output: output.to_string(),
            children,
        }
    }
}
impl fmt::Display for InferenceTree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.display_with_indent(f, 0)
    }
}
impl InferenceTree {
    fn display_with_indent(&self, f: &mut fmt::Formatter, indent: usize) -> fmt::Result {
        let prefix = "  ".repeat(indent);
        writeln!(
            f,
            "{}{}: {} => {}",
            prefix, self.rule, self.input, self.output
        )?;
        for child in &self.children {
            child.display_with_indent(f, indent + 1)?;
        }
        Ok(())
    }
}
pub struct TypeInference {
    counter: usize,
}
impl Default for TypeInference {
    fn default() -> Self {
        Self::new()
    }
}
#[allow(clippy::only_used_in_recursion)]
impl TypeInference {
    pub fn new() -> Self {
        Self { counter: 0 }
    }

    fn fresh_tyvar(&mut self) -> TyVar {
        let var = format!("t{}", self.counter);
        self.counter += 1;
        var
    }

    fn pretty_env(&self, env: &Env) -> String {
        if env.is_empty() {
            "{}".to_string()
        } else {
            let entries: Vec<String> = env.iter().map(|(k, v)| format!("{}: {}", k, v)).collect();
            format!("{{{}}}", entries.join(", "))
        }
    }

    fn pretty_subst(&self, subst: &Subst) -> String {
        if subst.is_empty() {
            "{}".to_string()
        } else {
            let entries: Vec<String> = subst.iter().map(|(k, v)| format!("{}/{}", v, k)).collect();
            format!("{{{}}}", entries.join(", "))
        }
    }

    fn apply_subst(&self, subst: &Subst, ty: &Type) -> Type {
        match ty {
            Type::Var(name) => subst.get(name).cloned().unwrap_or_else(|| ty.clone()),
            Type::Arrow(t1, t2) => Type::Arrow(
                Box::new(self.apply_subst(subst, t1)),
                Box::new(self.apply_subst(subst, t2)),
            ),
            Type::Tuple(types) => {
                Type::Tuple(types.iter().map(|t| self.apply_subst(subst, t)).collect())
            }
            Type::Int | Type::Bool => ty.clone(),
        }
    }

    fn apply_subst_scheme(&self, subst: &Subst, scheme: &Scheme) -> Scheme {
        // Remove bindings for quantified variables to avoid capture
        let mut filtered_subst = subst.clone();
        for var in &scheme.vars {
            filtered_subst.remove(var);
        }
        Scheme {
            vars: scheme.vars.clone(),
            ty: self.apply_subst(&filtered_subst, &scheme.ty),
        }
    }

    fn apply_subst_env(&self, subst: &Subst, env: &Env) -> Env {
        env.iter()
            .map(|(k, v)| (k.clone(), self.apply_subst_scheme(subst, v)))
            .collect()
    }

    fn compose_subst(&self, s1: &Subst, s2: &Subst) -> Subst {
        let mut result = s1.clone();
        for (k, v) in s2 {
            result.insert(k.clone(), self.apply_subst(s1, v));
        }
        result
    }

    fn free_type_vars(&self, ty: &Type) -> HashSet<TyVar> {
        match ty {
            Type::Var(name) => {
                let mut set = HashSet::new();
                set.insert(name.clone());
                set
            }
            Type::Arrow(t1, t2) => {
                let mut set = self.free_type_vars(t1);
                set.extend(self.free_type_vars(t2));
                set
            }
            Type::Tuple(types) => {
                let mut set = HashSet::new();
                for t in types {
                    set.extend(self.free_type_vars(t));
                }
                set
            }
            Type::Int | Type::Bool => HashSet::new(),
        }
    }

    fn free_type_vars_scheme(&self, scheme: &Scheme) -> HashSet<TyVar> {
        let mut set = self.free_type_vars(&scheme.ty);
        // Remove quantified variables
        for var in &scheme.vars {
            set.remove(var);
        }
        set
    }

    fn free_type_vars_env(&self, env: &Env) -> HashSet<TyVar> {
        let mut set = HashSet::new();
        for scheme in env.values() {
            set.extend(self.free_type_vars_scheme(scheme));
        }
        set
    }

    fn generalize(&self, env: &Env, ty: &Type) -> Scheme {
        let type_vars = self.free_type_vars(ty);
        let env_vars = self.free_type_vars_env(env);
        let mut free_vars: Vec<_> = type_vars.difference(&env_vars).cloned().collect();
        free_vars.sort(); // Sort for deterministic behavior

        Scheme {
            vars: free_vars,
            ty: ty.clone(),
        }
    }

    fn instantiate(&mut self, scheme: &Scheme) -> Type {
        // Create fresh type variables for each quantified variable
        let mut subst = HashMap::new();
        for var in &scheme.vars {
            let fresh = self.fresh_tyvar();
            subst.insert(var.clone(), Type::Var(fresh));
        }

        self.apply_subst(&subst, &scheme.ty)
    }

    fn occurs_check(&self, var: &TyVar, ty: &Type) -> bool {
        match ty {
            Type::Var(name) => name == var,
            Type::Arrow(t1, t2) => self.occurs_check(var, t1) || self.occurs_check(var, t2),
            Type::Tuple(types) => types.iter().any(|t| self.occurs_check(var, t)),
            Type::Int | Type::Bool => false,
        }
    }

    fn unify(&self, t1: &Type, t2: &Type) -> Result<(Subst, InferenceTree)> {
        let input = format!("{} ~ {}", t1, t2);

        match (t1, t2) {
            (Type::Int, Type::Int) | (Type::Bool, Type::Bool) => {
                let tree = InferenceTree::new("Unify-Base", &input, "{}", vec![]);
                Ok((HashMap::new(), tree))
            }
            (Type::Var(v), ty) | (ty, Type::Var(v)) => {
                if ty == &Type::Var(v.clone()) {
                    let tree = InferenceTree::new("Unify-Var-Same", &input, "{}", vec![]);
                    Ok((HashMap::new(), tree))
                } else if self.occurs_check(v, ty) {
                    Err(InferenceError::OccursCheck {
                        var: v.clone(),
                        ty: ty.clone(),
                    })
                } else {
                    let mut subst = HashMap::new();
                    subst.insert(v.clone(), ty.clone());
                    let output = format!("{{{}/{}}}", ty, v);
                    let tree = InferenceTree::new("Unify-Var", &input, &output, vec![]);
                    Ok((subst, tree))
                }
            }
            (Type::Arrow(a1, a2), Type::Arrow(b1, b2)) => {
                let (s1, tree1) = self.unify(a1, b1)?;
                let a2_subst = self.apply_subst(&s1, a2);
                let b2_subst = self.apply_subst(&s1, b2);
                let (s2, tree2) = self.unify(&a2_subst, &b2_subst)?;
                let final_subst = self.compose_subst(&s2, &s1);
                let output = self.pretty_subst(&final_subst);
                let tree = InferenceTree::new("Unify-Arrow", &input, &output, vec![tree1, tree2]);
                Ok((final_subst, tree))
            }
            (Type::Tuple(ts1), Type::Tuple(ts2)) => {
                if ts1.len() != ts2.len() {
                    return Err(InferenceError::TupleLengthMismatch {
                        left_len: ts1.len(),
                        right_len: ts2.len(),
                    });
                }

                let mut subst = HashMap::new();
                let mut trees = Vec::new();

                for (t1, t2) in ts1.iter().zip(ts2.iter()) {
                    let t1_subst = self.apply_subst(&subst, t1);
                    let t2_subst = self.apply_subst(&subst, t2);
                    let (s, tree) = self.unify(&t1_subst, &t2_subst)?;
                    subst = self.compose_subst(&s, &subst);
                    trees.push(tree);
                }

                let output = self.pretty_subst(&subst);
                let tree = InferenceTree::new("Unify-Tuple", &input, &output, trees);
                Ok((subst, tree))
            }
            _ => Err(InferenceError::UnificationFailure {
                expected: t1.clone(),
                actual: t2.clone(),
            }),
        }
    }

    pub fn infer(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
        match expr {
            Expr::Lit(Lit::Int(_)) => self.infer_lit_int(env, expr),
            Expr::Lit(Lit::Bool(_)) => self.infer_lit_bool(env, expr),
            Expr::Var(name) => self.infer_var(env, expr, name),
            Expr::Abs(param, body) => self.infer_abs(env, expr, param, body),
            Expr::App(func, arg) => self.infer_app(env, expr, func, arg),
            Expr::Let(var, value, body) => self.infer_let(env, expr, var, value, body),
            Expr::Tuple(exprs) => self.infer_tuple(env, expr, exprs),
        }
    }

    /// T-LitInt: ─────────────────
    ///           Γ ⊢ n : Int
    fn infer_lit_int(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);
        let tree = InferenceTree::new("T-Int", &input, "Int", vec![]);
        Ok((HashMap::new(), Type::Int, tree))
    }

    /// T-LitBool: ─────────────────
    ///            Γ ⊢ b : Bool
    fn infer_lit_bool(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);
        let tree = InferenceTree::new("T-Bool", &input, "Bool", vec![]);
        Ok((HashMap::new(), Type::Bool, tree))
    }

    /// T-Var: x : σ ∈ Γ    τ = inst(σ)
    ///        ─────────────────────────
    ///               Γ ⊢ x : τ
    fn infer_var(
        &mut self,
        env: &Env,
        expr: &Expr,
        name: &str,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        match env.get(name) {
            Some(scheme) => {
                let instantiated = self.instantiate(scheme);
                let output = format!("{}", instantiated);
                let tree = InferenceTree::new("T-Var", &input, &output, vec![]);
                Ok((HashMap::new(), instantiated, tree))
            }
            None => Err(InferenceError::UnboundVariable {
                name: name.to_string(),
            }),
        }
    }

    /// T-Lam: Γ, x : α ⊢ e : τ    α fresh
    ///        ─────────────────────────────
    ///           Γ ⊢ λx. e : α → τ
    fn infer_abs(
        &mut self,
        env: &Env,
        expr: &Expr,
        param: &str,
        body: &Expr,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let param_type = Type::Var(self.fresh_tyvar());
        let mut new_env = env.clone();
        // Insert a monomorphic scheme for the parameter
        let param_scheme = Scheme {
            vars: vec![],
            ty: param_type.clone(),
        };
        new_env.insert(param.to_string(), param_scheme);

        let (s1, body_type, tree1) = self.infer(&new_env, body)?;
        let param_type_subst = self.apply_subst(&s1, &param_type);
        let result_type = Type::Arrow(Box::new(param_type_subst), Box::new(body_type));

        let output = format!("{}", result_type);
        let tree = InferenceTree::new("T-Abs", &input, &output, vec![tree1]);
        Ok((s1, result_type, tree))
    }

    /// T-App: Γ ⊢ e₁ : τ₁    Γ ⊢ e₂ : τ₂    α fresh    S = unify(τ₁, τ₂ → α)
    ///        ──────────────────────────────────────────────────────────────
    ///                            Γ ⊢ e₁ e₂ : S(α)
    fn infer_app(
        &mut self,
        env: &Env,
        expr: &Expr,
        func: &Expr,
        arg: &Expr,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let result_type = Type::Var(self.fresh_tyvar());

        let (s1, func_type, tree1) = self.infer(env, func)?;
        let env_subst = self.apply_subst_env(&s1, env);
        let (s2, arg_type, tree2) = self.infer(&env_subst, arg)?;

        let func_type_subst = self.apply_subst(&s2, &func_type);
        let expected_func_type = Type::Arrow(Box::new(arg_type), Box::new(result_type.clone()));

        let (s3, tree3) = self.unify(&func_type_subst, &expected_func_type)?;

        let final_subst = self.compose_subst(&s3, &self.compose_subst(&s2, &s1));
        let final_type = self.apply_subst(&s3, &result_type);

        let output = format!("{}", final_type);
        let tree = InferenceTree::new("T-App", &input, &output, vec![tree1, tree2, tree3]);
        Ok((final_subst, final_type, tree))
    }

    /// T-Let: Γ ⊢ e₁ : τ₁    σ = gen(Γ, τ₁)    Γ, x : σ ⊢ e₂ : τ₂
    ///        ──────────────────────────────────────────────────────
    ///                     Γ ⊢ let x = e₁ in e₂ : τ₂
    fn infer_let(
        &mut self,
        env: &Env,
        expr: &Expr,
        var: &str,
        value: &Expr,
        body: &Expr,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let (s1, value_type, tree1) = self.infer(env, value)?;
        let env_subst = self.apply_subst_env(&s1, env);
        let generalized_type = self.generalize(&env_subst, &value_type);

        let mut new_env = env_subst;
        new_env.insert(var.to_string(), generalized_type);

        let (s2, body_type, tree2) = self.infer(&new_env, body)?;

        let final_subst = self.compose_subst(&s2, &s1);
        let output = format!("{}", body_type);
        let tree = InferenceTree::new("T-Let", &input, &output, vec![tree1, tree2]);
        Ok((final_subst, body_type, tree))
    }

    /// T-Tuple: Γ ⊢ e₁ : τ₁    ...    Γ ⊢ eₙ : τₙ
    ///          ─────────────────────────────────────
    ///              Γ ⊢ (e₁, ..., eₙ) : (τ₁, ..., τₙ)
    fn infer_tuple(
        &mut self,
        env: &Env,
        expr: &Expr,
        exprs: &[Expr],
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let mut subst = HashMap::new();
        let mut types = Vec::new();
        let mut trees = Vec::new();
        let mut current_env = env.clone();

        for expr in exprs {
            let (s, ty, tree) = self.infer(&current_env, expr)?;
            subst = self.compose_subst(&s, &subst);
            current_env = self.apply_subst_env(&s, &current_env);
            types.push(ty);
            trees.push(tree);
        }

        let result_type = Type::Tuple(types);
        let output = format!("{}", result_type);
        let tree = InferenceTree::new("T-Tuple", &input, &output, trees);
        Ok((subst, result_type, tree))
    }
}
pub fn run_inference(expr: &Expr) -> Result<InferenceTree> {
    let mut inference = TypeInference::new();
    let env = BTreeMap::new();
    let (_, _, tree) = inference.infer(&env, expr)?;
    Ok(tree)
}
pub fn infer_type_only(expr: &Expr) -> Result<Type> {
    let mut inference = TypeInference::new();
    let env = BTreeMap::new();
    let (_, ty, _) = inference.infer(&env, expr)?;
    Ok(ty)
}
#[derive(Debug)]
pub struct InferenceTree {
    pub rule: String,
    pub input: String,
    pub output: String,
    pub children: Vec<InferenceTree>,
}
}

These inference trees serve both debugging and educational purposes. They make explicit the implicit reasoning that Algorithm W performs, showing how type information flows through the program and how constraints get generated and solved.

The public interface provides both tree-generating and type-only versions of inference, supporting different use cases from interactive development to automated tooling.

#![allow(unused)]
fn main() {
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fmt;
use crate::ast::{Expr, Lit, Scheme, Type};
use crate::errors::{InferenceError, Result};
pub type TyVar = String;
pub type TmVar = String;
pub type Env = BTreeMap<TmVar, Scheme>;
pub type Subst = HashMap<TyVar, Type>;
#[derive(Debug)]
pub struct InferenceTree {
    pub rule: String,
    pub input: String,
    pub output: String,
    pub children: Vec<InferenceTree>,
}
impl InferenceTree {
    fn new(rule: &str, input: &str, output: &str, children: Vec<InferenceTree>) -> Self {
        Self {
            rule: rule.to_string(),
            input: input.to_string(),
            output: output.to_string(),
            children,
        }
    }
}
impl fmt::Display for InferenceTree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.display_with_indent(f, 0)
    }
}
impl InferenceTree {
    fn display_with_indent(&self, f: &mut fmt::Formatter, indent: usize) -> fmt::Result {
        let prefix = "  ".repeat(indent);
        writeln!(
            f,
            "{}{}: {} => {}",
            prefix, self.rule, self.input, self.output
        )?;
        for child in &self.children {
            child.display_with_indent(f, indent + 1)?;
        }
        Ok(())
    }
}
pub struct TypeInference {
    counter: usize,
}
impl Default for TypeInference {
    fn default() -> Self {
        Self::new()
    }
}
#[allow(clippy::only_used_in_recursion)]
impl TypeInference {
    pub fn new() -> Self {
        Self { counter: 0 }
    }

    fn fresh_tyvar(&mut self) -> TyVar {
        let var = format!("t{}", self.counter);
        self.counter += 1;
        var
    }

    fn pretty_env(&self, env: &Env) -> String {
        if env.is_empty() {
            "{}".to_string()
        } else {
            let entries: Vec<String> = env.iter().map(|(k, v)| format!("{}: {}", k, v)).collect();
            format!("{{{}}}", entries.join(", "))
        }
    }

    fn pretty_subst(&self, subst: &Subst) -> String {
        if subst.is_empty() {
            "{}".to_string()
        } else {
            let entries: Vec<String> = subst.iter().map(|(k, v)| format!("{}/{}", v, k)).collect();
            format!("{{{}}}", entries.join(", "))
        }
    }

    fn apply_subst(&self, subst: &Subst, ty: &Type) -> Type {
        match ty {
            Type::Var(name) => subst.get(name).cloned().unwrap_or_else(|| ty.clone()),
            Type::Arrow(t1, t2) => Type::Arrow(
                Box::new(self.apply_subst(subst, t1)),
                Box::new(self.apply_subst(subst, t2)),
            ),
            Type::Tuple(types) => {
                Type::Tuple(types.iter().map(|t| self.apply_subst(subst, t)).collect())
            }
            Type::Int | Type::Bool => ty.clone(),
        }
    }

    fn apply_subst_scheme(&self, subst: &Subst, scheme: &Scheme) -> Scheme {
        // Remove bindings for quantified variables to avoid capture
        let mut filtered_subst = subst.clone();
        for var in &scheme.vars {
            filtered_subst.remove(var);
        }
        Scheme {
            vars: scheme.vars.clone(),
            ty: self.apply_subst(&filtered_subst, &scheme.ty),
        }
    }

    fn apply_subst_env(&self, subst: &Subst, env: &Env) -> Env {
        env.iter()
            .map(|(k, v)| (k.clone(), self.apply_subst_scheme(subst, v)))
            .collect()
    }

    fn compose_subst(&self, s1: &Subst, s2: &Subst) -> Subst {
        let mut result = s1.clone();
        for (k, v) in s2 {
            result.insert(k.clone(), self.apply_subst(s1, v));
        }
        result
    }

    fn free_type_vars(&self, ty: &Type) -> HashSet<TyVar> {
        match ty {
            Type::Var(name) => {
                let mut set = HashSet::new();
                set.insert(name.clone());
                set
            }
            Type::Arrow(t1, t2) => {
                let mut set = self.free_type_vars(t1);
                set.extend(self.free_type_vars(t2));
                set
            }
            Type::Tuple(types) => {
                let mut set = HashSet::new();
                for t in types {
                    set.extend(self.free_type_vars(t));
                }
                set
            }
            Type::Int | Type::Bool => HashSet::new(),
        }
    }

    fn free_type_vars_scheme(&self, scheme: &Scheme) -> HashSet<TyVar> {
        let mut set = self.free_type_vars(&scheme.ty);
        // Remove quantified variables
        for var in &scheme.vars {
            set.remove(var);
        }
        set
    }

    fn free_type_vars_env(&self, env: &Env) -> HashSet<TyVar> {
        let mut set = HashSet::new();
        for scheme in env.values() {
            set.extend(self.free_type_vars_scheme(scheme));
        }
        set
    }

    fn generalize(&self, env: &Env, ty: &Type) -> Scheme {
        let type_vars = self.free_type_vars(ty);
        let env_vars = self.free_type_vars_env(env);
        let mut free_vars: Vec<_> = type_vars.difference(&env_vars).cloned().collect();
        free_vars.sort(); // Sort for deterministic behavior

        Scheme {
            vars: free_vars,
            ty: ty.clone(),
        }
    }

    fn instantiate(&mut self, scheme: &Scheme) -> Type {
        // Create fresh type variables for each quantified variable
        let mut subst = HashMap::new();
        for var in &scheme.vars {
            let fresh = self.fresh_tyvar();
            subst.insert(var.clone(), Type::Var(fresh));
        }

        self.apply_subst(&subst, &scheme.ty)
    }

    fn occurs_check(&self, var: &TyVar, ty: &Type) -> bool {
        match ty {
            Type::Var(name) => name == var,
            Type::Arrow(t1, t2) => self.occurs_check(var, t1) || self.occurs_check(var, t2),
            Type::Tuple(types) => types.iter().any(|t| self.occurs_check(var, t)),
            Type::Int | Type::Bool => false,
        }
    }

    fn unify(&self, t1: &Type, t2: &Type) -> Result<(Subst, InferenceTree)> {
        let input = format!("{} ~ {}", t1, t2);

        match (t1, t2) {
            (Type::Int, Type::Int) | (Type::Bool, Type::Bool) => {
                let tree = InferenceTree::new("Unify-Base", &input, "{}", vec![]);
                Ok((HashMap::new(), tree))
            }
            (Type::Var(v), ty) | (ty, Type::Var(v)) => {
                if ty == &Type::Var(v.clone()) {
                    let tree = InferenceTree::new("Unify-Var-Same", &input, "{}", vec![]);
                    Ok((HashMap::new(), tree))
                } else if self.occurs_check(v, ty) {
                    Err(InferenceError::OccursCheck {
                        var: v.clone(),
                        ty: ty.clone(),
                    })
                } else {
                    let mut subst = HashMap::new();
                    subst.insert(v.clone(), ty.clone());
                    let output = format!("{{{}/{}}}", ty, v);
                    let tree = InferenceTree::new("Unify-Var", &input, &output, vec![]);
                    Ok((subst, tree))
                }
            }
            (Type::Arrow(a1, a2), Type::Arrow(b1, b2)) => {
                let (s1, tree1) = self.unify(a1, b1)?;
                let a2_subst = self.apply_subst(&s1, a2);
                let b2_subst = self.apply_subst(&s1, b2);
                let (s2, tree2) = self.unify(&a2_subst, &b2_subst)?;
                let final_subst = self.compose_subst(&s2, &s1);
                let output = self.pretty_subst(&final_subst);
                let tree = InferenceTree::new("Unify-Arrow", &input, &output, vec![tree1, tree2]);
                Ok((final_subst, tree))
            }
            (Type::Tuple(ts1), Type::Tuple(ts2)) => {
                if ts1.len() != ts2.len() {
                    return Err(InferenceError::TupleLengthMismatch {
                        left_len: ts1.len(),
                        right_len: ts2.len(),
                    });
                }

                let mut subst = HashMap::new();
                let mut trees = Vec::new();

                for (t1, t2) in ts1.iter().zip(ts2.iter()) {
                    let t1_subst = self.apply_subst(&subst, t1);
                    let t2_subst = self.apply_subst(&subst, t2);
                    let (s, tree) = self.unify(&t1_subst, &t2_subst)?;
                    subst = self.compose_subst(&s, &subst);
                    trees.push(tree);
                }

                let output = self.pretty_subst(&subst);
                let tree = InferenceTree::new("Unify-Tuple", &input, &output, trees);
                Ok((subst, tree))
            }
            _ => Err(InferenceError::UnificationFailure {
                expected: t1.clone(),
                actual: t2.clone(),
            }),
        }
    }

    pub fn infer(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
        match expr {
            Expr::Lit(Lit::Int(_)) => self.infer_lit_int(env, expr),
            Expr::Lit(Lit::Bool(_)) => self.infer_lit_bool(env, expr),
            Expr::Var(name) => self.infer_var(env, expr, name),
            Expr::Abs(param, body) => self.infer_abs(env, expr, param, body),
            Expr::App(func, arg) => self.infer_app(env, expr, func, arg),
            Expr::Let(var, value, body) => self.infer_let(env, expr, var, value, body),
            Expr::Tuple(exprs) => self.infer_tuple(env, expr, exprs),
        }
    }

    /// T-LitInt: ─────────────────
    ///           Γ ⊢ n : Int
    fn infer_lit_int(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);
        let tree = InferenceTree::new("T-Int", &input, "Int", vec![]);
        Ok((HashMap::new(), Type::Int, tree))
    }

    /// T-LitBool: ─────────────────
    ///            Γ ⊢ b : Bool
    fn infer_lit_bool(&mut self, env: &Env, expr: &Expr) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);
        let tree = InferenceTree::new("T-Bool", &input, "Bool", vec![]);
        Ok((HashMap::new(), Type::Bool, tree))
    }

    /// T-Var: x : σ ∈ Γ    τ = inst(σ)
    ///        ─────────────────────────
    ///               Γ ⊢ x : τ
    fn infer_var(
        &mut self,
        env: &Env,
        expr: &Expr,
        name: &str,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        match env.get(name) {
            Some(scheme) => {
                let instantiated = self.instantiate(scheme);
                let output = format!("{}", instantiated);
                let tree = InferenceTree::new("T-Var", &input, &output, vec![]);
                Ok((HashMap::new(), instantiated, tree))
            }
            None => Err(InferenceError::UnboundVariable {
                name: name.to_string(),
            }),
        }
    }

    /// T-Lam: Γ, x : α ⊢ e : τ    α fresh
    ///        ─────────────────────────────
    ///           Γ ⊢ λx. e : α → τ
    fn infer_abs(
        &mut self,
        env: &Env,
        expr: &Expr,
        param: &str,
        body: &Expr,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let param_type = Type::Var(self.fresh_tyvar());
        let mut new_env = env.clone();
        // Insert a monomorphic scheme for the parameter
        let param_scheme = Scheme {
            vars: vec![],
            ty: param_type.clone(),
        };
        new_env.insert(param.to_string(), param_scheme);

        let (s1, body_type, tree1) = self.infer(&new_env, body)?;
        let param_type_subst = self.apply_subst(&s1, &param_type);
        let result_type = Type::Arrow(Box::new(param_type_subst), Box::new(body_type));

        let output = format!("{}", result_type);
        let tree = InferenceTree::new("T-Abs", &input, &output, vec![tree1]);
        Ok((s1, result_type, tree))
    }

    /// T-App: Γ ⊢ e₁ : τ₁    Γ ⊢ e₂ : τ₂    α fresh    S = unify(τ₁, τ₂ → α)
    ///        ──────────────────────────────────────────────────────────────
    ///                            Γ ⊢ e₁ e₂ : S(α)
    fn infer_app(
        &mut self,
        env: &Env,
        expr: &Expr,
        func: &Expr,
        arg: &Expr,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let result_type = Type::Var(self.fresh_tyvar());

        let (s1, func_type, tree1) = self.infer(env, func)?;
        let env_subst = self.apply_subst_env(&s1, env);
        let (s2, arg_type, tree2) = self.infer(&env_subst, arg)?;

        let func_type_subst = self.apply_subst(&s2, &func_type);
        let expected_func_type = Type::Arrow(Box::new(arg_type), Box::new(result_type.clone()));

        let (s3, tree3) = self.unify(&func_type_subst, &expected_func_type)?;

        let final_subst = self.compose_subst(&s3, &self.compose_subst(&s2, &s1));
        let final_type = self.apply_subst(&s3, &result_type);

        let output = format!("{}", final_type);
        let tree = InferenceTree::new("T-App", &input, &output, vec![tree1, tree2, tree3]);
        Ok((final_subst, final_type, tree))
    }

    /// T-Let: Γ ⊢ e₁ : τ₁    σ = gen(Γ, τ₁)    Γ, x : σ ⊢ e₂ : τ₂
    ///        ──────────────────────────────────────────────────────
    ///                     Γ ⊢ let x = e₁ in e₂ : τ₂
    fn infer_let(
        &mut self,
        env: &Env,
        expr: &Expr,
        var: &str,
        value: &Expr,
        body: &Expr,
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let (s1, value_type, tree1) = self.infer(env, value)?;
        let env_subst = self.apply_subst_env(&s1, env);
        let generalized_type = self.generalize(&env_subst, &value_type);

        let mut new_env = env_subst;
        new_env.insert(var.to_string(), generalized_type);

        let (s2, body_type, tree2) = self.infer(&new_env, body)?;

        let final_subst = self.compose_subst(&s2, &s1);
        let output = format!("{}", body_type);
        let tree = InferenceTree::new("T-Let", &input, &output, vec![tree1, tree2]);
        Ok((final_subst, body_type, tree))
    }

    /// T-Tuple: Γ ⊢ e₁ : τ₁    ...    Γ ⊢ eₙ : τₙ
    ///          ─────────────────────────────────────
    ///              Γ ⊢ (e₁, ..., eₙ) : (τ₁, ..., τₙ)
    fn infer_tuple(
        &mut self,
        env: &Env,
        expr: &Expr,
        exprs: &[Expr],
    ) -> Result<(Subst, Type, InferenceTree)> {
        let input = format!("{} ⊢ {} ⇒", self.pretty_env(env), expr);

        let mut subst = HashMap::new();
        let mut types = Vec::new();
        let mut trees = Vec::new();
        let mut current_env = env.clone();

        for expr in exprs {
            let (s, ty, tree) = self.infer(&current_env, expr)?;
            subst = self.compose_subst(&s, &subst);
            current_env = self.apply_subst_env(&s, &current_env);
            types.push(ty);
            trees.push(tree);
        }

        let result_type = Type::Tuple(types);
        let output = format!("{}", result_type);
        let tree = InferenceTree::new("T-Tuple", &input, &output, trees);
        Ok((subst, result_type, tree))
    }
}
pub fn run_inference(expr: &Expr) -> Result<InferenceTree> {
    let mut inference = TypeInference::new();
    let env = BTreeMap::new();
    let (_, _, tree) = inference.infer(&env, expr)?;
    Ok(tree)
}
pub fn infer_type_only(expr: &Expr) -> Result<Type> {
    let mut inference = TypeInference::new();
    let env = BTreeMap::new();
    let (_, ty, _) = inference.infer(&env, expr)?;
    Ok(ty)
}
}

Example Usage

To see Algorithm W in action, let’s type-check a polymorphic function that demonstrates let-polymorphism and generalization:

$ cargo run -- "let const = \\x -> \\y -> x in const 42 true"

This produces the following clean output showing the complete inference process:

Parsed expression: let const = λx.λy.x in const 42 true

Type inference successful!
Final type: Int

Inference trace:
T-Let: {} ⊢ let const = λx.λy.x in const 42 true ⇒ => Int
  T-Abs: {} ⊢ λx.λy.x ⇒ => t0 → t1 → t0
    T-Abs: {x: t0} ⊢ λy.x ⇒ => t1 → t0
      T-Var: {x: t0, y: t1} ⊢ x ⇒ => t0
  T-App: {const: forall t0 t1. t0 → t1 → t0} ⊢ const 42 true ⇒ => Int
    T-App: {const: forall t0 t1. t0 → t1 → t0} ⊢ const 42 ⇒ => t5 → Int
      T-Var: {const: forall t0 t1. t0 → t1 → t0} ⊢ const ⇒ => t4 → t5 → t4
      T-Int: {const: forall t0 t1. t0 → t1 → t0} ⊢ 42 ⇒ => Int
      Unify-Arrow: t4 → t5 → t4 ~ Int → t3 => {t5 → Int/t3, Int/t4}
        Unify-Var: t4 ~ Int => {Int/t4}
        Unify-Var: t5 → Int ~ t3 => {t5 → Int/t3}
    T-Bool: {const: forall t0 t1. t0 → t1 → t0} ⊢ true ⇒ => Bool
    Unify-Arrow: t5 → Int ~ Bool → t2 => {Int/t2, Bool/t5}
      Unify-Var: t5 ~ Bool => {Bool/t5}
      Unify-Var: Int ~ t2 => {Int/t2}

This trace shows several key aspects of Algorithm W:

1. Generalization: The lambda \x -> \y -> x initially gets type t0 → t1 → t0, but when bound to const in the let-expression, it’s generalized to ∀t0 t1. t0 → t1 → t0.

2. Instantiation: When const is used in the application, it gets instantiated with fresh type variables t4 and t5, allowing it to be used polymorphically.

3. Unification: The constraints from applying const to 42 and true get solved through unification, determining that t4 = Int, t5 = Bool, and the final result type is Int.

The final result is Int, showing that const 42 true correctly returns the first argument (42) regardless of the second argument’s type (true).

These interface functions demonstrate how Algorithm W can be embedded into larger systems. The tree-generating version supports educational tools and debuggers, while the type-only version provides the minimal interface needed for type checking during compilation.

Examples

The simplest cases involve literals and variables where type inference is straightforward. Integer literals have type Int, boolean literals have type Bool, and variables receive the types assigned to them in the typing environment.

Let’s examine how our implementation handles these basic cases using the test suite that validates our Algorithm W implementation.

#![allow(unused)]
fn main() {
Basic literals
42
true
false

Variables and identity
\x -> x
\f -> f
\x -> \y -> x
\x -> \y -> y

Application
(\x -> x) 42
(\x -> x) true
(\f -> \x -> f x) (\y -> y) 42

Let expressions
let x = 42 in x
let f = \x -> x in f
let id = \x -> x in id 42
let id = \x -> x in id true
let f = \x -> x in let g = \y -> y in f (g 42)

Tuples
(42, true)
(true, false, 42)
(\x -> (x, x)) 42
let pair = \x -> \y -> (x, y) in pair 42 true

Higher-order functions
\f -> \x -> f x
\f -> \x -> f (f x)
\f -> \g -> \x -> f (g x)
let twice = \f -> \x -> f (f x) in twice
let compose = \f -> \g -> \x -> f (g x) in compose

Complex examples
let K = \x -> \y -> x in K
let S = \f -> \g -> \x -> f x (g x) in S
let Y = \f -> (\x -> f (x x)) (\x -> f (x x)) in Y

Polymorphic examples (simplified for our basic implementation)
let id = \x -> x in (id, id)
let const = \x -> \y -> x in const
let flip = \f -> \x -> \y -> f y x in flip

Error cases (should fail)
\x -> x x
(\x -> x x) (\x -> x x)
}

These examples demonstrate the foundation of type inference. The literal 42 immediately receives type Int without any complex reasoning required. Boolean values like true and false similarly receive type Bool. The type system’s handling of these base cases forms the building blocks for more complex inference scenarios.

Function Types and Application

Function types represent the core of lambda calculus and functional programming. When we encounter a lambda abstraction, Algorithm W assigns a fresh type variable to the parameter and infers the type of the body. The resulting function type connects the parameter type with the return type through an arrow.

Function application drives the constraint generation that makes Algorithm W powerful. When applying a function to an argument, the algorithm generates constraints that the function’s type must be an arrow from the argument’s type to some result type. Unification then solves these constraints.

The identity function \x -> x provides the simplest example of polymorphic type inference. Algorithm W assigns a fresh type variable to the parameter x, then discovers that the body is just x itself. The resulting type t0 -> t0 captures the essence of the identity function: it accepts any type and returns the same type.

More complex function applications demonstrate how constraints propagate through the system. When we apply the identity function to the integer 42, the unification process discovers that the type variable t0 must be Int, yielding the final type Int for the entire expression.

Let Polymorphism

Let expressions introduce one of the most features of the Hindley-Milner type system: let-polymorphism. This mechanism allows variables bound in let expressions to be used with different types in different contexts, enabling flexible code reuse without sacrificing type safety.

The classic example involves binding the identity function in a let expression and then using it with different types. Algorithm W generalizes the type of the bound expression by abstracting over type variables that don’t appear free in the current environment. This generalization allows the same binding to be instantiated with fresh type variables at each use site.

Consider the expression let f = \x -> x in (f 42, f true). First, Algorithm W infers that the identity function has type t0 -> t0. During generalization, since t0 doesn’t appear in the environment, the system treats this as a polymorphic type that can be instantiated differently at each use.

When the algorithm encounters the first application f 42, it creates a fresh instance of the polymorphic type, say t1 -> t1, and unifies this with Int -> t2 (where t2 is the expected result type). This unification succeeds with t1 = Int and t2 = Int.

For the second application f true, Algorithm W creates another fresh instance t3 -> t3 and unifies this with Bool -> t4. This succeeds with t3 = Bool and t4 = Bool. The final result is a tuple type (Int, Bool), demonstrating how the same function can be used polymorphically.

Tuple Types and Complex Data Structures

Tuples provide a way to combine multiple values with potentially different types. Our Algorithm W implementation handles tuples by inferring the type of each component and combining them into a tuple type.

The expression (42, true) demonstrates basic tuple construction. Algorithm W infers that the first component has type Int and the second has type Bool, yielding the tuple type (Int, Bool). This extends naturally to nested tuples like ((1, 2), (true, false)) which receives type ((Int, Int), (Bool, Bool)).

Tuples interact interestingly with polymorphism. The expression let f = \x -> (x, x) in f 42 shows how a polymorphic function can construct tuples. The function f has type t0 -> (t0, t0), creating a tuple where both components have the same type as the input. When applied to 42, this yields type (Int, Int).

Type Error Detection

Algorithm W’s constraint-based approach makes it excellent at detecting type errors and providing meaningful error messages. When unification fails, the algorithm can identify exactly where and why types are incompatible.

Attempting to apply a non-function value like 42 true generates a type error because the integer 42 has type Int, but function application requires an arrow type. The unification of Int with Bool -> t0 fails, producing a clear error message.

More subtle errors arise from inconsistent uses of polymorphic functions. While Algorithm W handles let-polymorphism elegantly, it correctly rejects attempts to use functions in incompatible ways within the same scope.

Complex Inference Scenarios

Real-world programs often involve complex combinations of functions, applications, and data structures that test the full power of Algorithm W. These scenarios demonstrate how the algorithm handles intricate constraint propagation and substitution.

Higher-order functions provide particularly interesting examples. The expression \f -> \x -> f (f x) creates a function that applies another function twice. Algorithm W assigns fresh type variables and builds constraints that capture the relationships between all the types involved.

Let’s trace through this inference process. The outer lambda receives a fresh parameter type t0 for f. The inner lambda receives type t1 for x. The application f x requires that f has type t1 -> t2 for some fresh type t2. The outer application f (f x) then requires that f also has type t2 -> t3 for the final result type t3.

Unification resolves these constraints by discovering that t2 = t1 and the final type is (t1 -> t1) -> t1 -> t1. This captures the essence of function composition: given a function from some type to itself, produce a function that applies it twice.

The Y-Combinator

The Y-combinator represents a fundamental limitation of the Hindley-Milner type system. This famous fixed-point combinator cannot be typed in our system, illustrating an important boundary between what can be expressed with simple types and what requires more advanced type system features.

The Y-combinator is defined as:

$ cargo run -- "\\f -> (\\x -> f (x x)) (\\x -> f (x x))"
Type inference error: Occurs check failed: variable 't2' occurs in type t2 → t4

The failure occurs when attempting to type the self-application x x. When the type checker encounters this expression, it must assign the variable x a type that is simultaneously a function type (to be applied) and an argument type (to be passed to itself). This creates the constraint that some type variable t must equal t → τ for some type τ. The occurs check prevents this infinite type, ensuring the type system remains decidable and sound.

This limitation is not a bug but a fundamental characteristic of the Hindley-Milner type system. The Y-combinator requires more advanced type system features such as recursive types or explicit fixed-point operators. Real functional programming languages typically provide built-in recursion constructs like let rec that are handled specially by the type checker, avoiding the need for explicit fixed-point combinators at the term level.

The inability to type the Y-combinator illustrates an important trade-off in programming language design between expressiveness and decidability. While more powerful type systems exist that can handle self-application, they come at the cost of increased complexity and potentially undecidable type checking. Algorithm W chooses decidability and simplicity, making it practical for real programming languages while accepting certain expressiveness limitations.

Limitations: No Mutual Recursion

Our Algorithm W implementation has a significant limitation: it does not support mutual recursion. Mutually recursive definitions occur when two or more functions call each other, creating a circular dependency that requires careful handling during type inference.

Consider this example that would fail in our implementation:

let even = \n -> if n == 0 then true else odd (n - 1) in
let odd = \n -> if n == 0 then false else even (n - 1) in
odd 5

The problem is that when we encounter the definition of even, the function odd is not yet in scope, so the reference to odd in the body of even fails. Standard Algorithm W processes let-bindings sequentially, which makes mutual recursion impossible without additional machinery.

Supporting mutual recursion requires more techniques:

1. Dependency Analysis: The type checker must analyze the dependency graph of definitions to identify strongly connected components (groups of mutually recursive functions).

2. Simultaneous Inference: All functions in a mutually recursive group must be typed together, typically by assigning fresh type variables to each function initially, then solving all constraints simultaneously.

3. Generalization Delays: The generalization step (which introduces polymorphism) must be delayed until after all mutual dependencies are resolved.

These extensions significantly complicate the implementation and are beyond the scope of our simple Algorithm W demonstration.

Performance (Or Lack Thereof)

While Algorithm W is theoretically elegant, practical implementations must consider performance characteristics. The algorithm’s complexity depends on the size of expressions and the complexity of types involved. Most real programs involve relatively simple type constraints that Algorithm W handles efficiently.

Our implementation demonstrates the core algorithms without the optimizations found in production compilers. Industrial-strength implementations often employ techniques like type-directed compilation, constraint caching, and incremental type checking to handle large codebases efficiently.

System F Type Checker

While the simple lambda calculus with HM-style type schemes introduces safety and limited polymorphism, it is also very rigid and not very expressive. So now we move onto System F, also known as the polymorphic lambda calculus, smashes this limitation by introducing parametric polymorphism: the ability to write a single piece of code that is generic over types. This is the theoretical foundation for generics in languages like Rust, Java, and C++, and it represents a monumental leap in expressive power.

To achieve this, System F extends both the term and type languages. The most crucial addition to the type language is the universal quantifier, \(\forall\) (forall), which allows us to express the type of a polymorphic function.

\[ \begin{align*} \text{types} \quad \tau &::= \alpha \mid \tau_1 \to \tau_2 \mid \forall \alpha . \tau \mid \text{Int} \mid \text{Bool} \\ \text{expressions} \quad e &::= x \mid \lambda x : \tau . e \mid e_1 \ e_2 \mid \Lambda \alpha . e \mid e[\tau] \mid \dots \end{align*} \]

The type \(\forall \alpha . \tau\) is represented by Type::Forall(String, Box<Type>). The term language now includes two new constructs specifically for handling polymorphism:

1. Type Abstraction (\(\Lambda \alpha . e\)): This creates a polymorphic function. The capital lambda (\(\Lambda\)) signifies that we are abstracting over a type variable \(\alpha\), not a term variable. This expression is represented by Expr::TAbs(String, Box<Expr>).
2. Type Application (\(e[\tau]\)): This specializes a polymorphic function by applying it to a concrete type \(\tau\). It’s how we use a generic function. This is represented by Expr::TApp(Box<Expr>, Box<Type>).

A key difference from the untyped lambda calculus is that our term-level abstractions (\(\lambda\)) are now explicitly annotated: \(\lambda x : \tau . e\). The programmer must declare the type of the function’s parameter. This is a hallmark of System F; its power comes at the cost of full type inference, necessitating these annotations.

The quintessential example of System F’s power is the polymorphic identity function, id. We can now write a single identity function that works for any type \(\alpha\). We create it with a type abstraction and give its parameter x the generic type \(\alpha\):

\[ \text{id} = \Lambda \alpha . \lambda x : \alpha . x \]

The type of this function is \(\forall \alpha . \alpha \to \alpha\). This type says: “For all types \(\alpha\), I am a function that takes an argument of type \(\alpha\) and returns a value of type \(\alpha\).”

To use this polymorphic function, we apply it to a type using type application. For example, to get an identity function specifically for integers, we apply id to the type \(\text{Int}\):

\[ \text{id}[\text{Int}] \]

This is a computation that happens at the type-checking level. The result of this type application is a new, specialized expression, \(\lambda x : \text{Int} . x\), which has the corresponding specialized type \(\text{Int} \to \text{Int}\). We can then apply this specialized function to a value, like \((\lambda x : \text{Int} . x) \ 5\), which finally reduces to \(5\). This separation of type-level application (\(e[\tau]\)) and term-level application (\(e_1 \ e_2\)) is fundamental to System F. It provides a powerful and safe way to write generic code, but as we have seen, it forces the programmer to be more explicit with annotations, a tradeoff that directly motivates the bidirectional algorithms used in more modern language implementations.

Abstract Syntax Trees

Let’s look at the full abstract syntax tree for both the type language (enum Type) and the term language (enum Expr):

The Type Level

#![allow(unused)]
fn main() {
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    Var(String),
    App(Box<Expr>, Box<Expr>),
    Abs(String, Box<Type>, Box<Expr>),
    TApp(Box<Expr>, Box<Type>),
    TAbs(String, Box<Expr>),
    Ann(Box<Expr>, Box<Type>), // Type annotation: e : T
    LitInt(i64),
    LitBool(bool),
    // New constructs
    Let(String, Box<Expr>, Box<Expr>),           // let x = e1 in e2
    IfThenElse(Box<Expr>, Box<Expr>, Box<Expr>), // if e1 then e2 else e3
    BinOp(BinOp, Box<Expr>, Box<Expr>),          // Binary operations
}
#[derive(Debug, Clone, PartialEq)]
pub enum BinOp {
    // Arithmetic
    Add, // +
    Sub, // -
    Mul, // *
    Div, // /
    // Boolean
    And, // &&
    Or,  // ||
    // Comparison
    Eq, // ==
    Ne, // !=
    Lt, // <
    Le, // <=
    Gt, // >
    Ge, // >=
}
impl std::fmt::Display for Type {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Type::Var(name) => write!(f, "{}", name),
            Type::ETVar(name) => write!(f, "^{}", name),
            Type::Int => write!(f, "Int"),
            Type::Bool => write!(f, "Bool"),
            Type::Arrow(t1, t2) => {
                // Add parentheses around left side if it's an arrow to avoid ambiguity
                match t1.as_ref() {
                    Type::Arrow(_, _) => write!(f, "({}) -> {}", t1, t2),
                    _ => write!(f, "{} -> {}", t1, t2),
                }
            }
            Type::Forall(var, ty) => write!(f, "∀{}. {}", var, ty),
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    Var(String),                 // α (ordinary type variable)
    ETVar(String),               // ^α (existential type variable)
    Arrow(Box<Type>, Box<Type>), // A -> B
    Forall(String, Box<Type>),   // ∀α. A
    Int,                         // Int
    Bool,                        // Bool
}
}

The Type enum defines the grammar for all possible types in our language. It’s the vocabulary we use to describe our expressions.

• Type::Var(String): This represents a simple type variable, like \(\alpha\) or \(\beta\). These are placeholders for types that will be specified later, often used in polymorphic functions.

• Type::ETVar(String): This represents an “existential type variable,” often written as \(^\alpha\). These are a special kind of variable used internally by the type checker, particularly in bidirectional algorithms. They act as placeholders for an unknown type that the algorithm needs to infer.

• Type::Arrow(Box<Type>, Box<Type>): This is the function type, \(\tau_1 \to \tau_2\). It represents a function that takes an argument of the first type and returns a result of the second type.

• Type::Forall(String, Box<Type>): This is the universal quantifier, \(\forall \alpha . \tau\). It is the cornerstone of System F and represents a polymorphic type. It reads: “for all types \(\alpha\), the following type \(\tau\) holds.” The String is the name of the type variable \(\alpha\) being bound.

• Type::Int and Type::Bool: These are primitive, or base, types. They are concrete types that are built directly into the language, representing 64-bit integers and booleans, respectively.

The Value Level

The Expr enum defines the grammar for all runnable expressions or terms. This is the code that performs computations.

#![allow(unused)]
fn main() {
#[derive(Debug, Clone, PartialEq)]
pub enum BinOp {
    // Arithmetic
    Add, // +
    Sub, // -
    Mul, // *
    Div, // /
    // Boolean
    And, // &&
    Or,  // ||
    // Comparison
    Eq, // ==
    Ne, // !=
    Lt, // <
    Le, // <=
    Gt, // >
    Ge, // >=
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    Var(String),                 // α (ordinary type variable)
    ETVar(String),               // ^α (existential type variable)
    Arrow(Box<Type>, Box<Type>), // A -> B
    Forall(String, Box<Type>),   // ∀α. A
    Int,                         // Int
    Bool,                        // Bool
}
impl std::fmt::Display for Type {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Type::Var(name) => write!(f, "{}", name),
            Type::ETVar(name) => write!(f, "^{}", name),
            Type::Int => write!(f, "Int"),
            Type::Bool => write!(f, "Bool"),
            Type::Arrow(t1, t2) => {
                // Add parentheses around left side if it's an arrow to avoid ambiguity
                match t1.as_ref() {
                    Type::Arrow(_, _) => write!(f, "({}) -> {}", t1, t2),
                    _ => write!(f, "{} -> {}", t1, t2),
                }
            }
            Type::Forall(var, ty) => write!(f, "∀{}. {}", var, ty),
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    Var(String),
    App(Box<Expr>, Box<Expr>),
    Abs(String, Box<Type>, Box<Expr>),
    TApp(Box<Expr>, Box<Type>),
    TAbs(String, Box<Expr>),
    Ann(Box<Expr>, Box<Type>), // Type annotation: e : T
    LitInt(i64),
    LitBool(bool),
    // New constructs
    Let(String, Box<Expr>, Box<Expr>),           // let x = e1 in e2
    IfThenElse(Box<Expr>, Box<Expr>, Box<Expr>), // if e1 then e2 else e3
    BinOp(BinOp, Box<Expr>, Box<Expr>),          // Binary operations
}
}

• Var(String): A term-level variable, like x or f. It refers to a value that is in scope, such as a function parameter or a let-bound variable.

• App(Box<Expr>, Box<Expr>): This is function application, \(e_1 \ e_2\). It represents calling the function \(e_1\) with the argument \(e_2\).

• Abs(String, Box<Type>, Box<Expr>): This is a typed lambda abstraction, \(\lambda x : \tau . e\). It defines an anonymous function. Unlike in the pure lambda calculus, the parameter (String) must have an explicit type annotation (Box<Type>). The final Box<Expr> is the function’s body.

• TApp(Box<Expr>, Box<Type>): This is type application, \(e[\tau]\). This is the mechanism for specializing a polymorphic function. The expression \(e\) must have a Forall type, and this construct applies it to a specific type \(\tau\), effectively “filling in” the generic type parameter.

• TAbs(String, Box<Expr>): This is type abstraction, \(\Lambda \alpha . e\). This is how we create a polymorphic function. It introduces a new type variable (String) that can be used within the expression body (Box<Expr>).

• Ann(Box<Expr>, Box<Type>): This represents a type annotation, \(e : T\). It’s an explicit instruction to the type checker, asserting that the expression \(e\) should have the type \(T\). This is invaluable in a bidirectional system for guiding the inference process and resolving ambiguity.

• LitInt(i64) and LitBool(bool): These are literal values. They represent concrete, primitive values that are “wired into” the language, corresponding to the base types Int and Bool. They are the simplest form of expression, representing a constant value.

Implementation

Ok, now we move beyond the trivial type systems from the 1970s and into the fun stuff from the 1980s! System F’s bidirectional type checking represents a approach to handling polymorphic types without requiring full type annotations everywhere. The bidirectional approach splits type checking into two complementary modes: inference (synthesizing types from expressions) and checking (verifying that expressions conform to expected types). This division allows the system to gracefully handle situations where types are partially known or completely unknown, making the language more ergonomic while preserving type safety.

Typing Rules

Before diving into the implementation details, let’s establish the formal typing rules that govern System F. Buckle up, because we’re about to embark into the fun magical land of type-level wizardry! We’ll be introducing a few new symbols that might look intimidating at first, but they’re really not that scary once you get used to them.

• \( \Gamma \) (Gamma) - The typing context, which is like a dictionary that maps variables to their types and keeps track of what we know so far.

• \( \vdash \) (Turnstile) - The “proves” or “entails” symbol. When we write \( \Gamma \vdash e \Rightarrow A \), we’re saying “given context \( \Gamma \), expression \( e \) synthesizes type \( A \).”

• \( \Rightarrow \) (Double Right Arrow) - Inference mode, where we’re asking “what type does this expression have?” The type checker figures it out for us.

• \( \Leftarrow \) (Double Left Arrow) - Checking mode, where we’re saying “please verify this expression has the expected type.” We already know what type we want.

• \( \forall \) (Forall) - Universal quantification, meaning “for any type.” When we see \( \forall \alpha. A \), it means “for any type \( \alpha \), we have type \( A \).”

• \( \hat{\alpha} \) (Hat Alpha) - Existential type variables, which are like type-level unknowns that the system solves during inference. Think of them as placeholders that get filled in later.

• \( \bullet \) (Bullet) - The application judgment symbol used in our inference rules. When we write \( A \bullet e \Rightarrow B \), we’re saying “applying type \( A \) to expression \( e \) yields type \( B \).”

• \( <: \) (Subtype) - The subtyping relation, expressing that one type is “more specific” than another. For example, \( \text{Int} <: \forall \alpha. \alpha \) would mean Int is a subtype of the polymorphic type.

• \( [B/\alpha]A \) - Type substitution, replacing all occurrences of type variable \( \alpha \) with type \( B \) in type \( A \). This is how we instantiate polymorphic types.

Now that we’ve equipped ourselves with this symbolic toolkit, let’s see how these pieces combine to create the elegant machinery of System F type checking.

Basic Rules

The variable rule looks up types from the context:

\[ \frac{x : A \in \Gamma}{\Gamma \vdash x \Rightarrow A} \text{(T-Var)} \]

Application checks that the function type matches the argument:

\[ \frac{\Gamma \vdash e_1 \Rightarrow A \to B \quad \Gamma \vdash e_2 \Leftarrow A}{\Gamma \vdash e_1 ; e_2 \Rightarrow B} \text{(T-App)} \]

Lambda abstraction introduces a new variable binding:

\[ \frac{\Gamma, x : A \vdash e \Leftarrow B}{\Gamma \vdash \lambda x. e \Leftarrow A \to B} \text{(T-Abs)} \]

Polymorphic Rules

Universal introduction allows us to generalize over type variables:

\[ \frac{\Gamma, \alpha \vdash e \Leftarrow A}{\Gamma \vdash e \Leftarrow \forall \alpha. A} \text{(T-ForallI)} \]

Universal elimination instantiates polymorphic types:

\[ \frac{\Gamma \vdash e \Rightarrow \forall \alpha. A}{\Gamma \vdash e \Rightarrow [B/\alpha]A} \text{(T-ForallE)} \]

Type annotation allows switching from checking to inference mode:

\[ \frac{\Gamma \vdash e \Leftarrow A}{\Gamma \vdash (e : A) \Rightarrow A} \text{(T-Instr)} \]

Primitive Type Rules

Integer literals have type Int:

\[ \frac{}{\Gamma \vdash n \Rightarrow \text{Int}} \text{(T-LitInt)} \]

Boolean literals have type Bool:

\[ \frac{}{\Gamma \vdash \text{true} \Rightarrow \text{Bool}} \text{(T-LitBool)} \]

Control Flow Rules

Let bindings introduce local variables:

\[ \frac{\Gamma \vdash e_1 \Rightarrow A \quad \Gamma, x : A \vdash e_2 \Rightarrow B}{\Gamma \vdash \text{let } x = e_1 \text{ in } e_2 \Rightarrow B} \text{(T-Let)} \]

Conditional expressions require Bool conditions and matching branch types:

\[ \frac{\Gamma \vdash e_1 \Leftarrow \text{Bool} \quad \Gamma \vdash e_2 \Rightarrow A \quad \Gamma \vdash e_3 \Leftarrow A}{\Gamma \vdash \text{if } e_1 \text{ then } e_2 \text{ else } e_3 \Rightarrow A} \text{(T-If)} \]

Binary Operation Rules

Arithmetic operations take two integers and return an integer:

\[ \frac{\Gamma \vdash e_1 \Leftarrow \text{Int} \quad \Gamma \vdash e_2 \Leftarrow \text{Int}}{\Gamma \vdash e_1 \oplus e_2 \Rightarrow \text{Int}} \text{(T-Arith)} \]

Boolean operations take two booleans and return a boolean:

\[ \frac{\Gamma \vdash e_1 \Leftarrow \text{Bool} \quad \Gamma \vdash e_2 \Leftarrow \text{Bool}}{\Gamma \vdash e_1 \land e_2 \Rightarrow \text{Bool}} \text{(T-Bool)} \]

Comparison operations take two integers and return a boolean:

\[ \frac{\Gamma \vdash e_1 \Leftarrow \text{Int} \quad \Gamma \vdash e_2 \Leftarrow \text{Int}}{\Gamma \vdash e_1 < e_2 \Rightarrow \text{Bool}} \text{(T-Cmp)} \]

Equality operations are polymorphic and work on any type:

\[ \frac{\Gamma \vdash e_1 \Rightarrow A \quad \Gamma \vdash e_2 \Leftarrow A}{\Gamma \vdash e_1 = e_2 \Rightarrow \text{Bool}} \text{(T-Eq)} \]

Bidirectional Rules

The mode switch allows inference results to be checked:

\[ \frac{\Gamma \vdash e \Rightarrow A}{\Gamma \vdash e \Leftarrow A} \text{(T-Sub)} \]

Existential variables are introduced for unknown types:

\[ \frac{\Gamma, \hat{\alpha} \vdash e \Rightarrow A}{\Gamma \vdash e \Rightarrow [\hat{\alpha}/\alpha]A} \text{(T-InstL)} \]

Application Inference Rules

Application inference handles complex cases where the function type is not immediately known:

Application with arrow types: \[ \frac{\Gamma \vdash e_2 \Leftarrow A}{\Gamma \vdash A \to B \bullet e_2 \Rightarrow B} \text{(T-AppArrow)} \]

Application with existential variables: \[ \frac{\Gamma[\hat{\alpha} := \hat{\alpha_1} \to \hat{\alpha_2}], \hat{\alpha_1}, \hat{\alpha_2} \vdash e_2 \Leftarrow \hat{\alpha_1}}{\Gamma \vdash \hat{\alpha} \bullet e_2 \Rightarrow \hat{\alpha_2}} \text{(T-AppEVar)} \]

In these rules, \( \Rightarrow \) indicates inference mode (synthesizing a type), while \( \Leftarrow \) indicates checking mode (verifying against an expected type). The hat notation \( \hat{\alpha} \) denotes existential type variables that the system solves during inference.

Core Data Structures

Context and Environment Management

The bidirectional algorithm maintains a context that tracks multiple kinds of bindings and constraints. Our context system needs to handle not just term variables and their types, but also type variables, existential variables, and the relationships between them.

#![allow(unused)]
fn main() {
use std::collections::HashSet;
use std::fmt;
use crate::ast::{BinOp, Expr, Type};
use crate::errors::{TypeError, TypeResult};
pub type TyVar = String;
pub type TmVar = String;
impl fmt::Display for Entry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Entry::VarBnd(x, ty) => write!(f, "{}: {}", x, ty),
            Entry::TVarBnd(a) => write!(f, "{}", a),
            Entry::ETVarBnd(a) => write!(f, "^{}", a),
            Entry::SETVarBnd(a, ty) => write!(f, "^{} = {}", a, ty),
            Entry::Mark(a) => write!(f, "${}", a),
        }
    }
}
#[derive(Debug, Clone)]
pub struct Context(Vec<Entry>);
impl Default for Context {
    fn default() -> Self {
        Self::new()
    }
}
impl Context {
    pub fn new() -> Self {
        Context(Vec::new())
    }

    pub fn push(&mut self, entry: Entry) {
        self.0.push(entry);
    }

    pub fn find<F>(&self, predicate: F) -> Option<&Entry>
    where
        F: Fn(&Entry) -> bool, {
        self.0.iter().find(|entry| predicate(entry))
    }

    pub fn break3<F>(&self, predicate: F) -> (Vec<Entry>, Option<Entry>, Vec<Entry>)
    where
        F: Fn(&Entry) -> bool, {
        if let Some(pos) = self.0.iter().position(predicate) {
            let left = self.0[..pos].to_vec();
            let middle = self.0[pos].clone();
            let right = self.0[pos + 1..].to_vec();
            (left, Some(middle), right)
        } else {
            (self.0.clone(), None, Vec::new())
        }
    }

    pub fn from_parts(left: Vec<Entry>, middle: Entry, right: Vec<Entry>) -> Self {
        let mut ctx = left;
        ctx.push(middle);
        ctx.extend(right);
        Context(ctx)
    }
}
impl fmt::Display for Context {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let entries: Vec<String> = self.0.iter().rev().map(|e| e.to_string()).collect();
        write!(f, "{}", entries.join(", "))
    }
}
#[derive(Debug)]
pub struct InferenceTree {
    pub rule: String,
    pub input: String,
    pub output: String,
    pub children: Vec<InferenceTree>,
}
impl InferenceTree {
    fn new(rule: &str, input: &str, output: &str, children: Vec<InferenceTree>) -> Self {
        Self {
            rule: rule.to_string(),
            input: input.to_string(),
            output: output.to_string(),
            children,
        }
    }
}
impl fmt::Display for InferenceTree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.display_with_indent(f, 0)
    }
}
impl InferenceTree {
    fn display_with_indent(&self, f: &mut fmt::Formatter, indent: usize) -> fmt::Result {
        let prefix = "  ".repeat(indent);
        writeln!(
            f,
            "{}{}: {} => {}",
            prefix, self.rule, self.input, self.output
        )?;
        for child in &self.children {
            child.display_with_indent(f, indent + 1)?;
        }
        Ok(())
    }
}
pub struct BiDirectional {
    counter: usize,
}
impl Default for BiDirectional {
    fn default() -> Self {
        Self::new()
    }
}
#[allow(clippy::only_used_in_recursion)]
impl BiDirectional {
    pub fn new() -> Self {
        Self { counter: 0 }
    }

    fn fresh_tyvar(&mut self) -> TyVar {
        let var = format!("α{}", self.counter);
        self.counter += 1;
        var
    }

    fn is_mono(&self, ty: &Type) -> bool {
        match ty {
            Type::Int | Type::Bool | Type::Var(_) | Type::ETVar(_) => true,
            Type::Arrow(t1, t2) => self.is_mono(t1) && self.is_mono(t2),
            Type::Forall(_, _) => false,
        }
    }

    fn free_vars(&self, ty: &Type) -> HashSet<TyVar> {
        match ty {
            Type::Var(name) | Type::ETVar(name) => {
                let mut set = HashSet::new();
                set.insert(name.clone());
                set
            }
            Type::Arrow(t1, t2) => {
                let mut set = self.free_vars(t1);
                set.extend(self.free_vars(t2));
                set
            }
            Type::Forall(var, ty) => {
                let mut set = self.free_vars(ty);
                set.remove(var);
                set
            }
            Type::Int | Type::Bool => HashSet::new(),
        }
    }

    fn subst_type(&self, var: &TyVar, replacement: &Type, ty: &Type) -> Type {
        match ty {
            Type::Var(name) if name == var => replacement.clone(),
            Type::ETVar(name) if name == var => replacement.clone(),
            Type::Var(_) | Type::ETVar(_) | Type::Int | Type::Bool => ty.clone(),
            Type::Arrow(t1, t2) => Type::Arrow(
                Box::new(self.subst_type(var, replacement, t1)),
                Box::new(self.subst_type(var, replacement, t2)),
            ),
            Type::Forall(bound_var, body) => {
                if bound_var == var {
                    ty.clone() // Variable is shadowed
                } else {
                    Type::Forall(
                        bound_var.clone(),
                        Box::new(self.subst_type(var, replacement, body)),
                    )
                }
            }
        }
    }

    pub fn apply_ctx_type(&self, ctx: &Context, ty: &Type) -> Type {
        let mut current = ty.clone();
        let mut changed = true;

        // Keep applying substitutions until no more changes occur
        while changed {
            changed = false;
            let new_type = self.apply_ctx_type_once(ctx, &current);
            if new_type != current {
                changed = true;
                current = new_type;
            }
        }

        current
    }

    fn apply_ctx_type_once(&self, ctx: &Context, ty: &Type) -> Type {
        match ty {
            Type::ETVar(a) => {
                if let Some(Entry::SETVarBnd(_, replacement)) =
                    ctx.find(|entry| matches!(entry, Entry::SETVarBnd(name, _) if name == a))
                {
                    // Recursively apply substitutions to the replacement type
                    self.apply_ctx_type_once(ctx, replacement)
                } else {
                    ty.clone()
                }
            }
            Type::Arrow(t1, t2) => Type::Arrow(
                Box::new(self.apply_ctx_type_once(ctx, t1)),
                Box::new(self.apply_ctx_type_once(ctx, t2)),
            ),
            Type::Forall(var, body) => {
                Type::Forall(var.clone(), Box::new(self.apply_ctx_type_once(ctx, body)))
            }
            _ => ty.clone(),
        }
    }

    fn before(&self, ctx: &Context, a: &TyVar, b: &TyVar) -> bool {
        let pos_a = ctx
            .0
            .iter()
            .position(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
        let pos_b = ctx
            .0
            .iter()
            .position(|entry| matches!(entry, Entry::ETVarBnd(name) if name == b));

        match (pos_a, pos_b) {
            (Some(pa), Some(pb)) => pa > pb, // Later in the context means earlier in ordering
            _ => false,
        }
    }

    /// Check if a type variable occurs in a type (occurs check for unification)
    /// This prevents creating infinite types when solving ^α := τ by ensuring α
    /// ∉ τ
    fn occurs_check(&self, var: &TyVar, ty: &Type) -> bool {
        match ty {
            Type::Var(name) | Type::ETVar(name) => name == var,
            Type::Arrow(t1, t2) => self.occurs_check(var, t1) || self.occurs_check(var, t2),
            Type::Forall(bound_var, body) => {
                // If the variable is shadowed by the forall binding, it doesn't occur
                if bound_var == var {
                    false
                } else {
                    self.occurs_check(var, body)
                }
            }
            Type::Int | Type::Bool => false,
        }
    }

    // ========== TYPING RULE METHODS ==========
    // Each method implements a specific typing rule from System F

    /// T-Var: Variable lookup rule
    /// Γ, x:A ⊢ x ⇒ A
    fn infer_var(
        &self,
        ctx: &Context,
        x: &str,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        if let Some(Entry::VarBnd(_, ty)) =
            ctx.find(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x))
        {
            let output = format!("{} ⇒ {} ⊣ {}", input, ty, ctx);
            Ok((
                ty.clone(),
                ctx.clone(),
                InferenceTree::new("InfVar", input, &output, vec![]),
            ))
        } else {
            Err(TypeError::UnboundVariable {
                name: x.to_string(),
                expr: None,
            })
        }
    }

    /// T-LitInt: Integer literal rule
    /// ⊢ n ⇒ Int
    fn infer_lit_int(
        &self,
        ctx: &Context,
        _n: i64,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let output = format!("{} ⇒ Int ⊣ {}", input, ctx);
        Ok((
            Type::Int,
            ctx.clone(),
            InferenceTree::new("InfLitInt", input, &output, vec![]),
        ))
    }

    /// T-LitBool: Boolean literal rule
    /// ⊢ b ⇒ Bool
    fn infer_lit_bool(
        &self,
        ctx: &Context,
        _b: bool,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let output = format!("{} ⇒ Bool ⊣ {}", input, ctx);
        Ok((
            Type::Bool,
            ctx.clone(),
            InferenceTree::new("InfLitBool", input, &output, vec![]),
        ))
    }

    /// T-Abs: Lambda abstraction rule
    /// Γ,x:A ⊢ e ⇐ B
    /// ─────────────────────
    /// Γ ⊢ λx:A.e ⇒ A → B
    fn infer_abs(
        &mut self,
        ctx: &Context,
        x: &str,
        param_ty: &Type,
        body: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let b = self.fresh_tyvar();
        let mut new_ctx = ctx.clone();
        new_ctx.push(Entry::VarBnd(x.to_string(), param_ty.clone()));
        new_ctx.push(Entry::ETVarBnd(b.clone()));

        let (ctx1, tree) = self.check(&new_ctx, body, &Type::ETVar(b.clone()))?;
        let (left, _, right) =
            ctx1.break3(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x));
        // Preserve solved existential variable bindings from the left context
        let mut final_ctx_entries = left
            .into_iter()
            .filter(|entry| matches!(entry, Entry::SETVarBnd(_, _)))
            .collect::<Vec<_>>();
        final_ctx_entries.extend(right);
        let final_ctx = Context(final_ctx_entries);
        let result_ty = Type::Arrow(Box::new(param_ty.clone()), Box::new(Type::ETVar(b)));
        let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, final_ctx);
        Ok((
            result_ty,
            final_ctx,
            InferenceTree::new("InfLam", input, &output, vec![tree]),
        ))
    }

    /// T-App: Function application rule
    /// Γ ⊢ e1 ⇒ A   Γ ⊢ e2 • A ⇒⇒ C
    /// ────────────────────────────────
    /// Γ ⊢ e1 e2 ⇒ C
    fn infer_application(
        &mut self,
        ctx: &Context,
        func: &Expr,
        arg: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (func_ty, ctx1, tree1) = self.infer(ctx, func)?;
        let func_ty_applied = self.apply_ctx_type(&ctx1, &func_ty);
        let (result_ty, ctx2, tree2) = self.infer_app(&ctx1, &func_ty_applied, arg)?;
        let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, ctx2);
        Ok((
            result_ty,
            ctx2,
            InferenceTree::new("InfApp", input, &output, vec![tree1, tree2]),
        ))
    }

    /// T-Let: Let binding rule
    /// Γ ⊢ e1 ⇒ A   Γ,x:A ⊢ e2 ⇒ B
    /// ─────────────────────────────
    /// Γ ⊢ let x = e1 in e2 ⇒ B
    fn infer_let(
        &mut self,
        ctx: &Context,
        x: &str,
        e1: &Expr,
        e2: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (ty1, ctx1, tree1) = self.infer(ctx, e1)?;
        let mut new_ctx = ctx1.clone();
        new_ctx.push(Entry::VarBnd(x.to_string(), ty1));
        let (ty2, ctx2, tree2) = self.infer(&new_ctx, e2)?;
        let (left, _, right) =
            ctx2.break3(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x));
        // Preserve solved existential variable bindings from the left context
        let mut final_ctx_entries = left
            .into_iter()
            .filter(|entry| matches!(entry, Entry::SETVarBnd(_, _)))
            .collect::<Vec<_>>();
        final_ctx_entries.extend(right);
        let final_ctx = Context(final_ctx_entries);
        let output = format!("{} ⇒ {} ⊣ {}", input, ty2, final_ctx);
        Ok((
            ty2,
            final_ctx,
            InferenceTree::new("InfLet", input, &output, vec![tree1, tree2]),
        ))
    }

    /// T-If: Conditional rule
    /// Γ ⊢ e1 ⇐ Bool   Γ ⊢ e2 ⇒ A   Γ ⊢ e3 ⇒ A
    /// ──────────────────────────────────────────
    /// Γ ⊢ if e1 then e2 else e3 ⇒ A
    fn infer_if(
        &mut self,
        ctx: &Context,
        e1: &Expr,
        e2: &Expr,
        e3: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (ctx1, tree1) = self.check(ctx, e1, &Type::Bool)?;
        let (ty2, ctx2, tree2) = self.infer(&ctx1, e2)?;
        let (ty3, ctx3, tree3) = self.infer(&ctx2, e3)?;
        // Ensure both branches have the same type
        let (unified_ctx, tree_unify) = self.subtype(&ctx3, &ty2, &ty3)?;
        let output = format!("{} ⇒ {} ⊣ {}", input, ty2, unified_ctx);
        Ok((
            ty2,
            unified_ctx,
            InferenceTree::new(
                "InfIf",
                input,
                &output,
                vec![tree1, tree2, tree3, tree_unify],
            ),
        ))
    }

    /// T-BinOp: Binary operation rules
    fn infer_binop(
        &mut self,
        ctx: &Context,
        op: &BinOp,
        e1: &Expr,
        e2: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        match op {
            // T-Arith: Int → Int → Int
            BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div => {
                let (ctx1, tree1) = self.check(ctx, e1, &Type::Int)?;
                let (ctx2, tree2) = self.check(&ctx1, e2, &Type::Int)?;
                let output = format!("{} ⇒ Int ⊣ {}", input, ctx2);
                Ok((
                    Type::Int,
                    ctx2,
                    InferenceTree::new("InfArith", input, &output, vec![tree1, tree2]),
                ))
            }
            // T-Bool: Bool → Bool → Bool
            BinOp::And | BinOp::Or => {
                let (ctx1, tree1) = self.check(ctx, e1, &Type::Bool)?;
                let (ctx2, tree2) = self.check(&ctx1, e2, &Type::Bool)?;
                let output = format!("{} ⇒ Bool ⊣ {}", input, ctx2);
                Ok((
                    Type::Bool,
                    ctx2,
                    InferenceTree::new("InfBool", input, &output, vec![tree1, tree2]),
                ))
            }
            // T-Cmp: Int → Int → Bool
            BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge => {
                let (ctx1, tree1) = self.check(ctx, e1, &Type::Int)?;
                let (ctx2, tree2) = self.check(&ctx1, e2, &Type::Int)?;
                let output = format!("{} ⇒ Bool ⊣ {}", input, ctx2);
                Ok((
                    Type::Bool,
                    ctx2,
                    InferenceTree::new("InfCmp", input, &output, vec![tree1, tree2]),
                ))
            }
            // T-Eq: ∀α. α → α → Bool
            BinOp::Eq | BinOp::Ne => {
                let (ty1, ctx1, tree1) = self.infer(ctx, e1)?;
                let (ctx2, tree2) = self.check(&ctx1, e2, &ty1)?;
                let output = format!("{} ⇒ Bool ⊣ {}", input, ctx2);
                Ok((
                    Type::Bool,
                    ctx2,
                    InferenceTree::new("InfEq", input, &output, vec![tree1, tree2]),
                ))
            }
        }
    }

    fn inst_l(
        &mut self,
        ctx: &Context,
        a: &TyVar,
        ty: &Type,
    ) -> TypeResult<(Context, InferenceTree)> {
        let input = format!("{} ⊢ ^{} :=< {}", ctx, a, ty);

        match ty {
            Type::ETVar(b) if self.before(ctx, a, b) => {
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == b));
                let new_ctx = Context::from_parts(
                    left,
                    Entry::SETVarBnd(b.clone(), Type::ETVar(a.clone())),
                    right,
                );
                let output = format!("{}", new_ctx);
                Ok((
                    new_ctx,
                    InferenceTree::new("InstLReach", &input, &output, vec![]),
                ))
            }
            Type::Arrow(t1, t2) => {
                let a1 = self.fresh_tyvar();
                let a2 = self.fresh_tyvar();
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let arrow_type = Type::Arrow(
                    Box::new(Type::ETVar(a1.clone())),
                    Box::new(Type::ETVar(a2.clone())),
                );
                let mut new_ctx = left;
                new_ctx.push(Entry::SETVarBnd(a.clone(), arrow_type));
                new_ctx.push(Entry::ETVarBnd(a1.clone()));
                new_ctx.push(Entry::ETVarBnd(a2.clone()));
                new_ctx.extend(right);
                let ctx1 = Context(new_ctx);

                let (ctx2, tree1) = self.inst_r(&ctx1, t1, &a1)?;
                let t2_applied = self.apply_ctx_type(&ctx2, t2);
                let (ctx3, tree2) = self.inst_l(&ctx2, &a2, &t2_applied)?;

                let output = format!("{}", ctx3);
                Ok((
                    ctx3,
                    InferenceTree::new("InstLArr", &input, &output, vec![tree1, tree2]),
                ))
            }
            Type::Forall(b, t) => {
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::TVarBnd(b.clone()));
                let (ctx1, tree) = self.inst_l(&new_ctx, a, t)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == b));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("InstLAllR", &input, &output, vec![tree]),
                ))
            }
            _ if self.is_mono(ty) => {
                // Occurs check: ensure ^α doesn't occur in τ to prevent infinite types
                if self.occurs_check(a, ty) {
                    return Err(TypeError::OccursCheck {
                        var: a.clone(),
                        ty: ty.clone(),
                        expr: None,
                    });
                }
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let new_ctx =
                    Context::from_parts(left, Entry::SETVarBnd(a.clone(), ty.clone()), right);
                let output = format!("{}", new_ctx);
                Ok((
                    new_ctx,
                    InferenceTree::new("InstLSolve", &input, &output, vec![]),
                ))
            }
            _ => Err(TypeError::InstantiationError {
                var: a.clone(),
                ty: ty.clone(),
                expr: None,
            }),
        }
    }

    fn inst_r(
        &mut self,
        ctx: &Context,
        ty: &Type,
        a: &TyVar,
    ) -> TypeResult<(Context, InferenceTree)> {
        let input = format!("{} ⊢ {} :=< ^{}", ctx, ty, a);

        match ty {
            Type::ETVar(b) if self.before(ctx, a, b) => {
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == b));
                let new_ctx = Context::from_parts(
                    left,
                    Entry::SETVarBnd(b.clone(), Type::ETVar(a.clone())),
                    right,
                );
                let output = format!("{}", new_ctx);
                Ok((
                    new_ctx,
                    InferenceTree::new("InstRReach", &input, &output, vec![]),
                ))
            }
            Type::Arrow(t1, t2) => {
                let a1 = self.fresh_tyvar();
                let a2 = self.fresh_tyvar();
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let arrow_type = Type::Arrow(
                    Box::new(Type::ETVar(a1.clone())),
                    Box::new(Type::ETVar(a2.clone())),
                );
                let mut new_ctx = left;
                new_ctx.push(Entry::SETVarBnd(a.clone(), arrow_type));
                new_ctx.push(Entry::ETVarBnd(a1.clone()));
                new_ctx.push(Entry::ETVarBnd(a2.clone()));
                new_ctx.extend(right);
                let ctx1 = Context(new_ctx);

                let (ctx2, tree1) = self.inst_l(&ctx1, &a1, t1)?;
                let t2_applied = self.apply_ctx_type(&ctx2, t2);
                let (ctx3, tree2) = self.inst_r(&ctx2, &t2_applied, &a2)?;

                let output = format!("{}", ctx3);
                Ok((
                    ctx3,
                    InferenceTree::new("InstRArr", &input, &output, vec![tree1, tree2]),
                ))
            }
            Type::Forall(b, t) => {
                let subst_t = self.subst_type(b, &Type::ETVar(b.clone()), t);
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::ETVarBnd(b.clone()));
                new_ctx.push(Entry::Mark(b.clone()));
                let (ctx1, tree) = self.inst_r(&new_ctx, &subst_t, a)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::Mark(name) if name == b));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("InstRAllL", &input, &output, vec![tree]),
                ))
            }
            _ if self.is_mono(ty) => {
                // Occurs check: ensure ^α doesn't occur in τ to prevent infinite types
                if self.occurs_check(a, ty) {
                    return Err(TypeError::OccursCheck {
                        var: a.clone(),
                        ty: ty.clone(),
                        expr: None,
                    });
                }
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let new_ctx =
                    Context::from_parts(left, Entry::SETVarBnd(a.clone(), ty.clone()), right);
                let output = format!("{}", new_ctx);
                Ok((
                    new_ctx,
                    InferenceTree::new("InstRSolve", &input, &output, vec![]),
                ))
            }
            _ => Err(TypeError::InstantiationError {
                var: a.clone(),
                ty: ty.clone(),
                expr: None,
            }),
        }
    }

    fn subtype(
        &mut self,
        ctx: &Context,
        t1: &Type,
        t2: &Type,
    ) -> TypeResult<(Context, InferenceTree)> {
        let input = format!("{} ⊢ {} <: {}", ctx, t1, t2);

        match (t1, t2) {
            (Type::Int, Type::Int) | (Type::Bool, Type::Bool) => Ok((
                ctx.clone(),
                InferenceTree::new("SubRefl", &input, &format!("{}", ctx), vec![]),
            )),
            (Type::Var(a), Type::Var(b)) if a == b => Ok((
                ctx.clone(),
                InferenceTree::new("SubReflTVar", &input, &format!("{}", ctx), vec![]),
            )),
            (Type::ETVar(a), Type::ETVar(b)) if a == b => Ok((
                ctx.clone(),
                InferenceTree::new("SubReflETVar", &input, &format!("{}", ctx), vec![]),
            )),
            (Type::Arrow(a1, a2), Type::Arrow(b1, b2)) => {
                let (ctx1, tree1) = self.subtype(ctx, b1, a1)?; // Contravariant in argument
                let (ctx2, tree2) = self.subtype(&ctx1, a2, b2)?; // Covariant in result
                let output = format!("{}", ctx2);
                Ok((
                    ctx2,
                    InferenceTree::new("SubArr", &input, &output, vec![tree1, tree2]),
                ))
            }
            (_, Type::Forall(b, t2_body)) => {
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::TVarBnd(b.clone()));
                let (ctx1, tree) = self.subtype(&new_ctx, t1, t2_body)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == b));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("SubAllR", &input, &output, vec![tree]),
                ))
            }
            (Type::Forall(a, t1_body), _) => {
                let subst_t1 = self.subst_type(a, &Type::ETVar(a.clone()), t1_body);
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::ETVarBnd(a.clone()));
                new_ctx.push(Entry::Mark(a.clone()));
                let (ctx1, tree) = self.subtype(&new_ctx, &subst_t1, t2)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::Mark(name) if name == a));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("SubAllL", &input, &output, vec![tree]),
                ))
            }
            (Type::ETVar(a), _) if !self.free_vars(t2).contains(a) => {
                let (ctx1, tree) = self.inst_l(ctx, a, t2)?;
                let output = format!("{}", ctx1);
                Ok((
                    ctx1,
                    InferenceTree::new("SubInstL", &input, &output, vec![tree]),
                ))
            }
            (_, Type::ETVar(a)) if !self.free_vars(t1).contains(a) => {
                let (ctx1, tree) = self.inst_r(ctx, t1, a)?;
                let output = format!("{}", ctx1);
                Ok((
                    ctx1,
                    InferenceTree::new("SubInstR", &input, &output, vec![tree]),
                ))
            }
            _ => Err(TypeError::SubtypingError {
                left: t1.clone(),
                right: t2.clone(),
                expr: None,
            }),
        }
    }

    fn check(
        &mut self,
        ctx: &Context,
        expr: &Expr,
        ty: &Type,
    ) -> TypeResult<(Context, InferenceTree)> {
        let input = format!("{} ⊢ {:?} ⇐ {}", ctx, expr, ty);

        match (expr, ty) {
            (Expr::LitInt(_), Type::Int) => Ok((
                ctx.clone(),
                InferenceTree::new("ChkLitInt", &input, &format!("{}", ctx), vec![]),
            )),
            (Expr::LitBool(_), Type::Bool) => Ok((
                ctx.clone(),
                InferenceTree::new("ChkLitBool", &input, &format!("{}", ctx), vec![]),
            )),
            (Expr::Abs(x, _param_ty, body), Type::Arrow(expected_param, result_ty)) => {
                // For simplicity, we assume the parameter type matches
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::VarBnd(x.clone(), *expected_param.clone()));
                let (ctx1, tree) = self.check(&new_ctx, body, result_ty)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("ChkLam", &input, &output, vec![tree]),
                ))
            }
            (_, Type::Forall(a, ty_body)) => {
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::TVarBnd(a.clone()));
                let (ctx1, tree) = self.check(&new_ctx, expr, ty_body)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == a));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("ChkAll", &input, &output, vec![tree]),
                ))
            }
            _ => {
                // Fallback to inference + subtyping
                let (inferred_ty, ctx1, tree1) = self.infer(ctx, expr)?;
                let inferred_applied = self.apply_ctx_type(&ctx1, &inferred_ty);
                let ty_applied = self.apply_ctx_type(&ctx1, ty);
                let (ctx2, tree2) = self.subtype(&ctx1, &inferred_applied, &ty_applied)?;
                let output = format!("{}", ctx2);
                Ok((
                    ctx2,
                    InferenceTree::new("ChkSub", &input, &output, vec![tree1, tree2]),
                ))
            }
        }
    }

    pub fn infer(
        &mut self,
        ctx: &Context,
        expr: &Expr,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let input = format!("{} ⊢ {:?}", ctx, expr);

        match expr {
            Expr::Var(x) => self.infer_var(ctx, x, &input),
            Expr::Ann(expr, ty) => self.infer_ann(ctx, expr, ty, &input),
            Expr::LitInt(n) => self.infer_lit_int(ctx, *n, &input),
            Expr::LitBool(b) => self.infer_lit_bool(ctx, *b, &input),
            Expr::Abs(x, param_ty, body) => self.infer_abs(ctx, x, param_ty, body, &input),
            Expr::App(func, arg) => self.infer_application(ctx, func, arg, &input),
            Expr::TAbs(a, body) => self.infer_tabs(ctx, a, body, &input),
            Expr::TApp(func, ty_arg) => self.infer_tapp(ctx, func, ty_arg, &input),
            Expr::Let(x, e1, e2) => self.infer_let(ctx, x, e1, e2, &input),
            Expr::IfThenElse(e1, e2, e3) => self.infer_if(ctx, e1, e2, e3, &input),
            Expr::BinOp(op, e1, e2) => self.infer_binop(ctx, op, e1, e2, &input),
        }
    }

    /// T-Instr: Type annotation rule
    /// Γ ⊢ e ⇐ A
    /// ──────────────────
    /// Γ ⊢ (e : A) ⇒ A
    fn infer_ann(
        &mut self,
        ctx: &Context,
        expr: &Expr,
        ty: &Type,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (ctx1, tree) = self.check(ctx, expr, ty)?;
        let output = format!("{} ⇒ {} ⊣ {}", input, ty, ctx1);
        Ok((
            ty.clone(),
            ctx1,
            InferenceTree::new("InfAnn", input, &output, vec![tree]),
        ))
    }

    /// T-ForallI: Type abstraction rule
    /// Γ, α ⊢ e ⇒ A
    /// ──────────────────────
    /// Γ ⊢ Λα. e ⇒ ∀α. A
    fn infer_tabs(
        &mut self,
        ctx: &Context,
        a: &str,
        body: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let mut new_ctx = ctx.clone();
        new_ctx.push(Entry::TVarBnd(a.to_string()));
        let (body_ty, ctx1, tree) = self.infer(&new_ctx, body)?;

        // Apply context substitutions to resolve existential variables before removing
        // type binding
        let resolved_body_ty = self.apply_ctx_type(&ctx1, &body_ty);

        let (left, _, right) =
            ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == a));
        // Preserve solved existential variable bindings from the left context
        let mut final_ctx_entries = left
            .into_iter()
            .filter(|entry| matches!(entry, Entry::SETVarBnd(_, _)))
            .collect::<Vec<_>>();
        final_ctx_entries.extend(right);
        let final_ctx = Context(final_ctx_entries);
        let result_ty = Type::Forall(a.to_string(), Box::new(resolved_body_ty));
        let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, final_ctx);
        Ok((
            result_ty,
            final_ctx,
            InferenceTree::new("InfTAbs", input, &output, vec![tree]),
        ))
    }

    /// T-ForallE: Type application rule
    /// Γ ⊢ e ⇒ ∀α. A
    /// ──────────────────────
    /// Γ ⊢ e[B] ⇒ [B/α]A
    fn infer_tapp(
        &mut self,
        ctx: &Context,
        func: &Expr,
        ty_arg: &Type,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (func_ty, ctx1, tree1) = self.infer(ctx, func)?;
        match func_ty {
            Type::Forall(a, body_ty) => {
                let result_ty = self.subst_type(&a, ty_arg, &body_ty);
                let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, ctx1);
                Ok((
                    result_ty,
                    ctx1,
                    InferenceTree::new("InfTApp", input, &output, vec![tree1]),
                ))
            }
            _ => Err(TypeError::TypeApplicationError {
                actual: func_ty.clone(),
                expr: None,
            }),
        }
    }

    fn infer_app(
        &mut self,
        ctx: &Context,
        func_ty: &Type,
        arg: &Expr,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let input = format!("{} ⊢ {:?} • {}", ctx, arg, func_ty);

        match func_ty {
            Type::Arrow(param_ty, result_ty) => {
                let (ctx1, tree) = self.check(ctx, arg, param_ty)?;
                let output = format!("{} ⇒⇒ {} ⊣ {}", input, result_ty, ctx1);
                Ok((
                    result_ty.as_ref().clone(),
                    ctx1,
                    InferenceTree::new("InfAppArr", &input, &output, vec![tree]),
                ))
            }
            Type::ETVar(a) => {
                let a1 = self.fresh_tyvar();
                let a2 = self.fresh_tyvar();
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let arrow_type = Type::Arrow(
                    Box::new(Type::ETVar(a1.clone())),
                    Box::new(Type::ETVar(a2.clone())),
                );
                let mut new_ctx = left;
                new_ctx.push(Entry::SETVarBnd(a.clone(), arrow_type));
                new_ctx.push(Entry::ETVarBnd(a1.clone()));
                new_ctx.push(Entry::ETVarBnd(a2.clone()));
                new_ctx.extend(right);
                let ctx1 = Context(new_ctx);

                let (ctx2, tree) = self.check(&ctx1, arg, &Type::ETVar(a1))?;
                let output = format!("{} ⇒⇒ ^{} ⊣ {}", input, a2, ctx2);
                Ok((
                    Type::ETVar(a2),
                    ctx2,
                    InferenceTree::new("InfAppETVar", &input, &output, vec![tree]),
                ))
            }
            _ => Err(TypeError::ApplicationTypeError {
                actual: func_ty.clone(),
                expr: None,
            }),
        }
    }
}
pub fn run_bidirectional(expr: &Expr) -> TypeResult<(Type, Context, InferenceTree)> {
    let mut bi = BiDirectional::new();
    let ctx = Context::new();
    let (ty, final_ctx, tree) = bi.infer(&ctx, expr)?;
    // Apply the final context to resolve existential variables
    let resolved_ty = bi.apply_ctx_type(&final_ctx, &ty);
    Ok((resolved_ty, final_ctx, tree))
}
#[derive(Debug, Clone)]
pub enum Entry {
    VarBnd(TmVar, Type),    // x: A
    TVarBnd(TyVar),         // α
    ETVarBnd(TyVar),        // ^α
    SETVarBnd(TyVar, Type), // ^α = τ
    Mark(TyVar),            // $α
}
}

The context entries represent different kinds of information the type checker needs to track throughout the bidirectional inference process. Variable bindings, represented as VarBnd(TmVar, Type), associate term variables with their types, such as recording that x has type Int in the current scope. Type variable bindings, denoted as TVarBnd(TyVar), introduce type variables into scope, such as the \( \alpha \) that appears in universal quantification \( \forall \alpha. \ldots \).

Existential type variable bindings, written as ETVarBnd(TyVar), introduce existential type variables that represent unknown types to be determined through constraint solving. Solved existential type variable bindings, represented as SETVarBnd(TyVar, Type), record the concrete solutions discovered for existential variables during the inference process. Finally, marker entries, denoted as Mark(TyVar), mark the beginning of a scope to enable proper garbage collection of type variables when their scope is exited.

The context itself maintains these entries in a stack-like structure where order matters crucially for scoping and variable resolution:

#![allow(unused)]
fn main() {
use std::collections::HashSet;
use std::fmt;
use crate::ast::{BinOp, Expr, Type};
use crate::errors::{TypeError, TypeResult};
pub type TyVar = String;
pub type TmVar = String;
#[derive(Debug, Clone)]
pub enum Entry {
    VarBnd(TmVar, Type),    // x: A
    TVarBnd(TyVar),         // α
    ETVarBnd(TyVar),        // ^α
    SETVarBnd(TyVar, Type), // ^α = τ
    Mark(TyVar),            // $α
}
impl fmt::Display for Entry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Entry::VarBnd(x, ty) => write!(f, "{}: {}", x, ty),
            Entry::TVarBnd(a) => write!(f, "{}", a),
            Entry::ETVarBnd(a) => write!(f, "^{}", a),
            Entry::SETVarBnd(a, ty) => write!(f, "^{} = {}", a, ty),
            Entry::Mark(a) => write!(f, "${}", a),
        }
    }
}
impl Default for Context {
    fn default() -> Self {
        Self::new()
    }
}
impl Context {
    pub fn new() -> Self {
        Context(Vec::new())
    }

    pub fn push(&mut self, entry: Entry) {
        self.0.push(entry);
    }

    pub fn find<F>(&self, predicate: F) -> Option<&Entry>
    where
        F: Fn(&Entry) -> bool, {
        self.0.iter().find(|entry| predicate(entry))
    }

    pub fn break3<F>(&self, predicate: F) -> (Vec<Entry>, Option<Entry>, Vec<Entry>)
    where
        F: Fn(&Entry) -> bool, {
        if let Some(pos) = self.0.iter().position(predicate) {
            let left = self.0[..pos].to_vec();
            let middle = self.0[pos].clone();
            let right = self.0[pos + 1..].to_vec();
            (left, Some(middle), right)
        } else {
            (self.0.clone(), None, Vec::new())
        }
    }

    pub fn from_parts(left: Vec<Entry>, middle: Entry, right: Vec<Entry>) -> Self {
        let mut ctx = left;
        ctx.push(middle);
        ctx.extend(right);
        Context(ctx)
    }
}
impl fmt::Display for Context {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let entries: Vec<String> = self.0.iter().rev().map(|e| e.to_string()).collect();
        write!(f, "{}", entries.join(", "))
    }
}
#[derive(Debug)]
pub struct InferenceTree {
    pub rule: String,
    pub input: String,
    pub output: String,
    pub children: Vec<InferenceTree>,
}
impl InferenceTree {
    fn new(rule: &str, input: &str, output: &str, children: Vec<InferenceTree>) -> Self {
        Self {
            rule: rule.to_string(),
            input: input.to_string(),
            output: output.to_string(),
            children,
        }
    }
}
impl fmt::Display for InferenceTree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.display_with_indent(f, 0)
    }
}
impl InferenceTree {
    fn display_with_indent(&self, f: &mut fmt::Formatter, indent: usize) -> fmt::Result {
        let prefix = "  ".repeat(indent);
        writeln!(
            f,
            "{}{}: {} => {}",
            prefix, self.rule, self.input, self.output
        )?;
        for child in &self.children {
            child.display_with_indent(f, indent + 1)?;
        }
        Ok(())
    }
}
pub struct BiDirectional {
    counter: usize,
}
impl Default for BiDirectional {
    fn default() -> Self {
        Self::new()
    }
}
#[allow(clippy::only_used_in_recursion)]
impl BiDirectional {
    pub fn new() -> Self {
        Self { counter: 0 }
    }

    fn fresh_tyvar(&mut self) -> TyVar {
        let var = format!("α{}", self.counter);
        self.counter += 1;
        var
    }

    fn is_mono(&self, ty: &Type) -> bool {
        match ty {
            Type::Int | Type::Bool | Type::Var(_) | Type::ETVar(_) => true,
            Type::Arrow(t1, t2) => self.is_mono(t1) && self.is_mono(t2),
            Type::Forall(_, _) => false,
        }
    }

    fn free_vars(&self, ty: &Type) -> HashSet<TyVar> {
        match ty {
            Type::Var(name) | Type::ETVar(name) => {
                let mut set = HashSet::new();
                set.insert(name.clone());
                set
            }
            Type::Arrow(t1, t2) => {
                let mut set = self.free_vars(t1);
                set.extend(self.free_vars(t2));
                set
            }
            Type::Forall(var, ty) => {
                let mut set = self.free_vars(ty);
                set.remove(var);
                set
            }
            Type::Int | Type::Bool => HashSet::new(),
        }
    }

    fn subst_type(&self, var: &TyVar, replacement: &Type, ty: &Type) -> Type {
        match ty {
            Type::Var(name) if name == var => replacement.clone(),
            Type::ETVar(name) if name == var => replacement.clone(),
            Type::Var(_) | Type::ETVar(_) | Type::Int | Type::Bool => ty.clone(),
            Type::Arrow(t1, t2) => Type::Arrow(
                Box::new(self.subst_type(var, replacement, t1)),
                Box::new(self.subst_type(var, replacement, t2)),
            ),
            Type::Forall(bound_var, body) => {
                if bound_var == var {
                    ty.clone() // Variable is shadowed
                } else {
                    Type::Forall(
                        bound_var.clone(),
                        Box::new(self.subst_type(var, replacement, body)),
                    )
                }
            }
        }
    }

    pub fn apply_ctx_type(&self, ctx: &Context, ty: &Type) -> Type {
        let mut current = ty.clone();
        let mut changed = true;

        // Keep applying substitutions until no more changes occur
        while changed {
            changed = false;
            let new_type = self.apply_ctx_type_once(ctx, &current);
            if new_type != current {
                changed = true;
                current = new_type;
            }
        }

        current
    }

    fn apply_ctx_type_once(&self, ctx: &Context, ty: &Type) -> Type {
        match ty {
            Type::ETVar(a) => {
                if let Some(Entry::SETVarBnd(_, replacement)) =
                    ctx.find(|entry| matches!(entry, Entry::SETVarBnd(name, _) if name == a))
                {
                    // Recursively apply substitutions to the replacement type
                    self.apply_ctx_type_once(ctx, replacement)
                } else {
                    ty.clone()
                }
            }
            Type::Arrow(t1, t2) => Type::Arrow(
                Box::new(self.apply_ctx_type_once(ctx, t1)),
                Box::new(self.apply_ctx_type_once(ctx, t2)),
            ),
            Type::Forall(var, body) => {
                Type::Forall(var.clone(), Box::new(self.apply_ctx_type_once(ctx, body)))
            }
            _ => ty.clone(),
        }
    }

    fn before(&self, ctx: &Context, a: &TyVar, b: &TyVar) -> bool {
        let pos_a = ctx
            .0
            .iter()
            .position(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
        let pos_b = ctx
            .0
            .iter()
            .position(|entry| matches!(entry, Entry::ETVarBnd(name) if name == b));

        match (pos_a, pos_b) {
            (Some(pa), Some(pb)) => pa > pb, // Later in the context means earlier in ordering
            _ => false,
        }
    }

    /// Check if a type variable occurs in a type (occurs check for unification)
    /// This prevents creating infinite types when solving ^α := τ by ensuring α
    /// ∉ τ
    fn occurs_check(&self, var: &TyVar, ty: &Type) -> bool {
        match ty {
            Type::Var(name) | Type::ETVar(name) => name == var,
            Type::Arrow(t1, t2) => self.occurs_check(var, t1) || self.occurs_check(var, t2),
            Type::Forall(bound_var, body) => {
                // If the variable is shadowed by the forall binding, it doesn't occur
                if bound_var == var {
                    false
                } else {
                    self.occurs_check(var, body)
                }
            }
            Type::Int | Type::Bool => false,
        }
    }

    // ========== TYPING RULE METHODS ==========
    // Each method implements a specific typing rule from System F

    /// T-Var: Variable lookup rule
    /// Γ, x:A ⊢ x ⇒ A
    fn infer_var(
        &self,
        ctx: &Context,
        x: &str,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        if let Some(Entry::VarBnd(_, ty)) =
            ctx.find(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x))
        {
            let output = format!("{} ⇒ {} ⊣ {}", input, ty, ctx);
            Ok((
                ty.clone(),
                ctx.clone(),
                InferenceTree::new("InfVar", input, &output, vec![]),
            ))
        } else {
            Err(TypeError::UnboundVariable {
                name: x.to_string(),
                expr: None,
            })
        }
    }

    /// T-LitInt: Integer literal rule
    /// ⊢ n ⇒ Int
    fn infer_lit_int(
        &self,
        ctx: &Context,
        _n: i64,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let output = format!("{} ⇒ Int ⊣ {}", input, ctx);
        Ok((
            Type::Int,
            ctx.clone(),
            InferenceTree::new("InfLitInt", input, &output, vec![]),
        ))
    }

    /// T-LitBool: Boolean literal rule
    /// ⊢ b ⇒ Bool
    fn infer_lit_bool(
        &self,
        ctx: &Context,
        _b: bool,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let output = format!("{} ⇒ Bool ⊣ {}", input, ctx);
        Ok((
            Type::Bool,
            ctx.clone(),
            InferenceTree::new("InfLitBool", input, &output, vec![]),
        ))
    }

    /// T-Abs: Lambda abstraction rule
    /// Γ,x:A ⊢ e ⇐ B
    /// ─────────────────────
    /// Γ ⊢ λx:A.e ⇒ A → B
    fn infer_abs(
        &mut self,
        ctx: &Context,
        x: &str,
        param_ty: &Type,
        body: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let b = self.fresh_tyvar();
        let mut new_ctx = ctx.clone();
        new_ctx.push(Entry::VarBnd(x.to_string(), param_ty.clone()));
        new_ctx.push(Entry::ETVarBnd(b.clone()));

        let (ctx1, tree) = self.check(&new_ctx, body, &Type::ETVar(b.clone()))?;
        let (left, _, right) =
            ctx1.break3(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x));
        // Preserve solved existential variable bindings from the left context
        let mut final_ctx_entries = left
            .into_iter()
            .filter(|entry| matches!(entry, Entry::SETVarBnd(_, _)))
            .collect::<Vec<_>>();
        final_ctx_entries.extend(right);
        let final_ctx = Context(final_ctx_entries);
        let result_ty = Type::Arrow(Box::new(param_ty.clone()), Box::new(Type::ETVar(b)));
        let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, final_ctx);
        Ok((
            result_ty,
            final_ctx,
            InferenceTree::new("InfLam", input, &output, vec![tree]),
        ))
    }

    /// T-App: Function application rule
    /// Γ ⊢ e1 ⇒ A   Γ ⊢ e2 • A ⇒⇒ C
    /// ────────────────────────────────
    /// Γ ⊢ e1 e2 ⇒ C
    fn infer_application(
        &mut self,
        ctx: &Context,
        func: &Expr,
        arg: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (func_ty, ctx1, tree1) = self.infer(ctx, func)?;
        let func_ty_applied = self.apply_ctx_type(&ctx1, &func_ty);
        let (result_ty, ctx2, tree2) = self.infer_app(&ctx1, &func_ty_applied, arg)?;
        let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, ctx2);
        Ok((
            result_ty,
            ctx2,
            InferenceTree::new("InfApp", input, &output, vec![tree1, tree2]),
        ))
    }

    /// T-Let: Let binding rule
    /// Γ ⊢ e1 ⇒ A   Γ,x:A ⊢ e2 ⇒ B
    /// ─────────────────────────────
    /// Γ ⊢ let x = e1 in e2 ⇒ B
    fn infer_let(
        &mut self,
        ctx: &Context,
        x: &str,
        e1: &Expr,
        e2: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (ty1, ctx1, tree1) = self.infer(ctx, e1)?;
        let mut new_ctx = ctx1.clone();
        new_ctx.push(Entry::VarBnd(x.to_string(), ty1));
        let (ty2, ctx2, tree2) = self.infer(&new_ctx, e2)?;
        let (left, _, right) =
            ctx2.break3(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x));
        // Preserve solved existential variable bindings from the left context
        let mut final_ctx_entries = left
            .into_iter()
            .filter(|entry| matches!(entry, Entry::SETVarBnd(_, _)))
            .collect::<Vec<_>>();
        final_ctx_entries.extend(right);
        let final_ctx = Context(final_ctx_entries);
        let output = format!("{} ⇒ {} ⊣ {}", input, ty2, final_ctx);
        Ok((
            ty2,
            final_ctx,
            InferenceTree::new("InfLet", input, &output, vec![tree1, tree2]),
        ))
    }

    /// T-If: Conditional rule
    /// Γ ⊢ e1 ⇐ Bool   Γ ⊢ e2 ⇒ A   Γ ⊢ e3 ⇒ A
    /// ──────────────────────────────────────────
    /// Γ ⊢ if e1 then e2 else e3 ⇒ A
    fn infer_if(
        &mut self,
        ctx: &Context,
        e1: &Expr,
        e2: &Expr,
        e3: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (ctx1, tree1) = self.check(ctx, e1, &Type::Bool)?;
        let (ty2, ctx2, tree2) = self.infer(&ctx1, e2)?;
        let (ty3, ctx3, tree3) = self.infer(&ctx2, e3)?;
        // Ensure both branches have the same type
        let (unified_ctx, tree_unify) = self.subtype(&ctx3, &ty2, &ty3)?;
        let output = format!("{} ⇒ {} ⊣ {}", input, ty2, unified_ctx);
        Ok((
            ty2,
            unified_ctx,
            InferenceTree::new(
                "InfIf",
                input,
                &output,
                vec![tree1, tree2, tree3, tree_unify],
            ),
        ))
    }

    /// T-BinOp: Binary operation rules
    fn infer_binop(
        &mut self,
        ctx: &Context,
        op: &BinOp,
        e1: &Expr,
        e2: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        match op {
            // T-Arith: Int → Int → Int
            BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div => {
                let (ctx1, tree1) = self.check(ctx, e1, &Type::Int)?;
                let (ctx2, tree2) = self.check(&ctx1, e2, &Type::Int)?;
                let output = format!("{} ⇒ Int ⊣ {}", input, ctx2);
                Ok((
                    Type::Int,
                    ctx2,
                    InferenceTree::new("InfArith", input, &output, vec![tree1, tree2]),
                ))
            }
            // T-Bool: Bool → Bool → Bool
            BinOp::And | BinOp::Or => {
                let (ctx1, tree1) = self.check(ctx, e1, &Type::Bool)?;
                let (ctx2, tree2) = self.check(&ctx1, e2, &Type::Bool)?;
                let output = format!("{} ⇒ Bool ⊣ {}", input, ctx2);
                Ok((
                    Type::Bool,
                    ctx2,
                    InferenceTree::new("InfBool", input, &output, vec![tree1, tree2]),
                ))
            }
            // T-Cmp: Int → Int → Bool
            BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge => {
                let (ctx1, tree1) = self.check(ctx, e1, &Type::Int)?;
                let (ctx2, tree2) = self.check(&ctx1, e2, &Type::Int)?;
                let output = format!("{} ⇒ Bool ⊣ {}", input, ctx2);
                Ok((
                    Type::Bool,
                    ctx2,
                    InferenceTree::new("InfCmp", input, &output, vec![tree1, tree2]),
                ))
            }
            // T-Eq: ∀α. α → α → Bool
            BinOp::Eq | BinOp::Ne => {
                let (ty1, ctx1, tree1) = self.infer(ctx, e1)?;
                let (ctx2, tree2) = self.check(&ctx1, e2, &ty1)?;
                let output = format!("{} ⇒ Bool ⊣ {}", input, ctx2);
                Ok((
                    Type::Bool,
                    ctx2,
                    InferenceTree::new("InfEq", input, &output, vec![tree1, tree2]),
                ))
            }
        }
    }

    fn inst_l(
        &mut self,
        ctx: &Context,
        a: &TyVar,
        ty: &Type,
    ) -> TypeResult<(Context, InferenceTree)> {
        let input = format!("{} ⊢ ^{} :=< {}", ctx, a, ty);

        match ty {
            Type::ETVar(b) if self.before(ctx, a, b) => {
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == b));
                let new_ctx = Context::from_parts(
                    left,
                    Entry::SETVarBnd(b.clone(), Type::ETVar(a.clone())),
                    right,
                );
                let output = format!("{}", new_ctx);
                Ok((
                    new_ctx,
                    InferenceTree::new("InstLReach", &input, &output, vec![]),
                ))
            }
            Type::Arrow(t1, t2) => {
                let a1 = self.fresh_tyvar();
                let a2 = self.fresh_tyvar();
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let arrow_type = Type::Arrow(
                    Box::new(Type::ETVar(a1.clone())),
                    Box::new(Type::ETVar(a2.clone())),
                );
                let mut new_ctx = left;
                new_ctx.push(Entry::SETVarBnd(a.clone(), arrow_type));
                new_ctx.push(Entry::ETVarBnd(a1.clone()));
                new_ctx.push(Entry::ETVarBnd(a2.clone()));
                new_ctx.extend(right);
                let ctx1 = Context(new_ctx);

                let (ctx2, tree1) = self.inst_r(&ctx1, t1, &a1)?;
                let t2_applied = self.apply_ctx_type(&ctx2, t2);
                let (ctx3, tree2) = self.inst_l(&ctx2, &a2, &t2_applied)?;

                let output = format!("{}", ctx3);
                Ok((
                    ctx3,
                    InferenceTree::new("InstLArr", &input, &output, vec![tree1, tree2]),
                ))
            }
            Type::Forall(b, t) => {
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::TVarBnd(b.clone()));
                let (ctx1, tree) = self.inst_l(&new_ctx, a, t)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == b));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("InstLAllR", &input, &output, vec![tree]),
                ))
            }
            _ if self.is_mono(ty) => {
                // Occurs check: ensure ^α doesn't occur in τ to prevent infinite types
                if self.occurs_check(a, ty) {
                    return Err(TypeError::OccursCheck {
                        var: a.clone(),
                        ty: ty.clone(),
                        expr: None,
                    });
                }
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let new_ctx =
                    Context::from_parts(left, Entry::SETVarBnd(a.clone(), ty.clone()), right);
                let output = format!("{}", new_ctx);
                Ok((
                    new_ctx,
                    InferenceTree::new("InstLSolve", &input, &output, vec![]),
                ))
            }
            _ => Err(TypeError::InstantiationError {
                var: a.clone(),
                ty: ty.clone(),
                expr: None,
            }),
        }
    }

    fn inst_r(
        &mut self,
        ctx: &Context,
        ty: &Type,
        a: &TyVar,
    ) -> TypeResult<(Context, InferenceTree)> {
        let input = format!("{} ⊢ {} :=< ^{}", ctx, ty, a);

        match ty {
            Type::ETVar(b) if self.before(ctx, a, b) => {
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == b));
                let new_ctx = Context::from_parts(
                    left,
                    Entry::SETVarBnd(b.clone(), Type::ETVar(a.clone())),
                    right,
                );
                let output = format!("{}", new_ctx);
                Ok((
                    new_ctx,
                    InferenceTree::new("InstRReach", &input, &output, vec![]),
                ))
            }
            Type::Arrow(t1, t2) => {
                let a1 = self.fresh_tyvar();
                let a2 = self.fresh_tyvar();
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let arrow_type = Type::Arrow(
                    Box::new(Type::ETVar(a1.clone())),
                    Box::new(Type::ETVar(a2.clone())),
                );
                let mut new_ctx = left;
                new_ctx.push(Entry::SETVarBnd(a.clone(), arrow_type));
                new_ctx.push(Entry::ETVarBnd(a1.clone()));
                new_ctx.push(Entry::ETVarBnd(a2.clone()));
                new_ctx.extend(right);
                let ctx1 = Context(new_ctx);

                let (ctx2, tree1) = self.inst_l(&ctx1, &a1, t1)?;
                let t2_applied = self.apply_ctx_type(&ctx2, t2);
                let (ctx3, tree2) = self.inst_r(&ctx2, &t2_applied, &a2)?;

                let output = format!("{}", ctx3);
                Ok((
                    ctx3,
                    InferenceTree::new("InstRArr", &input, &output, vec![tree1, tree2]),
                ))
            }
            Type::Forall(b, t) => {
                let subst_t = self.subst_type(b, &Type::ETVar(b.clone()), t);
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::ETVarBnd(b.clone()));
                new_ctx.push(Entry::Mark(b.clone()));
                let (ctx1, tree) = self.inst_r(&new_ctx, &subst_t, a)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::Mark(name) if name == b));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("InstRAllL", &input, &output, vec![tree]),
                ))
            }
            _ if self.is_mono(ty) => {
                // Occurs check: ensure ^α doesn't occur in τ to prevent infinite types
                if self.occurs_check(a, ty) {
                    return Err(TypeError::OccursCheck {
                        var: a.clone(),
                        ty: ty.clone(),
                        expr: None,
                    });
                }
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let new_ctx =
                    Context::from_parts(left, Entry::SETVarBnd(a.clone(), ty.clone()), right);
                let output = format!("{}", new_ctx);
                Ok((
                    new_ctx,
                    InferenceTree::new("InstRSolve", &input, &output, vec![]),
                ))
            }
            _ => Err(TypeError::InstantiationError {
                var: a.clone(),
                ty: ty.clone(),
                expr: None,
            }),
        }
    }

    fn subtype(
        &mut self,
        ctx: &Context,
        t1: &Type,
        t2: &Type,
    ) -> TypeResult<(Context, InferenceTree)> {
        let input = format!("{} ⊢ {} <: {}", ctx, t1, t2);

        match (t1, t2) {
            (Type::Int, Type::Int) | (Type::Bool, Type::Bool) => Ok((
                ctx.clone(),
                InferenceTree::new("SubRefl", &input, &format!("{}", ctx), vec![]),
            )),
            (Type::Var(a), Type::Var(b)) if a == b => Ok((
                ctx.clone(),
                InferenceTree::new("SubReflTVar", &input, &format!("{}", ctx), vec![]),
            )),
            (Type::ETVar(a), Type::ETVar(b)) if a == b => Ok((
                ctx.clone(),
                InferenceTree::new("SubReflETVar", &input, &format!("{}", ctx), vec![]),
            )),
            (Type::Arrow(a1, a2), Type::Arrow(b1, b2)) => {
                let (ctx1, tree1) = self.subtype(ctx, b1, a1)?; // Contravariant in argument
                let (ctx2, tree2) = self.subtype(&ctx1, a2, b2)?; // Covariant in result
                let output = format!("{}", ctx2);
                Ok((
                    ctx2,
                    InferenceTree::new("SubArr", &input, &output, vec![tree1, tree2]),
                ))
            }
            (_, Type::Forall(b, t2_body)) => {
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::TVarBnd(b.clone()));
                let (ctx1, tree) = self.subtype(&new_ctx, t1, t2_body)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == b));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("SubAllR", &input, &output, vec![tree]),
                ))
            }
            (Type::Forall(a, t1_body), _) => {
                let subst_t1 = self.subst_type(a, &Type::ETVar(a.clone()), t1_body);
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::ETVarBnd(a.clone()));
                new_ctx.push(Entry::Mark(a.clone()));
                let (ctx1, tree) = self.subtype(&new_ctx, &subst_t1, t2)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::Mark(name) if name == a));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("SubAllL", &input, &output, vec![tree]),
                ))
            }
            (Type::ETVar(a), _) if !self.free_vars(t2).contains(a) => {
                let (ctx1, tree) = self.inst_l(ctx, a, t2)?;
                let output = format!("{}", ctx1);
                Ok((
                    ctx1,
                    InferenceTree::new("SubInstL", &input, &output, vec![tree]),
                ))
            }
            (_, Type::ETVar(a)) if !self.free_vars(t1).contains(a) => {
                let (ctx1, tree) = self.inst_r(ctx, t1, a)?;
                let output = format!("{}", ctx1);
                Ok((
                    ctx1,
                    InferenceTree::new("SubInstR", &input, &output, vec![tree]),
                ))
            }
            _ => Err(TypeError::SubtypingError {
                left: t1.clone(),
                right: t2.clone(),
                expr: None,
            }),
        }
    }

    fn check(
        &mut self,
        ctx: &Context,
        expr: &Expr,
        ty: &Type,
    ) -> TypeResult<(Context, InferenceTree)> {
        let input = format!("{} ⊢ {:?} ⇐ {}", ctx, expr, ty);

        match (expr, ty) {
            (Expr::LitInt(_), Type::Int) => Ok((
                ctx.clone(),
                InferenceTree::new("ChkLitInt", &input, &format!("{}", ctx), vec![]),
            )),
            (Expr::LitBool(_), Type::Bool) => Ok((
                ctx.clone(),
                InferenceTree::new("ChkLitBool", &input, &format!("{}", ctx), vec![]),
            )),
            (Expr::Abs(x, _param_ty, body), Type::Arrow(expected_param, result_ty)) => {
                // For simplicity, we assume the parameter type matches
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::VarBnd(x.clone(), *expected_param.clone()));
                let (ctx1, tree) = self.check(&new_ctx, body, result_ty)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("ChkLam", &input, &output, vec![tree]),
                ))
            }
            (_, Type::Forall(a, ty_body)) => {
                let mut new_ctx = ctx.clone();
                new_ctx.push(Entry::TVarBnd(a.clone()));
                let (ctx1, tree) = self.check(&new_ctx, expr, ty_body)?;
                let (_, _, right) =
                    ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == a));
                let final_ctx = Context(right);
                let output = format!("{}", final_ctx);
                Ok((
                    final_ctx,
                    InferenceTree::new("ChkAll", &input, &output, vec![tree]),
                ))
            }
            _ => {
                // Fallback to inference + subtyping
                let (inferred_ty, ctx1, tree1) = self.infer(ctx, expr)?;
                let inferred_applied = self.apply_ctx_type(&ctx1, &inferred_ty);
                let ty_applied = self.apply_ctx_type(&ctx1, ty);
                let (ctx2, tree2) = self.subtype(&ctx1, &inferred_applied, &ty_applied)?;
                let output = format!("{}", ctx2);
                Ok((
                    ctx2,
                    InferenceTree::new("ChkSub", &input, &output, vec![tree1, tree2]),
                ))
            }
        }
    }

    pub fn infer(
        &mut self,
        ctx: &Context,
        expr: &Expr,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let input = format!("{} ⊢ {:?}", ctx, expr);

        match expr {
            Expr::Var(x) => self.infer_var(ctx, x, &input),
            Expr::Ann(expr, ty) => self.infer_ann(ctx, expr, ty, &input),
            Expr::LitInt(n) => self.infer_lit_int(ctx, *n, &input),
            Expr::LitBool(b) => self.infer_lit_bool(ctx, *b, &input),
            Expr::Abs(x, param_ty, body) => self.infer_abs(ctx, x, param_ty, body, &input),
            Expr::App(func, arg) => self.infer_application(ctx, func, arg, &input),
            Expr::TAbs(a, body) => self.infer_tabs(ctx, a, body, &input),
            Expr::TApp(func, ty_arg) => self.infer_tapp(ctx, func, ty_arg, &input),
            Expr::Let(x, e1, e2) => self.infer_let(ctx, x, e1, e2, &input),
            Expr::IfThenElse(e1, e2, e3) => self.infer_if(ctx, e1, e2, e3, &input),
            Expr::BinOp(op, e1, e2) => self.infer_binop(ctx, op, e1, e2, &input),
        }
    }

    /// T-Instr: Type annotation rule
    /// Γ ⊢ e ⇐ A
    /// ──────────────────
    /// Γ ⊢ (e : A) ⇒ A
    fn infer_ann(
        &mut self,
        ctx: &Context,
        expr: &Expr,
        ty: &Type,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (ctx1, tree) = self.check(ctx, expr, ty)?;
        let output = format!("{} ⇒ {} ⊣ {}", input, ty, ctx1);
        Ok((
            ty.clone(),
            ctx1,
            InferenceTree::new("InfAnn", input, &output, vec![tree]),
        ))
    }

    /// T-ForallI: Type abstraction rule
    /// Γ, α ⊢ e ⇒ A
    /// ──────────────────────
    /// Γ ⊢ Λα. e ⇒ ∀α. A
    fn infer_tabs(
        &mut self,
        ctx: &Context,
        a: &str,
        body: &Expr,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let mut new_ctx = ctx.clone();
        new_ctx.push(Entry::TVarBnd(a.to_string()));
        let (body_ty, ctx1, tree) = self.infer(&new_ctx, body)?;

        // Apply context substitutions to resolve existential variables before removing
        // type binding
        let resolved_body_ty = self.apply_ctx_type(&ctx1, &body_ty);

        let (left, _, right) =
            ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == a));
        // Preserve solved existential variable bindings from the left context
        let mut final_ctx_entries = left
            .into_iter()
            .filter(|entry| matches!(entry, Entry::SETVarBnd(_, _)))
            .collect::<Vec<_>>();
        final_ctx_entries.extend(right);
        let final_ctx = Context(final_ctx_entries);
        let result_ty = Type::Forall(a.to_string(), Box::new(resolved_body_ty));
        let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, final_ctx);
        Ok((
            result_ty,
            final_ctx,
            InferenceTree::new("InfTAbs", input, &output, vec![tree]),
        ))
    }

    /// T-ForallE: Type application rule
    /// Γ ⊢ e ⇒ ∀α. A
    /// ──────────────────────
    /// Γ ⊢ e[B] ⇒ [B/α]A
    fn infer_tapp(
        &mut self,
        ctx: &Context,
        func: &Expr,
        ty_arg: &Type,
        input: &str,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let (func_ty, ctx1, tree1) = self.infer(ctx, func)?;
        match func_ty {
            Type::Forall(a, body_ty) => {
                let result_ty = self.subst_type(&a, ty_arg, &body_ty);
                let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, ctx1);
                Ok((
                    result_ty,
                    ctx1,
                    InferenceTree::new("InfTApp", input, &output, vec![tree1]),
                ))
            }
            _ => Err(TypeError::TypeApplicationError {
                actual: func_ty.clone(),
                expr: None,
            }),
        }
    }

    fn infer_app(
        &mut self,
        ctx: &Context,
        func_ty: &Type,
        arg: &Expr,
    ) -> TypeResult<(Type, Context, InferenceTree)> {
        let input = format!("{} ⊢ {:?} • {}", ctx, arg, func_ty);

        match func_ty {
            Type::Arrow(param_ty, result_ty) => {
                let (ctx1, tree) = self.check(ctx, arg, param_ty)?;
                let output = format!("{} ⇒⇒ {} ⊣ {}", input, result_ty, ctx1);
                Ok((
                    result_ty.as_ref().clone(),
                    ctx1,
                    InferenceTree::new("InfAppArr", &input, &output, vec![tree]),
                ))
            }
            Type::ETVar(a) => {
                let a1 = self.fresh_tyvar();
                let a2 = self.fresh_tyvar();
                let (left, _, right) =
                    ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
                let arrow_type = Type::Arrow(
                    Box::new(Type::ETVar(a1.clone())),
                    Box::new(Type::ETVar(a2.clone())),
                );
                let mut new_ctx = left;
                new_ctx.push(Entry::SETVarBnd(a.clone(), arrow_type));
                new_ctx.push(Entry::ETVarBnd(a1.clone()));
                new_ctx.push(Entry::ETVarBnd(a2.clone()));
                new_ctx.extend(right);
                let ctx1 = Context(new_ctx);

                let (ctx2, tree) = self.check(&ctx1, arg, &Type::ETVar(a1))?;
                let output = format!("{} ⇒⇒ ^{} ⊣ {}", input, a2, ctx2);
                Ok((
                    Type::ETVar(a2),
                    ctx2,
                    InferenceTree::new("InfAppETVar", &input, &output, vec![tree]),
                ))
            }
            _ => Err(TypeError::ApplicationTypeError {
                actual: func_ty.clone(),
                expr: None,
            }),
        }
    }
}
pub fn run_bidirectional(expr: &Expr) -> TypeResult<(Type, Context, InferenceTree)> {
    let mut bi = BiDirectional::new();
    let ctx = Context::new();
    let (ty, final_ctx, tree) = bi.infer(&ctx, expr)?;
    // Apply the final context to resolve existential variables
    let resolved_ty = bi.apply_ctx_type(&final_ctx, &ty);
    Ok((resolved_ty, final_ctx, tree))
}
#[derive(Debug, Clone)]
pub struct Context(Vec<Entry>);
}

Context management requires careful attention to scoping rules. When we enter a polymorphic function or introduce new type variables, we must ensure they are properly cleaned up when we exit their scope. The context provides methods for breaking it apart and reconstructing it to handle these operations efficiently.

Type Substitution and Application

Type substitution forms the computational heart of our System F implementation. When we instantiate a polymorphic type or solve existential variables, we need to systematically replace type variables with concrete types throughout complex type expressions.

#![allow(unused)]
fn main() {
fn subst_type(&self, var: &TyVar, replacement: &Type, ty: &Type) -> Type {
    match ty {
        Type::Var(name) if name == var => replacement.clone(),
        Type::ETVar(name) if name == var => replacement.clone(),
        Type::Var(_) | Type::ETVar(_) | Type::Int | Type::Bool => ty.clone(),
        Type::Arrow(t1, t2) => Type::Arrow(
            Box::new(self.subst_type(var, replacement, t1)),
            Box::new(self.subst_type(var, replacement, t2)),
        ),
        Type::Forall(bound_var, body) => {
            if bound_var == var {
                ty.clone() // Variable is shadowed
            } else {
                Type::Forall(
                    bound_var.clone(),
                    Box::new(self.subst_type(var, replacement, body)),
                )
            }
        }
    }
}
}

Substitution must handle variable capture correctly. When substituting into a Forall type, we must ensure that the bound variable doesn’t conflict with the replacement type. This is analogous to alpha-conversion in lambda calculus but operates at the type level.

Context application extends substitution by applying the current state of existential variable solutions to a type:

#![allow(unused)]
fn main() {
pub fn apply_ctx_type(&self, ctx: &Context, ty: &Type) -> Type {
    let mut current = ty.clone();
    let mut changed = true;

    // Keep applying substitutions until no more changes occur
    while changed {
        changed = false;
        let new_type = self.apply_ctx_type_once(ctx, &current);
        if new_type != current {
            changed = true;
            current = new_type;
        }
    }

    current
}
}

This operation is essential because our algorithm incrementally builds up solutions to existential variables. As we learn more about unknown types, we need to propagate this information through all the types we’re working with.

Bidirectional Algorithm Core

Inference Rules

The inference mode synthesizes types from expressions. Our implementation uses a modular approach where the main inference function delegates to specialized methods for each syntactic form.

#![allow(unused)]
fn main() {
pub fn infer(
    &mut self,
    ctx: &Context,
    expr: &Expr,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let input = format!("{} ⊢ {:?}", ctx, expr);

    match expr {
        Expr::Var(x) => self.infer_var(ctx, x, &input),
        Expr::Ann(expr, ty) => self.infer_ann(ctx, expr, ty, &input),
        Expr::LitInt(n) => self.infer_lit_int(ctx, *n, &input),
        Expr::LitBool(b) => self.infer_lit_bool(ctx, *b, &input),
        Expr::Abs(x, param_ty, body) => self.infer_abs(ctx, x, param_ty, body, &input),
        Expr::App(func, arg) => self.infer_application(ctx, func, arg, &input),
        Expr::TAbs(a, body) => self.infer_tabs(ctx, a, body, &input),
        Expr::TApp(func, ty_arg) => self.infer_tapp(ctx, func, ty_arg, &input),
        Expr::Let(x, e1, e2) => self.infer_let(ctx, x, e1, e2, &input),
        Expr::IfThenElse(e1, e2, e3) => self.infer_if(ctx, e1, e2, e3, &input),
        Expr::BinOp(op, e1, e2) => self.infer_binop(ctx, op, e1, e2, &input),
    }
}
}

The inference function delegates to specialized methods that implement individual typing rules:

Variable Lookup follows the T-Var rule:

\[ \frac{x : A \in \Gamma}{\Gamma \vdash x \Rightarrow A} \text{(T-Var)} \]

#![allow(unused)]
fn main() {
/// T-Var: Variable lookup rule
/// Γ, x:A ⊢ x ⇒ A
fn infer_var(
    &self,
    ctx: &Context,
    x: &str,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    if let Some(Entry::VarBnd(_, ty)) =
        ctx.find(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x))
    {
        let output = format!("{} ⇒ {} ⊣ {}", input, ty, ctx);
        Ok((
            ty.clone(),
            ctx.clone(),
            InferenceTree::new("InfVar", input, &output, vec![]),
        ))
    } else {
        Err(TypeError::UnboundVariable {
            name: x.to_string(),
            expr: None,
        })
    }
}
}

Integer Literals use the T-LitInt rule:

\[ \frac{}{\Gamma \vdash n \Rightarrow \text{Int}} \text{(T-LitInt)} \]

#![allow(unused)]
fn main() {
/// T-LitInt: Integer literal rule
/// ⊢ n ⇒ Int
fn infer_lit_int(
    &self,
    ctx: &Context,
    _n: i64,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let output = format!("{} ⇒ Int ⊣ {}", input, ctx);
    Ok((
        Type::Int,
        ctx.clone(),
        InferenceTree::new("InfLitInt", input, &output, vec![]),
    ))
}
}

Boolean Literals use the T-LitBool rule:

\[ \frac{}{\Gamma \vdash \text{true} \Rightarrow \text{Bool}} \text{(T-LitBool)} \]

#![allow(unused)]
fn main() {
/// T-LitBool: Boolean literal rule
/// ⊢ b ⇒ Bool
fn infer_lit_bool(
    &self,
    ctx: &Context,
    _b: bool,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let output = format!("{} ⇒ Bool ⊣ {}", input, ctx);
    Ok((
        Type::Bool,
        ctx.clone(),
        InferenceTree::new("InfLitBool", input, &output, vec![]),
    ))
}
}

Lambda Abstraction implements the T-Abs rule:

\[ \frac{\Gamma, x : A \vdash e \Leftarrow B}{\Gamma \vdash \lambda x. e \Leftarrow A \to B} \text{(T-Abs)} \]

#![allow(unused)]
fn main() {
/// T-Abs: Lambda abstraction rule
/// Γ,x:A ⊢ e ⇐ B
/// ─────────────────────
/// Γ ⊢ λx:A.e ⇒ A → B
fn infer_abs(
    &mut self,
    ctx: &Context,
    x: &str,
    param_ty: &Type,
    body: &Expr,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let b = self.fresh_tyvar();
    let mut new_ctx = ctx.clone();
    new_ctx.push(Entry::VarBnd(x.to_string(), param_ty.clone()));
    new_ctx.push(Entry::ETVarBnd(b.clone()));

    let (ctx1, tree) = self.check(&new_ctx, body, &Type::ETVar(b.clone()))?;
    let (left, _, right) =
        ctx1.break3(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x));
    // Preserve solved existential variable bindings from the left context
    let mut final_ctx_entries = left
        .into_iter()
        .filter(|entry| matches!(entry, Entry::SETVarBnd(_, _)))
        .collect::<Vec<_>>();
    final_ctx_entries.extend(right);
    let final_ctx = Context(final_ctx_entries);
    let result_ty = Type::Arrow(Box::new(param_ty.clone()), Box::new(Type::ETVar(b)));
    let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, final_ctx);
    Ok((
        result_ty,
        final_ctx,
        InferenceTree::new("InfLam", input, &output, vec![tree]),
    ))
}
}

Function Application uses the T-App rule:

\[ \frac{\Gamma \vdash e_1 \Rightarrow A \to B \quad \Gamma \vdash e_2 \Leftarrow A}{\Gamma \vdash e_1 ; e_2 \Rightarrow B} \text{(T-App)} \]

#![allow(unused)]
fn main() {
/// T-App: Function application rule
/// Γ ⊢ e1 ⇒ A   Γ ⊢ e2 • A ⇒⇒ C
/// ────────────────────────────────
/// Γ ⊢ e1 e2 ⇒ C
fn infer_application(
    &mut self,
    ctx: &Context,
    func: &Expr,
    arg: &Expr,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let (func_ty, ctx1, tree1) = self.infer(ctx, func)?;
    let func_ty_applied = self.apply_ctx_type(&ctx1, &func_ty);
    let (result_ty, ctx2, tree2) = self.infer_app(&ctx1, &func_ty_applied, arg)?;
    let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, ctx2);
    Ok((
        result_ty,
        ctx2,
        InferenceTree::new("InfApp", input, &output, vec![tree1, tree2]),
    ))
}
}

Let Bindings implement the T-Let rule:

\[ \frac{\Gamma \vdash e_1 \Rightarrow A \quad \Gamma, x : A \vdash e_2 \Rightarrow B}{\Gamma \vdash \text{let } x = e_1 \text{ in } e_2 \Rightarrow B} \text{(T-Let)} \]

#![allow(unused)]
fn main() {
/// T-Let: Let binding rule
/// Γ ⊢ e1 ⇒ A   Γ,x:A ⊢ e2 ⇒ B
/// ─────────────────────────────
/// Γ ⊢ let x = e1 in e2 ⇒ B
fn infer_let(
    &mut self,
    ctx: &Context,
    x: &str,
    e1: &Expr,
    e2: &Expr,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let (ty1, ctx1, tree1) = self.infer(ctx, e1)?;
    let mut new_ctx = ctx1.clone();
    new_ctx.push(Entry::VarBnd(x.to_string(), ty1));
    let (ty2, ctx2, tree2) = self.infer(&new_ctx, e2)?;
    let (left, _, right) =
        ctx2.break3(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x));
    // Preserve solved existential variable bindings from the left context
    let mut final_ctx_entries = left
        .into_iter()
        .filter(|entry| matches!(entry, Entry::SETVarBnd(_, _)))
        .collect::<Vec<_>>();
    final_ctx_entries.extend(right);
    let final_ctx = Context(final_ctx_entries);
    let output = format!("{} ⇒ {} ⊣ {}", input, ty2, final_ctx);
    Ok((
        ty2,
        final_ctx,
        InferenceTree::new("InfLet", input, &output, vec![tree1, tree2]),
    ))
}
}

Conditional Expressions use the T-If rule:

\[ \frac{\Gamma \vdash e_1 \Leftarrow \text{Bool} \quad \Gamma \vdash e_2 \Rightarrow A \quad \Gamma \vdash e_3 \Leftarrow A}{\Gamma \vdash \text{if } e_1 \text{ then } e_2 \text{ else } e_3 \Rightarrow A} \text{(T-If)} \]

#![allow(unused)]
fn main() {
/// T-If: Conditional rule
/// Γ ⊢ e1 ⇐ Bool   Γ ⊢ e2 ⇒ A   Γ ⊢ e3 ⇒ A
/// ──────────────────────────────────────────
/// Γ ⊢ if e1 then e2 else e3 ⇒ A
fn infer_if(
    &mut self,
    ctx: &Context,
    e1: &Expr,
    e2: &Expr,
    e3: &Expr,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let (ctx1, tree1) = self.check(ctx, e1, &Type::Bool)?;
    let (ty2, ctx2, tree2) = self.infer(&ctx1, e2)?;
    let (ty3, ctx3, tree3) = self.infer(&ctx2, e3)?;
    // Ensure both branches have the same type
    let (unified_ctx, tree_unify) = self.subtype(&ctx3, &ty2, &ty3)?;
    let output = format!("{} ⇒ {} ⊣ {}", input, ty2, unified_ctx);
    Ok((
        ty2,
        unified_ctx,
        InferenceTree::new(
            "InfIf",
            input,
            &output,
            vec![tree1, tree2, tree3, tree_unify],
        ),
    ))
}
}

Binary Operations implement the T-BinOp rules (T-Arith, T-Bool, T-Cmp, T-Eq):

\[ \frac{\Gamma \vdash e_1 \Leftarrow \text{Int} \quad \Gamma \vdash e_2 \Leftarrow \text{Int}}{\Gamma \vdash e_1 \oplus e_2 \Rightarrow \text{Int}} \text{(T-Arith)} \]

\[ \frac{\Gamma \vdash e_1 \Leftarrow \text{Bool} \quad \Gamma \vdash e_2 \Leftarrow \text{Bool}}{\Gamma \vdash e_1 \land e_2 \Rightarrow \text{Bool}} \text{(T-Bool)} \]

\[ \frac{\Gamma \vdash e_1 \Leftarrow \text{Int} \quad \Gamma \vdash e_2 \Leftarrow \text{Int}}{\Gamma \vdash e_1 < e_2 \Rightarrow \text{Bool}} \text{(T-Cmp)} \]

\[ \frac{\Gamma \vdash e_1 \Rightarrow A \quad \Gamma \vdash e_2 \Leftarrow A}{\Gamma \vdash e_1 = e_2 \Rightarrow \text{Bool}} \text{(T-Eq)} \]

#![allow(unused)]
fn main() {
/// T-BinOp: Binary operation rules
fn infer_binop(
    &mut self,
    ctx: &Context,
    op: &BinOp,
    e1: &Expr,
    e2: &Expr,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    match op {
        // T-Arith: Int → Int → Int
        BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div => {
            let (ctx1, tree1) = self.check(ctx, e1, &Type::Int)?;
            let (ctx2, tree2) = self.check(&ctx1, e2, &Type::Int)?;
            let output = format!("{} ⇒ Int ⊣ {}", input, ctx2);
            Ok((
                Type::Int,
                ctx2,
                InferenceTree::new("InfArith", input, &output, vec![tree1, tree2]),
            ))
        }
        // T-Bool: Bool → Bool → Bool
        BinOp::And | BinOp::Or => {
            let (ctx1, tree1) = self.check(ctx, e1, &Type::Bool)?;
            let (ctx2, tree2) = self.check(&ctx1, e2, &Type::Bool)?;
            let output = format!("{} ⇒ Bool ⊣ {}", input, ctx2);
            Ok((
                Type::Bool,
                ctx2,
                InferenceTree::new("InfBool", input, &output, vec![tree1, tree2]),
            ))
        }
        // T-Cmp: Int → Int → Bool
        BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge => {
            let (ctx1, tree1) = self.check(ctx, e1, &Type::Int)?;
            let (ctx2, tree2) = self.check(&ctx1, e2, &Type::Int)?;
            let output = format!("{} ⇒ Bool ⊣ {}", input, ctx2);
            Ok((
                Type::Bool,
                ctx2,
                InferenceTree::new("InfCmp", input, &output, vec![tree1, tree2]),
            ))
        }
        // T-Eq: ∀α. α → α → Bool
        BinOp::Eq | BinOp::Ne => {
            let (ty1, ctx1, tree1) = self.infer(ctx, e1)?;
            let (ctx2, tree2) = self.check(&ctx1, e2, &ty1)?;
            let output = format!("{} ⇒ Bool ⊣ {}", input, ctx2);
            Ok((
                Type::Bool,
                ctx2,
                InferenceTree::new("InfEq", input, &output, vec![tree1, tree2]),
            ))
        }
    }
}
}

Type Annotations use the T-Instr rule:

\[ \frac{\Gamma \vdash e \Leftarrow A}{\Gamma \vdash (e : A) \Rightarrow A} \text{(T-Instr)} \]

#![allow(unused)]
fn main() {
/// T-Instr: Type annotation rule
/// Γ ⊢ e ⇐ A
/// ──────────────────
/// Γ ⊢ (e : A) ⇒ A
fn infer_ann(
    &mut self,
    ctx: &Context,
    expr: &Expr,
    ty: &Type,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let (ctx1, tree) = self.check(ctx, expr, ty)?;
    let output = format!("{} ⇒ {} ⊣ {}", input, ty, ctx1);
    Ok((
        ty.clone(),
        ctx1,
        InferenceTree::new("InfAnn", input, &output, vec![tree]),
    ))
}
}

Type Abstraction implements the T-TAbs rule:

\[ \frac{\Gamma, \alpha \vdash e \Rightarrow A}{\Gamma \vdash \Lambda \alpha. e \Rightarrow \forall \alpha. A} \text{(T-TAbs)} \]

#![allow(unused)]
fn main() {
/// T-ForallI: Type abstraction rule
/// Γ, α ⊢ e ⇒ A
/// ──────────────────────
/// Γ ⊢ Λα. e ⇒ ∀α. A
fn infer_tabs(
    &mut self,
    ctx: &Context,
    a: &str,
    body: &Expr,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let mut new_ctx = ctx.clone();
    new_ctx.push(Entry::TVarBnd(a.to_string()));
    let (body_ty, ctx1, tree) = self.infer(&new_ctx, body)?;

    // Apply context substitutions to resolve existential variables before removing
    // type binding
    let resolved_body_ty = self.apply_ctx_type(&ctx1, &body_ty);

    let (left, _, right) =
        ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == a));
    // Preserve solved existential variable bindings from the left context
    let mut final_ctx_entries = left
        .into_iter()
        .filter(|entry| matches!(entry, Entry::SETVarBnd(_, _)))
        .collect::<Vec<_>>();
    final_ctx_entries.extend(right);
    let final_ctx = Context(final_ctx_entries);
    let result_ty = Type::Forall(a.to_string(), Box::new(resolved_body_ty));
    let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, final_ctx);
    Ok((
        result_ty,
        final_ctx,
        InferenceTree::new("InfTAbs", input, &output, vec![tree]),
    ))
}
}

Type Application uses the T-TApp rule:

\[ \frac{\Gamma \vdash e \Rightarrow \forall \alpha. A}{\Gamma \vdash e[B] \Rightarrow [B/\alpha]A} \text{(T-TApp)} \]

#![allow(unused)]
fn main() {
/// T-ForallE: Type application rule
/// Γ ⊢ e ⇒ ∀α. A
/// ──────────────────────
/// Γ ⊢ e[B] ⇒ [B/α]A
fn infer_tapp(
    &mut self,
    ctx: &Context,
    func: &Expr,
    ty_arg: &Type,
    input: &str,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let (func_ty, ctx1, tree1) = self.infer(ctx, func)?;
    match func_ty {
        Type::Forall(a, body_ty) => {
            let result_ty = self.subst_type(&a, ty_arg, &body_ty);
            let output = format!("{} ⇒ {} ⊣ {}", input, result_ty, ctx1);
            Ok((
                result_ty,
                ctx1,
                InferenceTree::new("InfTApp", input, &output, vec![tree1]),
            ))
        }
        _ => Err(TypeError::TypeApplicationError {
            actual: func_ty.clone(),
            expr: None,
        }),
    }
}
}

Each method includes comments linking it to the formal typing rule it implements, making the correspondence between theory and implementation explicit.

Checking Rules

The checking mode verifies that expressions conform to expected types. This is where the algorithm can make progress even when types aren’t fully determined.

#![allow(unused)]
fn main() {
fn check(
    &mut self,
    ctx: &Context,
    expr: &Expr,
    ty: &Type,
) -> TypeResult<(Context, InferenceTree)> {
    let input = format!("{} ⊢ {:?} ⇐ {}", ctx, expr, ty);

    match (expr, ty) {
        (Expr::LitInt(_), Type::Int) => Ok((
            ctx.clone(),
            InferenceTree::new("ChkLitInt", &input, &format!("{}", ctx), vec![]),
        )),
        (Expr::LitBool(_), Type::Bool) => Ok((
            ctx.clone(),
            InferenceTree::new("ChkLitBool", &input, &format!("{}", ctx), vec![]),
        )),
        (Expr::Abs(x, _param_ty, body), Type::Arrow(expected_param, result_ty)) => {
            // For simplicity, we assume the parameter type matches
            let mut new_ctx = ctx.clone();
            new_ctx.push(Entry::VarBnd(x.clone(), *expected_param.clone()));
            let (ctx1, tree) = self.check(&new_ctx, body, result_ty)?;
            let (_, _, right) =
                ctx1.break3(|entry| matches!(entry, Entry::VarBnd(name, _) if name == x));
            let final_ctx = Context(right);
            let output = format!("{}", final_ctx);
            Ok((
                final_ctx,
                InferenceTree::new("ChkLam", &input, &output, vec![tree]),
            ))
        }
        (_, Type::Forall(a, ty_body)) => {
            let mut new_ctx = ctx.clone();
            new_ctx.push(Entry::TVarBnd(a.clone()));
            let (ctx1, tree) = self.check(&new_ctx, expr, ty_body)?;
            let (_, _, right) =
                ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == a));
            let final_ctx = Context(right);
            let output = format!("{}", final_ctx);
            Ok((
                final_ctx,
                InferenceTree::new("ChkAll", &input, &output, vec![tree]),
            ))
        }
        _ => {
            // Fallback to inference + subtyping
            let (inferred_ty, ctx1, tree1) = self.infer(ctx, expr)?;
            let inferred_applied = self.apply_ctx_type(&ctx1, &inferred_ty);
            let ty_applied = self.apply_ctx_type(&ctx1, ty);
            let (ctx2, tree2) = self.subtype(&ctx1, &inferred_applied, &ty_applied)?;
            let output = format!("{}", ctx2);
            Ok((
                ctx2,
                InferenceTree::new("ChkSub", &input, &output, vec![tree1, tree2]),
            ))
        }
    }
}
}

Checking mode provides specialized handling that takes advantage of known type information to guide the inference process more efficiently. For lambda expressions, when checking \( \lambda x:\tau_1. e \) against type \( \tau_1 \to \tau_2 \), the algorithm can immediately verify that the parameter types match and then recursively check the body \( e \) against the expected return type \( \tau_2 \). This direct decomposition avoids the need to synthesize types and then verify compatibility.

When checking against universal types \( \forall\alpha. \tau \), the algorithm introduces the type variable \( \alpha \) into the context and then checks the expression against the instantiated type \( \tau \). This approach ensures that the universal quantification is handled correctly while maintaining the scoping discipline required for sound type checking. When direct checking strategies are not applicable to the current expression and expected type combination, the algorithm falls back to a synthesis-plus-subtyping approach, where it first synthesizes a type for the expression and then verifies that this synthesized type is a subtype of the expected type.

Subtyping and Instantiation

Subtyping Relations

System F includes a subtyping system that handles the relationships between polymorphic types. The key insight is that ∀α. τ is more general than any specific instantiation of τ.

#![allow(unused)]
fn main() {
fn subtype(
    &mut self,
    ctx: &Context,
    t1: &Type,
    t2: &Type,
) -> TypeResult<(Context, InferenceTree)> {
    let input = format!("{} ⊢ {} <: {}", ctx, t1, t2);

    match (t1, t2) {
        (Type::Int, Type::Int) | (Type::Bool, Type::Bool) => Ok((
            ctx.clone(),
            InferenceTree::new("SubRefl", &input, &format!("{}", ctx), vec![]),
        )),
        (Type::Var(a), Type::Var(b)) if a == b => Ok((
            ctx.clone(),
            InferenceTree::new("SubReflTVar", &input, &format!("{}", ctx), vec![]),
        )),
        (Type::ETVar(a), Type::ETVar(b)) if a == b => Ok((
            ctx.clone(),
            InferenceTree::new("SubReflETVar", &input, &format!("{}", ctx), vec![]),
        )),
        (Type::Arrow(a1, a2), Type::Arrow(b1, b2)) => {
            let (ctx1, tree1) = self.subtype(ctx, b1, a1)?; // Contravariant in argument
            let (ctx2, tree2) = self.subtype(&ctx1, a2, b2)?; // Covariant in result
            let output = format!("{}", ctx2);
            Ok((
                ctx2,
                InferenceTree::new("SubArr", &input, &output, vec![tree1, tree2]),
            ))
        }
        (_, Type::Forall(b, t2_body)) => {
            let mut new_ctx = ctx.clone();
            new_ctx.push(Entry::TVarBnd(b.clone()));
            let (ctx1, tree) = self.subtype(&new_ctx, t1, t2_body)?;
            let (_, _, right) =
                ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == b));
            let final_ctx = Context(right);
            let output = format!("{}", final_ctx);
            Ok((
                final_ctx,
                InferenceTree::new("SubAllR", &input, &output, vec![tree]),
            ))
        }
        (Type::Forall(a, t1_body), _) => {
            let subst_t1 = self.subst_type(a, &Type::ETVar(a.clone()), t1_body);
            let mut new_ctx = ctx.clone();
            new_ctx.push(Entry::ETVarBnd(a.clone()));
            new_ctx.push(Entry::Mark(a.clone()));
            let (ctx1, tree) = self.subtype(&new_ctx, &subst_t1, t2)?;
            let (_, _, right) =
                ctx1.break3(|entry| matches!(entry, Entry::Mark(name) if name == a));
            let final_ctx = Context(right);
            let output = format!("{}", final_ctx);
            Ok((
                final_ctx,
                InferenceTree::new("SubAllL", &input, &output, vec![tree]),
            ))
        }
        (Type::ETVar(a), _) if !self.free_vars(t2).contains(a) => {
            let (ctx1, tree) = self.inst_l(ctx, a, t2)?;
            let output = format!("{}", ctx1);
            Ok((
                ctx1,
                InferenceTree::new("SubInstL", &input, &output, vec![tree]),
            ))
        }
        (_, Type::ETVar(a)) if !self.free_vars(t1).contains(a) => {
            let (ctx1, tree) = self.inst_r(ctx, t1, a)?;
            let output = format!("{}", ctx1);
            Ok((
                ctx1,
                InferenceTree::new("SubInstR", &input, &output, vec![tree]),
            ))
        }
        _ => Err(TypeError::SubtypingError {
            left: t1.clone(),
            right: t2.clone(),
            expr: None,
        }),
    }
}
}

The subtyping rules capture several essential relationships that govern type compatibility in System F. Function types exhibit the classic contravariant-covariant pattern, where a function that accepts more general arguments and returns more specific results is considered a subtype of a function with more specific argument requirements and more general return types. This means that a function of type \( A_1 \to B_1 \) is a subtype of \( A_2 \to B_2 \) when \( A_2 \leq A_1 \) (contravariant in arguments) and \( B_1 \leq B_2 \) (covariant in results).

Universal quantification follows the principle that \( \forall\alpha. \tau_1 \leq \tau_2 \) holds if \( \tau_1 \leq \tau_2 \) when \( \alpha \) is instantiated with a fresh existential variable, effectively allowing polymorphic types to be related through their instantiations. When subtyping involves existential variables, the algorithm must solve constraints about what these variables should be instantiated to in order to satisfy the subtyping relationship, often leading to the generation of additional constraints that propagate through the system.

Variable Instantiation

The instantiation judgments handle the core complexity of polymorphic type checking. When we have constraints like ^α ≤ τ (an existential variable should be at most as general as some type), we need to find appropriate instantiations.

#![allow(unused)]
fn main() {
fn inst_l(
    &mut self,
    ctx: &Context,
    a: &TyVar,
    ty: &Type,
) -> TypeResult<(Context, InferenceTree)> {
    let input = format!("{} ⊢ ^{} :=< {}", ctx, a, ty);

    match ty {
        Type::ETVar(b) if self.before(ctx, a, b) => {
            let (left, _, right) =
                ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == b));
            let new_ctx = Context::from_parts(
                left,
                Entry::SETVarBnd(b.clone(), Type::ETVar(a.clone())),
                right,
            );
            let output = format!("{}", new_ctx);
            Ok((
                new_ctx,
                InferenceTree::new("InstLReach", &input, &output, vec![]),
            ))
        }
        Type::Arrow(t1, t2) => {
            let a1 = self.fresh_tyvar();
            let a2 = self.fresh_tyvar();
            let (left, _, right) =
                ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
            let arrow_type = Type::Arrow(
                Box::new(Type::ETVar(a1.clone())),
                Box::new(Type::ETVar(a2.clone())),
            );
            let mut new_ctx = left;
            new_ctx.push(Entry::SETVarBnd(a.clone(), arrow_type));
            new_ctx.push(Entry::ETVarBnd(a1.clone()));
            new_ctx.push(Entry::ETVarBnd(a2.clone()));
            new_ctx.extend(right);
            let ctx1 = Context(new_ctx);

            let (ctx2, tree1) = self.inst_r(&ctx1, t1, &a1)?;
            let t2_applied = self.apply_ctx_type(&ctx2, t2);
            let (ctx3, tree2) = self.inst_l(&ctx2, &a2, &t2_applied)?;

            let output = format!("{}", ctx3);
            Ok((
                ctx3,
                InferenceTree::new("InstLArr", &input, &output, vec![tree1, tree2]),
            ))
        }
        Type::Forall(b, t) => {
            let mut new_ctx = ctx.clone();
            new_ctx.push(Entry::TVarBnd(b.clone()));
            let (ctx1, tree) = self.inst_l(&new_ctx, a, t)?;
            let (_, _, right) =
                ctx1.break3(|entry| matches!(entry, Entry::TVarBnd(name) if name == b));
            let final_ctx = Context(right);
            let output = format!("{}", final_ctx);
            Ok((
                final_ctx,
                InferenceTree::new("InstLAllR", &input, &output, vec![tree]),
            ))
        }
        _ if self.is_mono(ty) => {
            // Occurs check: ensure ^α doesn't occur in τ to prevent infinite types
            if self.occurs_check(a, ty) {
                return Err(TypeError::OccursCheck {
                    var: a.clone(),
                    ty: ty.clone(),
                    expr: None,
                });
            }
            let (left, _, right) =
                ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
            let new_ctx =
                Context::from_parts(left, Entry::SETVarBnd(a.clone(), ty.clone()), right);
            let output = format!("{}", new_ctx);
            Ok((
                new_ctx,
                InferenceTree::new("InstLSolve", &input, &output, vec![]),
            ))
        }
        _ => Err(TypeError::InstantiationError {
            var: a.clone(),
            ty: ty.clone(),
            expr: None,
        }),
    }
}
}

Left instantiation (inst_l) handles cases where the existential variable is on the left side of a constraint. This typically means we’re looking for the most general type that satisfies the constraint.

#![allow(unused)]
fn main() {
fn inst_r(
    &mut self,
    ctx: &Context,
    ty: &Type,
    a: &TyVar,
) -> TypeResult<(Context, InferenceTree)> {
    let input = format!("{} ⊢ {} :=< ^{}", ctx, ty, a);

    match ty {
        Type::ETVar(b) if self.before(ctx, a, b) => {
            let (left, _, right) =
                ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == b));
            let new_ctx = Context::from_parts(
                left,
                Entry::SETVarBnd(b.clone(), Type::ETVar(a.clone())),
                right,
            );
            let output = format!("{}", new_ctx);
            Ok((
                new_ctx,
                InferenceTree::new("InstRReach", &input, &output, vec![]),
            ))
        }
        Type::Arrow(t1, t2) => {
            let a1 = self.fresh_tyvar();
            let a2 = self.fresh_tyvar();
            let (left, _, right) =
                ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
            let arrow_type = Type::Arrow(
                Box::new(Type::ETVar(a1.clone())),
                Box::new(Type::ETVar(a2.clone())),
            );
            let mut new_ctx = left;
            new_ctx.push(Entry::SETVarBnd(a.clone(), arrow_type));
            new_ctx.push(Entry::ETVarBnd(a1.clone()));
            new_ctx.push(Entry::ETVarBnd(a2.clone()));
            new_ctx.extend(right);
            let ctx1 = Context(new_ctx);

            let (ctx2, tree1) = self.inst_l(&ctx1, &a1, t1)?;
            let t2_applied = self.apply_ctx_type(&ctx2, t2);
            let (ctx3, tree2) = self.inst_r(&ctx2, &t2_applied, &a2)?;

            let output = format!("{}", ctx3);
            Ok((
                ctx3,
                InferenceTree::new("InstRArr", &input, &output, vec![tree1, tree2]),
            ))
        }
        Type::Forall(b, t) => {
            let subst_t = self.subst_type(b, &Type::ETVar(b.clone()), t);
            let mut new_ctx = ctx.clone();
            new_ctx.push(Entry::ETVarBnd(b.clone()));
            new_ctx.push(Entry::Mark(b.clone()));
            let (ctx1, tree) = self.inst_r(&new_ctx, &subst_t, a)?;
            let (_, _, right) =
                ctx1.break3(|entry| matches!(entry, Entry::Mark(name) if name == b));
            let final_ctx = Context(right);
            let output = format!("{}", final_ctx);
            Ok((
                final_ctx,
                InferenceTree::new("InstRAllL", &input, &output, vec![tree]),
            ))
        }
        _ if self.is_mono(ty) => {
            // Occurs check: ensure ^α doesn't occur in τ to prevent infinite types
            if self.occurs_check(a, ty) {
                return Err(TypeError::OccursCheck {
                    var: a.clone(),
                    ty: ty.clone(),
                    expr: None,
                });
            }
            let (left, _, right) =
                ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
            let new_ctx =
                Context::from_parts(left, Entry::SETVarBnd(a.clone(), ty.clone()), right);
            let output = format!("{}", new_ctx);
            Ok((
                new_ctx,
                InferenceTree::new("InstRSolve", &input, &output, vec![]),
            ))
        }
        _ => Err(TypeError::InstantiationError {
            var: a.clone(),
            ty: ty.clone(),
            expr: None,
        }),
    }
}
}

Right instantiation (inst_r) handles the dual case where the existential variable is on the right. This typically means we’re looking for the most specific type that satisfies the constraint.

The instantiation algorithms include careful handling of several complex scenarios that arise during constraint solving. The reach relationship occurs when two existential variables are constrained against each other, requiring the algorithm to determine which variable should be solved in terms of the other while maintaining the proper ordering constraints. Arrow type instantiation requires breaking function types apart into their component argument and return types, creating separate instantiation constraints for each component that must be solved consistently.

The interaction between instantiation and universal quantification presents particular challenges, as the algorithm must ensure that polymorphic types are instantiated correctly while preserving the scoping discipline that prevents type variables from escaping their intended scope. These cases require constraint management to ensure that all relationships are maintained throughout the solving process.

Occurs Check

A critical component of sound type inference is the occurs check, which prevents infinite types from arising during unification. When solving constraints like ^α := τ, we must ensure that α does not occur within τ, as this would create a cyclic type.

#![allow(unused)]
fn main() {
/// Check if a type variable occurs in a type (occurs check for unification)
/// This prevents creating infinite types when solving ^α := τ by ensuring α
/// ∉ τ
fn occurs_check(&self, var: &TyVar, ty: &Type) -> bool {
    match ty {
        Type::Var(name) | Type::ETVar(name) => name == var,
        Type::Arrow(t1, t2) => self.occurs_check(var, t1) || self.occurs_check(var, t2),
        Type::Forall(bound_var, body) => {
            // If the variable is shadowed by the forall binding, it doesn't occur
            if bound_var == var {
                false
            } else {
                self.occurs_check(var, body)
            }
        }
        Type::Int | Type::Bool => false,
    }
}
}

The occurs check is applied during the InstLSolve and InstRSolve cases of instantiation. Without this check, the type system could accept programs that would lead to infinite types, violating the decidability of type checking.

Application Inference

Function application in System F requires careful handling because the function type might not be immediately apparent. The infer_app judgment implements the T-AppArrow and T-AppEVar rules defined earlier:

#![allow(unused)]
fn main() {
fn infer_app(
    &mut self,
    ctx: &Context,
    func_ty: &Type,
    arg: &Expr,
) -> TypeResult<(Type, Context, InferenceTree)> {
    let input = format!("{} ⊢ {:?} • {}", ctx, arg, func_ty);

    match func_ty {
        Type::Arrow(param_ty, result_ty) => {
            let (ctx1, tree) = self.check(ctx, arg, param_ty)?;
            let output = format!("{} ⇒⇒ {} ⊣ {}", input, result_ty, ctx1);
            Ok((
                result_ty.as_ref().clone(),
                ctx1,
                InferenceTree::new("InfAppArr", &input, &output, vec![tree]),
            ))
        }
        Type::ETVar(a) => {
            let a1 = self.fresh_tyvar();
            let a2 = self.fresh_tyvar();
            let (left, _, right) =
                ctx.break3(|entry| matches!(entry, Entry::ETVarBnd(name) if name == a));
            let arrow_type = Type::Arrow(
                Box::new(Type::ETVar(a1.clone())),
                Box::new(Type::ETVar(a2.clone())),
            );
            let mut new_ctx = left;
            new_ctx.push(Entry::SETVarBnd(a.clone(), arrow_type));
            new_ctx.push(Entry::ETVarBnd(a1.clone()));
            new_ctx.push(Entry::ETVarBnd(a2.clone()));
            new_ctx.extend(right);
            let ctx1 = Context(new_ctx);

            let (ctx2, tree) = self.check(&ctx1, arg, &Type::ETVar(a1))?;
            let output = format!("{} ⇒⇒ ^{} ⊣ {}", input, a2, ctx2);
            Ok((
                Type::ETVar(a2),
                ctx2,
                InferenceTree::new("InfAppETVar", &input, &output, vec![tree]),
            ))
        }
        _ => Err(TypeError::ApplicationTypeError {
            actual: func_ty.clone(),
            expr: None,
        }),
    }
}
}

The implementation handles two core application scenarios through distinct inference rules. The T-AppArrow rule applies when the function has a known arrow type \( A \to B \), allowing the algorithm to check the argument against \( A \) and return \( B \) as the result type. This straightforward case corresponds to the Type::Arrow pattern in the implementation and represents the standard function application scenario.

The T-AppEVar rule handles the more complex case where the function type is an existential variable \( \hat{\alpha} \). In this situation, the algorithm instantiates the existential variable as \( \hat{\alpha_1} \to \hat{\alpha_2} \) with fresh existential variables, then checks the argument against \( \hat{\alpha_1} \) and returns \( \hat{\alpha_2} \) as the result type. This corresponds to the Type::ETVar case and enables type inference even when the function type is initially unknown.

When the function has a polymorphic Forall type, the instantiation is handled through the subtyping mechanism using the SubAllL rule rather than directly in application inference. This design choice ensures soundness by routing polymorphic instantiation through the well-established subtyping infrastructure and follows the standard bidirectional algorithm design patterns.

Error Handling

Our implementation provides comprehensive error reporting that distinguishes between parse errors and type errors. Parse errors use source-located reporting with ariadne, while type errors provide contextual information about the expressions being typed.

#![allow(unused)]
fn main() {
use ariadne::{Color, Label, Report, ReportKind};
use thiserror::Error;
use crate::ast::Type;
/// Source span for error reporting
#[derive(Debug, Clone, PartialEq)]
pub struct Span {
    pub start: usize,
    pub end: usize,
}
impl Span {
    pub fn new(start: usize, end: usize) -> Self {
        Self { start, end }
    }

    pub fn point(pos: usize) -> Self {
        Self {
            start: pos,
            end: pos + 1,
        }
    }
}
impl TypeError {}
impl std::fmt::Display for TypeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TypeError::UnboundVariable { name, expr } => {
                write!(f, "Variable '{}' not found in context", name)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::ApplicationTypeError { actual, expr } => {
                write!(f, "Expected function type in application, got {}", actual)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::TypeApplicationError { actual, expr } => {
                write!(
                    f,
                    "Expected forall type in type application, got {}",
                    actual
                )?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::OccursCheck { var, ty, expr } => {
                write!(
                    f,
                    "Occurs check failed: variable '{}' occurs in type {}",
                    var, ty
                )?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::SubtypingError { left, right, expr } => {
                write!(f, "No matching rule for subtyping {} <: {}", left, right)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::InstantiationError { var, ty, expr } => {
                write!(f, "No matching rule for instantiation {} :=< {}", var, ty)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::CheckingError { expected, expr } => {
                write!(f, "No matching rule for checking against type {}", expected)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::InferenceError { expr } => {
                write!(f, "No matching rule for inference")?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::ContextError { message, expr } => {
                write!(f, "Context error: {}", message)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
        }
    }
}
#[derive(Error, Debug)]
#[allow(dead_code)]
pub enum ParseError {
    #[error("Parse error: {message}")]
    LalrpopError { message: String, span: Span },
}
impl ParseError {
    pub fn to_ariadne_report<'a>(
        &self,
        filename: &'a str,
    ) -> Report<'a, (&'a str, std::ops::Range<usize>)> {
        match self {
            ParseError::LalrpopError { message, span } => {
                Report::build(ReportKind::Error, filename, span.start)
                    .with_message("Parse Error")
                    .with_label(
                        Label::new((filename, span.start..span.end))
                            .with_message(message)
                            .with_color(Color::Red),
                    )
                    .finish()
            }
        }
    }
}
pub type TypeResult<T> = Result<T, TypeError>;
#[allow(dead_code)]
pub type ParseResult<T> = Result<T, ParseError>;
#[derive(Debug)]
#[allow(dead_code, clippy::enum_variant_names)]
pub enum TypeError {
    UnboundVariable {
        name: String,
        expr: Option<String>,
    },

    ApplicationTypeError {
        actual: Type,
        expr: Option<String>,
    },

    TypeApplicationError {
        actual: Type,
        expr: Option<String>,
    },

    OccursCheck {
        var: String,
        ty: Type,
        expr: Option<String>,
    },

    SubtypingError {
        left: Type,
        right: Type,
        expr: Option<String>,
    },

    InstantiationError {
        var: String,
        ty: Type,
        expr: Option<String>,
    },

    CheckingError {
        expected: Type,
        expr: Option<String>,
    },

    InferenceError {
        expr: Option<String>,
    },

    ContextError {
        message: String,
        expr: Option<String>,
    },
}
}

The type error system includes expression context to help developers understand where failures occur during type checking. Each error variant includes an expr field that stores the expression being typed when the error occurred, providing valuable debugging information.

Parse errors receive enhanced treatment with source location information for precise error reporting:

#![allow(unused)]
fn main() {
use ariadne::{Color, Label, Report, ReportKind};
use thiserror::Error;
use crate::ast::Type;
/// Source span for error reporting
#[derive(Debug, Clone, PartialEq)]
pub struct Span {
    pub start: usize,
    pub end: usize,
}
impl Span {
    pub fn new(start: usize, end: usize) -> Self {
        Self { start, end }
    }

    pub fn point(pos: usize) -> Self {
        Self {
            start: pos,
            end: pos + 1,
        }
    }
}
#[derive(Debug)]
#[allow(dead_code, clippy::enum_variant_names)]
pub enum TypeError {
    UnboundVariable {
        name: String,
        expr: Option<String>,
    },

    ApplicationTypeError {
        actual: Type,
        expr: Option<String>,
    },

    TypeApplicationError {
        actual: Type,
        expr: Option<String>,
    },

    OccursCheck {
        var: String,
        ty: Type,
        expr: Option<String>,
    },

    SubtypingError {
        left: Type,
        right: Type,
        expr: Option<String>,
    },

    InstantiationError {
        var: String,
        ty: Type,
        expr: Option<String>,
    },

    CheckingError {
        expected: Type,
        expr: Option<String>,
    },

    InferenceError {
        expr: Option<String>,
    },

    ContextError {
        message: String,
        expr: Option<String>,
    },
}
impl TypeError {}
impl std::fmt::Display for TypeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TypeError::UnboundVariable { name, expr } => {
                write!(f, "Variable '{}' not found in context", name)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::ApplicationTypeError { actual, expr } => {
                write!(f, "Expected function type in application, got {}", actual)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::TypeApplicationError { actual, expr } => {
                write!(
                    f,
                    "Expected forall type in type application, got {}",
                    actual
                )?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::OccursCheck { var, ty, expr } => {
                write!(
                    f,
                    "Occurs check failed: variable '{}' occurs in type {}",
                    var, ty
                )?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::SubtypingError { left, right, expr } => {
                write!(f, "No matching rule for subtyping {} <: {}", left, right)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::InstantiationError { var, ty, expr } => {
                write!(f, "No matching rule for instantiation {} :=< {}", var, ty)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::CheckingError { expected, expr } => {
                write!(f, "No matching rule for checking against type {}", expected)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::InferenceError { expr } => {
                write!(f, "No matching rule for inference")?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
            TypeError::ContextError { message, expr } => {
                write!(f, "Context error: {}", message)?;
                if let Some(expr) = expr {
                    write!(f, "\n  When typing expression: {}", expr)?;
                }
                Ok(())
            }
        }
    }
}
impl ParseError {
    pub fn to_ariadne_report<'a>(
        &self,
        filename: &'a str,
    ) -> Report<'a, (&'a str, std::ops::Range<usize>)> {
        match self {
            ParseError::LalrpopError { message, span } => {
                Report::build(ReportKind::Error, filename, span.start)
                    .with_message("Parse Error")
                    .with_label(
                        Label::new((filename, span.start..span.end))
                            .with_message(message)
                            .with_color(Color::Red),
                    )
                    .finish()
            }
        }
    }
}
pub type TypeResult<T> = Result<T, TypeError>;
#[allow(dead_code)]
pub type ParseResult<T> = Result<T, ParseError>;
#[derive(Error, Debug)]
#[allow(dead_code)]
pub enum ParseError {
    #[error("Parse error: {message}")]
    LalrpopError { message: String, span: Span },
}
}

This dual approach ensures that syntax errors receive precise source location feedback while type errors focus on logical relationships between expressions and types.

End-to-end

To demonstrate the complete System F implementation with all the modular typing rules working together, consider this lambda expression that doubles its integer argument:

$ cargo run -- check "\x : Int -> x + x"

The system produces the following output showing the complete inference process:

Parsed expression: Abs("x", Int, BinOp(Add, Var("x"), Var("x")))

Type checking successful!
Final type: Int -> Int

InfLam:  ⊢ Abs("x", Int, BinOp(Add, Var("x"), Var("x"))) =>  ⊢ Abs("x", Int, BinOp(Add, Var("x"), Var("x"))) ⇒ Int -> ^α0 ⊣ ^α0 = Int
  ChkSub: ^α0, x: Int ⊢ BinOp(Add, Var("x"), Var("x")) ⇐ ^α0 => ^α0 = Int, x: Int
    InfArith: ^α0, x: Int ⊢ BinOp(Add, Var("x"), Var("x")) => ^α0, x: Int ⊢ BinOp(Add, Var("x"), Var("x")) ⇒ Int ⊣ ^α0, x: Int
      ChkSub: ^α0, x: Int ⊢ Var("x") ⇐ Int => ^α0, x: Int
        InfVar: ^α0, x: Int ⊢ Var("x") => ^α0, x: Int ⊢ Var("x") ⇒ Int ⊣ ^α0, x: Int
        SubRefl: ^α0, x: Int ⊢ Int <: Int => ^α0, x: Int
      ChkSub: ^α0, x: Int ⊢ Var("x") ⇐ Int => ^α0, x: Int
        InfVar: ^α0, x: Int ⊢ Var("x") => ^α0, x: Int ⊢ Var("x") ⇒ Int ⊣ ^α0, x: Int
        SubRefl: ^α0, x: Int ⊢ Int <: Int => ^α0, x: Int
    SubInstR: ^α0, x: Int ⊢ Int <: ^α0 => ^α0 = Int, x: Int
      InstRSolve: ^α0, x: Int ⊢ Int :=< ^α0 => ^α0 = Int, x: Int

The final result shows Final type: Int -> Int, correctly inferring that this lambda expression is a function from integers to integers! The existential variable ^α0 in the proof tree gets resolved to Int through the constraint solving process, and the final type application ensures all existential variables are properly substituted in the output.

For non-trivial expressions, the constraint solving process may involve more complex reasoning and may require additional inference steps but using the proof tree method we can visualize the inference steps and understand the constraints being solved. This is a very powerful technique.

Mission Accomplished

And there we have it! A complete bidirectional type checker for System F with polymorphic types, existential variables, and constraint solving! The algorithm handles the complex interplay between synthesis and checking modes, managing existential variables and subtyping relationships with the precision needed for a production-quality type system.

The bidirectional approach transforms what could be an overwhelming inference problem into a systematic, decidable process that gives us both powerful expressiveness and reliable type safety. System F’s polymorphism opens up a whole new world of type-safe programming that feels almost magical once you see it in action!

Examples

System F’s power lies in its ability to express polymorphic computations that work across different types while maintaining type safety. This chapter explores a range of examples that demonstrate the expressive capabilities of our implementation, from simple polymorphic functions to complex higher-rank polymorphism scenarios.

Our examples are drawn from real test cases in the implementation, showing both the syntax our parser accepts and the types our bidirectional algorithm infers. These examples illustrate the key features of System F and demonstrate how the type system guides program construction.

Basic Polymorphic Functions

The fundamental building block of System F is the polymorphic identity function. This simple example demonstrates universal quantification and type application in their most basic form.

Let’s examine how our implementation handles the core examples from the test suite:

#![allow(unused)]
fn main() {
Basic literals
42
true
false

Type annotations
\x : Int -> x
\x : Bool -> x
(\x : Int -> x) 42
(\x : Bool -> x) true

Polymorphic identity using forall
forall a. \x : a -> x

Type application with brackets
(forall a. \x : a -> x) [Int]
(forall a. \x : a -> x) [Bool]

Higher-rank polymorphism examples  
\f : (forall a. a -> a) -> f
(\f : (forall a. a -> a) -> f) (forall a. \x : a -> x)

Nested type abstractions
forall a. forall b. \x : a -> \y : b -> x
forall a. forall b. \x : a -> \y : b -> y

Complex polymorphic examples
forall a. \f : (a -> a) -> \x : a -> f x

Function types
\f : (Int -> Int) -> \x : Int -> f x
(\f : (Int -> Int) -> \x : Int -> f x) (\y : Int -> y) 42

Error cases should be handled by the type checker
(\x : Int -> x) true
}

These test cases reveal several important aspects of our System F implementation:

Literal Types and Basic Functions

The simplest cases involve literal values and monomorphic functions, where our bidirectional algorithm produces clean, precise types:

• 42 : Int - Integer literals immediately receive the base type Int
• true : Bool and false : Bool - Boolean literals similarly receive their base types
• \x : Int -> x : Int -> Int - Monomorphic lambda expressions receive concrete function types

The type inference algorithm resolves all constraints to produce the most specific types possible. When we provide explicit type annotations, the algorithm can determine the complete function type without ambiguity.

Type Annotations and Explicit Typing

System F requires explicit type annotations on lambda parameters, and our bidirectional algorithm infers complete function types:

• \x : Int -> x produces the type Int -> Int
• (\x : Int -> x) 42 applies this function to produce type Int

These examples show how the algorithm handles the interaction between explicit annotations and type inference. The parameter type is fixed by annotation, and the algorithm infers that the return type must match the parameter type since we’re returning the parameter unchanged.

Universal Quantification in Practice

The true power of System F emerges with universal quantification. The polymorphic identity function demonstrates this clearly:

• forall a. \x : a -> x : ∀a. a -> a

This example shows several important aspects of our implementation:

1. Type Abstraction: The forall a. syntax introduces a type variable into scope
2. Polymorphic Parameters: The parameter x : a uses the quantified type variable
3. Complete Type Inference: The algorithm correctly resolves all type constraints to produce the expected polymorphic type

Type Application and Instantiation

Type application allows us to specialize polymorphic functions to specific types:

• (forall a. \x : a -> x) [Int] : Int -> Int
• (forall a. \x : a -> x) [Bool] : Bool -> Bool

Both applications produce the expected concrete function types, showing that our algorithm correctly instantiates the polymorphic function with the requested types and resolves all type constraints to their final forms.

Higher-Rank Polymorphism

System F supports higher-rank polymorphism, where polymorphic types can appear in argument positions. This creates functions that accept polymorphic arguments:

• \f : (forall a. a -> a) -> f : ^α0 -> ^α1

This function accepts any argument of type forall a. a -> a (a polymorphic identity function) and returns it unchanged. The higher-rank nature means the caller must provide a function that works for all types, not just some specific type.

The application example demonstrates this in action:

• (\f : (forall a. a -> a) -> f) (forall a. \x : a -> x) : ^α0

Here we pass the polymorphic identity function to a function that expects exactly that type signature. This showcases the precision and expressiveness of System F’s type system.

Nested Type Abstractions

System F allows multiple type abstractions to be nested, creating functions that are polymorphic in multiple type parameters:

• forall a. forall b. \x : a -> \y : b -> x : ∀a. ∀b. a -> b -> a
• forall a. forall b. \x : a -> \y : b -> y : ∀a. ∀b. a -> b -> b

These examples create functions that:

1. Are polymorphic in two types a and b
2. Accept arguments of those respective types
3. Return either the first or second argument

The K and S combinators from combinatory logic can be expressed naturally in this style, showing System F’s expressiveness for abstract computation. As a fun aside, you express every other function in terms of just K and S, although in practice this is more of an theoretical curio because its unimaginably inefficient.

Complex Polymorphism

More examples demonstrate how System F handles complex combinations of polymorphism and function application:

• forall a. \f : (a -> a) -> \x : a -> f x : ∀a. (a -> a) -> a -> a

This creates a function that:

1. Is polymorphic in type a
2. Accepts a function f : a -> a (an endomorphism on a)
3. Accepts a value x : a
4. Applies the function to the value

This pattern is fundamental in functional programming and demonstrates how System F naturally expresses higher-order polymorphic functions.

System F-ω

System F-ω stands as one of the most type systems in the theoretical foundations of programming languages, combining parametric polymorphism with higher-kinded types to enable powerful abstraction mechanisms. This system forms the theoretical backbone for modern functional programming languages like Haskell and provides the expressive power needed for advanced programming patterns including functors, monads, and generic programming.

To understand System F-ω’s significance, we must first examine its place within the broader landscape of type systems through the systematic classification known as the lambda cube.

The Lambda Cube: A Systematic Classification

The lambda cube, introduced by Henk Barendregt, provides a systematic way to understand how different type systems relate to each other by considering three independent dimensions of abstraction. Each dimension of the cube corresponds to a different kind of dependency between terms and types, where “dependency” refers to the capacity of a term or type to bind a term or type.

The three orthogonal axes of the lambda cube correspond to:

\(\to\)-axis (Types depending on Terms): Enables dependent types where the structure of types can depend on the values of terms. This allows types like Vector n where the type carries information about term-level values.

\(\uparrow\)-axis (Terms depending on Types): Enables polymorphism where terms can abstract over and depend on type parameters. This allows functions like identity : ∀α. α \\to α that work uniformly across all types.

\(\nearrow\)-axis (Types depending on Types): Enables type operators where types can abstract over and depend on other types. This allows type constructors like Maybe : * \\to * that take types as arguments and produce new types.

The eight vertices of the cube emerge from the different ways to combine these three dimensions of dependency. Each vertex represents a different typed system, obtained by enabling or disabling each form of abstraction:

Simply Typed Lambda Calculus (λ→): The foundation of the cube, supporting only term abstraction. Functions can abstract over terms (λx:τ. e) but types remain fixed and simple.

System F (λ2): Adds polymorphism through type abstraction, enabling terms to abstract over types (Λα. e). This introduces parametric polymorphism where functions like identity can work uniformly across all types.

System Fω (λω): Introduces higher-kinded types by adding type operators, allowing types to abstract over types (λα:κ. τ). This enables abstraction over type constructors like Maybe and List.

Lambda P (λP): Adds dependent types where types can depend on terms. This allows types like Vector n where the length n is a term-level value, enabling precise specifications of data structure properties.

System F-ω (λ2ω): Combines polymorphism with higher-kinded types, enabling both term abstraction over types and type abstraction over types. This is our target system, providing the expressiveness needed for modern functional programming.

System Fω-P (λωP): Combines higher-kinded types with dependent types, allowing type-level computation that can depend on term-level values.

System FP (λ2P): Combines polymorphism with dependent types, enabling functions that are parametric over both types and values while allowing types to depend on those values.

Calculus of Constructions (λ2ωP): The most expressive corner, combining all three forms of abstraction. This system underlies proof assistants like Coq and enables types that can express arbitrary logical propositions.

The lambda cube demonstrates that these eight systems form a natural hierarchy, with each system being a conservative extension of those below it in the ordering. System F-ω sits at a sweet spot, providing significant expressive power while maintaining decidable type checking and relatively straightforward implementation techniques.

Where System F introduced universal quantification over types (∀α. τ), System F-ω extends this to quantification over type constructors of arbitrary kinds. This enables us to write functions that are polymorphic not just over types like Int or Bool, but over type constructors like Maybe or List, and even higher-order type constructors that take multiple type arguments.

The type system is stratified by kinds, which classify types just as types classify terms:

\[ \begin{align*} \text{kinds} \quad \kappa &::= \star \mid \kappa_1 \to \kappa_2 \\ \text{types} \quad \tau &::= \alpha \mid \tau_1 \to \tau_2 \mid \forall \alpha:\kappa. \tau \mid \tau_1 \tau_2 \mid \lambda \alpha:\kappa. \tau \\ \text{terms} \quad e &::= x \mid \lambda x:\tau. e \mid e_1 e_2 \mid \Lambda \alpha:\kappa. e \mid e[\tau] \end{align*} \]

This hierarchy extends naturally upward: just as we need types to classify terms and kinds to classify types, we could in principle introduce sorts to classify kinds.

While System F-ω stops at the kind level for practical reasons, the theoretical framework suggests an infinite tower: terms have types, types have kinds, kinds have sorts, sorts have supersorts, and so on. This infinite hierarchy, known as the cumulative hierarchy in type theory, reflects the fundamental principle that every mathematical object must be classified by something at a higher level of abstraction. The aleph notation acknowledges this infinite ascent while recognizing that practical type systems must terminate the hierarchy at some finite level. More on this when we get to dependent types later!

Higher-Kinded Types

The key innovation of System F-ω is the kind system. While ordinary types like Int and Bool have kind * (pronounced “star”), type constructors have function kinds:

• Maybe has kind * -> * (takes a type, returns a type)
• Either has kind * -> * -> * (takes two types, returns a type)
• A hypothetical StateT might have kind * -> (* -> *) -> * -> *

This stratification enables precise reasoning about type-level functions while maintaining decidable type checking.

Our implementation demonstrates these concepts through a rich surface language that compiles to a core System F-ω calculus. The surface syntax provides familiar algebraic data types and pattern matching:

data Maybe a = Nothing | Just a
data Either a b = Left a | Right b
data List a = Nil | Cons a (List a)

These declarations create type constructors of appropriate kinds. Maybe becomes a type-level function that, when applied to a type like Int, produces the concrete type Maybe Int.

Type-Level Computation

System F-ω supports type-level computation through type application and type abstraction. While our implementation focuses on the foundational mechanisms, the theoretical system allows arbitrary computation at the type level:

• Type Application: F τ applies type constructor F to type τ
• Type Abstraction: λα:κ. τ creates a type-level function

This enables type-level programming, though our implementation focuses on the essential features needed for practical programming language design.

Polymorphic Data Structures

The combination of universal quantification and higher-kinded types enables elegant expression of polymorphic data structures and functions:

map :: forall a b. (a -> b) -> List a -> List b
foldRight :: forall a b. (a -> b -> b) -> b -> List a -> b

These functions work uniformly across all types, demonstrating the power of parametric polymorphism in System F-ω.

Implementation Architecture

Our System F-ω implementation consists of several key components working together:

• Surface Language: A user-friendly syntax with algebraic data types, pattern matching, and type inference
• Core Language: An explicit System F-ω calculus with kinds, type abstractions, and applications
• Elaboration: Translation from surface to core, inserting implicit type arguments and abstractions
• Type Inference: A bidirectional algorithm based on the DK worklist approach
• Kind Inference: Automatic inference of kinds for type constructors and type expressions

The implementation demonstrates that type systems can be made practical through careful algorithm design and implementation techniques.

Our implementation uses bidirectional type checking, which splits the problem into two complementary modes:

• Synthesis (⇒): Given an expression, determine its type
• Checking (⇐): Given an expression and expected type, verify compatibility

This approach handles the complexity of higher-rank polymorphism and type applications while maintaining decidability and providing good error messages.

Language Design

Our System F-ω implementation employs a two-layer architecture that separates user-facing syntax from the internal representation used by the type checker. This design pattern, common in compilers, allows us to provide an ergonomic programming experience while maintaining a clean theoretical foundation for type checking algorithms.

The surface language offers familiar syntax with algebraic data types, pattern matching, and implicit type inference. The core language provides an explicit representation of System F-ω with kinds, type abstractions, and applications. Translation between these layers handles the complex process of inserting implicit type arguments and managing the type-level computations that System F-ω enables.

The Haskell Hairball

Before diving into our clean System F-ω design, we must acknowledge the elephant in the room: Haskell, the language that wore the hairshirt for two decades and somehow convinced a generation of programmers that this constituted virtue. Our implementation deliberately avoids the accumulated cruft that has transformed what began as an elegant research language into something resembling a Lovecraftian nightmare of interacting language extensions.

Haskell represents a fascinating case study in how good intentions, academic enthusiasm, and the sunk cost fallacy can combine to create somethign is both a beautiful and a trainwreck at the same time. Consider the design decisions that seemed reasonable at the time but now serve as warnings to future language designers:

• Call-by-need evaluation: Lazy evaluation sounds theoretically elegant until you discover that reasoning about space and time complexity requires divination skills, and debugging memory leaks involves consulting tea leaves about thunk accumulation patterns

• LARPing dependent types: Rather than implementing actual dependent types, Haskell developed an increasingly baroque system of type-level programming that lets you pretend you have dependent types while maintaining all the complexity and none of the theoretical guarantees

• Language extension proliferation: What started as a few modest extensions has metastasized into a catalog of over 200 language pragmas, creating not one language but a thousand mutually incompatible dialects that share only syntax

• Infinite RAM Requirements: The GHC compiler requires approximately the computational resources of a small Dyson swarm to bootstrap itself, making it effectively impossible to port to new architectures without access to industrial-scale computing infrastructure

Haskell, while historically important and technically very interesting, is an excellent guide for how not to design a language.

Surface Language Syntax

The surface language provides a Haskell-like syntax that programmers can work with naturally. Users write programs without explicit type abstractions or applications, and the compiler inserts these automatically during elaboration to the core language.

Data Type Declarations

#![allow(unused)]
fn main() {
use std::fmt;
#[derive(Debug, Clone, PartialEq)]
pub struct Module {
    pub declarations: Vec<Declaration>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct Constructor {
    pub name: String,
    pub fields: Vec<Type>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TypeScheme {
    pub quantified: Vec<String>, // forall a b. ...
    pub ty: Type,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    /// Type variable: a, b
    Var(String),
    /// Type constructor: Int, Bool, Color
    Con(String),
    /// Function type: a -> b
    Arrow(Box<Type>, Box<Type>),
    /// Type application: List a, Maybe Int
    App(Box<Type>, Box<Type>),
    /// Record type: { x :: Int, y :: Int }
    Record(Vec<(String, Type)>),
    /// Forall type: forall a b. a -> b
    Forall(Vec<String>, Box<Type>),
}
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    /// Variable: x, y
    Var(String),
    /// Integer literal: 42
    LitInt(i64),
    /// Lambda: \x -> e
    Lambda { param: String, body: Box<Expr> },
    /// Function application: f x
    App { func: Box<Expr>, arg: Box<Expr> },
    /// Conditional: if cond then e1 else e2
    If {
        cond: Box<Expr>,
        then_branch: Box<Expr>,
        else_branch: Box<Expr>,
    },
    /// Pattern matching: match e { p1 -> e1, p2 -> e2 }
    Match {
        expr: Box<Expr>,
        arms: Vec<MatchArm>,
    },
    /// Binary operation: e1 + e2
    BinOp {
        op: BinOp,
        left: Box<Expr>,
        right: Box<Expr>,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct MatchArm {
    pub pattern: Pattern,
    pub body: Expr,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Pattern {
    /// Wildcard pattern: _
    Wildcard,
    /// Variable pattern: x
    Var(String),
    /// Constructor pattern: Red, Cons x xs
    Constructor { name: String, args: Vec<Pattern> },
}
#[derive(Debug, Clone, PartialEq)]
pub enum BinOp {
    Add,
    Sub,
    Mul,
    Div,
    // Comparison operations
    Lt,
    Le,
}
impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Type::Var(name) | Type::Con(name) => write!(f, "{}", name),
            Type::Arrow(t1, t2) => write!(f, "{} -> {}", t1, t2),
            Type::App(t1, t2) => write!(f, "{} {}", t1, t2),
            Type::Record(fields) => {
                write!(f, "{{ ")?;
                for (i, (name, ty)) in fields.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{} :: {}", name, ty)?;
                }
                write!(f, " }}")
            }
            Type::Forall(vars, ty) => {
                write!(f, "forall ")?;
                for (i, var) in vars.iter().enumerate() {
                    if i > 0 {
                        write!(f, " ")?;
                    }
                    write!(f, "{}", var)?;
                }
                write!(f, ". {}", ty)
            }
        }
    }
}
impl fmt::Display for TypeScheme {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.quantified.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "forall {}. {}", self.quantified.join(" "), self.ty)
        }
    }
}
impl fmt::Display for BinOp {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            BinOp::Add => write!(f, "+"),
            BinOp::Sub => write!(f, "-"),
            BinOp::Mul => write!(f, "*"),
            BinOp::Div => write!(f, "/"),
            BinOp::Lt => write!(f, "<"),
            BinOp::Le => write!(f, "<="),
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Declaration {
    /// Data type declaration: data Color = Red | Blue
    Data {
        name: String,
        type_params: Vec<String>,
        constructors: Vec<Constructor>,
    },
    /// Function type signature: fib :: Int -> Int
    TypeSig {
        name: String,
        type_scheme: TypeScheme,
    },
    /// Function definition: fib n = match n { 0 -> 0, ... }
    FunDef {
        name: String,
        params: Vec<String>,
        body: Expr,
    },
}
}

Algebraic data types form the foundation of our surface language. These declarations create both type constructors and value constructors:

data Bool = True | False;
data Maybe a = Nothing | Just a;
data Either a b = Left a | Right b;
data List a = Nil | Cons a (List a);

Each data declaration introduces:

• A type constructor (like Maybe or List) that can be applied to type arguments
• Value constructors (like Nothing, Just, Nil, Cons) for creating values of the type
• Implicit kind information determined by the number and usage of type parameters

The surface language allows natural type parameter syntax where data List a automatically implies that List has kind * -> *.

Type Signatures and Schemes

#![allow(unused)]
fn main() {
use std::fmt;
#[derive(Debug, Clone, PartialEq)]
pub struct Module {
    pub declarations: Vec<Declaration>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Declaration {
    /// Data type declaration: data Color = Red | Blue
    Data {
        name: String,
        type_params: Vec<String>,
        constructors: Vec<Constructor>,
    },
    /// Function type signature: fib :: Int -> Int
    TypeSig {
        name: String,
        type_scheme: TypeScheme,
    },
    /// Function definition: fib n = match n { 0 -> 0, ... }
    FunDef {
        name: String,
        params: Vec<String>,
        body: Expr,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct Constructor {
    pub name: String,
    pub fields: Vec<Type>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    /// Type variable: a, b
    Var(String),
    /// Type constructor: Int, Bool, Color
    Con(String),
    /// Function type: a -> b
    Arrow(Box<Type>, Box<Type>),
    /// Type application: List a, Maybe Int
    App(Box<Type>, Box<Type>),
    /// Record type: { x :: Int, y :: Int }
    Record(Vec<(String, Type)>),
    /// Forall type: forall a b. a -> b
    Forall(Vec<String>, Box<Type>),
}
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    /// Variable: x, y
    Var(String),
    /// Integer literal: 42
    LitInt(i64),
    /// Lambda: \x -> e
    Lambda { param: String, body: Box<Expr> },
    /// Function application: f x
    App { func: Box<Expr>, arg: Box<Expr> },
    /// Conditional: if cond then e1 else e2
    If {
        cond: Box<Expr>,
        then_branch: Box<Expr>,
        else_branch: Box<Expr>,
    },
    /// Pattern matching: match e { p1 -> e1, p2 -> e2 }
    Match {
        expr: Box<Expr>,
        arms: Vec<MatchArm>,
    },
    /// Binary operation: e1 + e2
    BinOp {
        op: BinOp,
        left: Box<Expr>,
        right: Box<Expr>,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct MatchArm {
    pub pattern: Pattern,
    pub body: Expr,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Pattern {
    /// Wildcard pattern: _
    Wildcard,
    /// Variable pattern: x
    Var(String),
    /// Constructor pattern: Red, Cons x xs
    Constructor { name: String, args: Vec<Pattern> },
}
#[derive(Debug, Clone, PartialEq)]
pub enum BinOp {
    Add,
    Sub,
    Mul,
    Div,
    // Comparison operations
    Lt,
    Le,
}
impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Type::Var(name) | Type::Con(name) => write!(f, "{}", name),
            Type::Arrow(t1, t2) => write!(f, "{} -> {}", t1, t2),
            Type::App(t1, t2) => write!(f, "{} {}", t1, t2),
            Type::Record(fields) => {
                write!(f, "{{ ")?;
                for (i, (name, ty)) in fields.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{} :: {}", name, ty)?;
                }
                write!(f, " }}")
            }
            Type::Forall(vars, ty) => {
                write!(f, "forall ")?;
                for (i, var) in vars.iter().enumerate() {
                    if i > 0 {
                        write!(f, " ")?;
                    }
                    write!(f, "{}", var)?;
                }
                write!(f, ". {}", ty)
            }
        }
    }
}
impl fmt::Display for TypeScheme {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.quantified.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "forall {}. {}", self.quantified.join(" "), self.ty)
        }
    }
}
impl fmt::Display for BinOp {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            BinOp::Add => write!(f, "+"),
            BinOp::Sub => write!(f, "-"),
            BinOp::Mul => write!(f, "*"),
            BinOp::Div => write!(f, "/"),
            BinOp::Lt => write!(f, "<"),
            BinOp::Le => write!(f, "<="),
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub struct TypeScheme {
    pub quantified: Vec<String>, // forall a b. ...
    pub ty: Type,
}
}

Type schemes provide explicit polymorphic type signatures using familiar quantifier syntax:

identity :: a -> a;
const :: a -> b -> a;
map :: (a -> b) -> List a -> List b;

The surface language uses implicit quantification - any free type variables in a type signature are automatically universally quantified. This provides the convenience of languages like Haskell while maintaining the theoretical precision of System F-ω.

Surface Types

#![allow(unused)]
fn main() {
use std::fmt;
#[derive(Debug, Clone, PartialEq)]
pub struct Module {
    pub declarations: Vec<Declaration>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Declaration {
    /// Data type declaration: data Color = Red | Blue
    Data {
        name: String,
        type_params: Vec<String>,
        constructors: Vec<Constructor>,
    },
    /// Function type signature: fib :: Int -> Int
    TypeSig {
        name: String,
        type_scheme: TypeScheme,
    },
    /// Function definition: fib n = match n { 0 -> 0, ... }
    FunDef {
        name: String,
        params: Vec<String>,
        body: Expr,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct Constructor {
    pub name: String,
    pub fields: Vec<Type>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TypeScheme {
    pub quantified: Vec<String>, // forall a b. ...
    pub ty: Type,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    /// Variable: x, y
    Var(String),
    /// Integer literal: 42
    LitInt(i64),
    /// Lambda: \x -> e
    Lambda { param: String, body: Box<Expr> },
    /// Function application: f x
    App { func: Box<Expr>, arg: Box<Expr> },
    /// Conditional: if cond then e1 else e2
    If {
        cond: Box<Expr>,
        then_branch: Box<Expr>,
        else_branch: Box<Expr>,
    },
    /// Pattern matching: match e { p1 -> e1, p2 -> e2 }
    Match {
        expr: Box<Expr>,
        arms: Vec<MatchArm>,
    },
    /// Binary operation: e1 + e2
    BinOp {
        op: BinOp,
        left: Box<Expr>,
        right: Box<Expr>,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct MatchArm {
    pub pattern: Pattern,
    pub body: Expr,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Pattern {
    /// Wildcard pattern: _
    Wildcard,
    /// Variable pattern: x
    Var(String),
    /// Constructor pattern: Red, Cons x xs
    Constructor { name: String, args: Vec<Pattern> },
}
#[derive(Debug, Clone, PartialEq)]
pub enum BinOp {
    Add,
    Sub,
    Mul,
    Div,
    // Comparison operations
    Lt,
    Le,
}
impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Type::Var(name) | Type::Con(name) => write!(f, "{}", name),
            Type::Arrow(t1, t2) => write!(f, "{} -> {}", t1, t2),
            Type::App(t1, t2) => write!(f, "{} {}", t1, t2),
            Type::Record(fields) => {
                write!(f, "{{ ")?;
                for (i, (name, ty)) in fields.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{} :: {}", name, ty)?;
                }
                write!(f, " }}")
            }
            Type::Forall(vars, ty) => {
                write!(f, "forall ")?;
                for (i, var) in vars.iter().enumerate() {
                    if i > 0 {
                        write!(f, " ")?;
                    }
                    write!(f, "{}", var)?;
                }
                write!(f, ". {}", ty)
            }
        }
    }
}
impl fmt::Display for TypeScheme {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.quantified.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "forall {}. {}", self.quantified.join(" "), self.ty)
        }
    }
}
impl fmt::Display for BinOp {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            BinOp::Add => write!(f, "+"),
            BinOp::Sub => write!(f, "-"),
            BinOp::Mul => write!(f, "*"),
            BinOp::Div => write!(f, "/"),
            BinOp::Lt => write!(f, "<"),
            BinOp::Le => write!(f, "<="),
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    /// Type variable: a, b
    Var(String),
    /// Type constructor: Int, Bool, Color
    Con(String),
    /// Function type: a -> b
    Arrow(Box<Type>, Box<Type>),
    /// Type application: List a, Maybe Int
    App(Box<Type>, Box<Type>),
    /// Record type: { x :: Int, y :: Int }
    Record(Vec<(String, Type)>),
    /// Forall type: forall a b. a -> b
    Forall(Vec<String>, Box<Type>),
}
}

The surface type system includes:

• Type variables (a, b) for generic parameters
• Type constructors (Int, Bool, List) for concrete and parameterized types
• Function types (a -> b) with right-associative arrow notation
• Type application (List a, Maybe Int) for applying type constructors

Notably absent from the surface language are explicit type abstractions, type applications to terms, and kind annotations. These are inferred and inserted automatically during elaboration.

Expression Language

#![allow(unused)]
fn main() {
use std::fmt;
#[derive(Debug, Clone, PartialEq)]
pub struct Module {
    pub declarations: Vec<Declaration>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Declaration {
    /// Data type declaration: data Color = Red | Blue
    Data {
        name: String,
        type_params: Vec<String>,
        constructors: Vec<Constructor>,
    },
    /// Function type signature: fib :: Int -> Int
    TypeSig {
        name: String,
        type_scheme: TypeScheme,
    },
    /// Function definition: fib n = match n { 0 -> 0, ... }
    FunDef {
        name: String,
        params: Vec<String>,
        body: Expr,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct Constructor {
    pub name: String,
    pub fields: Vec<Type>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TypeScheme {
    pub quantified: Vec<String>, // forall a b. ...
    pub ty: Type,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    /// Type variable: a, b
    Var(String),
    /// Type constructor: Int, Bool, Color
    Con(String),
    /// Function type: a -> b
    Arrow(Box<Type>, Box<Type>),
    /// Type application: List a, Maybe Int
    App(Box<Type>, Box<Type>),
    /// Record type: { x :: Int, y :: Int }
    Record(Vec<(String, Type)>),
    /// Forall type: forall a b. a -> b
    Forall(Vec<String>, Box<Type>),
}
#[derive(Debug, Clone, PartialEq)]
pub struct MatchArm {
    pub pattern: Pattern,
    pub body: Expr,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Pattern {
    /// Wildcard pattern: _
    Wildcard,
    /// Variable pattern: x
    Var(String),
    /// Constructor pattern: Red, Cons x xs
    Constructor { name: String, args: Vec<Pattern> },
}
#[derive(Debug, Clone, PartialEq)]
pub enum BinOp {
    Add,
    Sub,
    Mul,
    Div,
    // Comparison operations
    Lt,
    Le,
}
impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Type::Var(name) | Type::Con(name) => write!(f, "{}", name),
            Type::Arrow(t1, t2) => write!(f, "{} -> {}", t1, t2),
            Type::App(t1, t2) => write!(f, "{} {}", t1, t2),
            Type::Record(fields) => {
                write!(f, "{{ ")?;
                for (i, (name, ty)) in fields.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{} :: {}", name, ty)?;
                }
                write!(f, " }}")
            }
            Type::Forall(vars, ty) => {
                write!(f, "forall ")?;
                for (i, var) in vars.iter().enumerate() {
                    if i > 0 {
                        write!(f, " ")?;
                    }
                    write!(f, "{}", var)?;
                }
                write!(f, ". {}", ty)
            }
        }
    }
}
impl fmt::Display for TypeScheme {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.quantified.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "forall {}. {}", self.quantified.join(" "), self.ty)
        }
    }
}
impl fmt::Display for BinOp {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            BinOp::Add => write!(f, "+"),
            BinOp::Sub => write!(f, "-"),
            BinOp::Mul => write!(f, "*"),
            BinOp::Div => write!(f, "/"),
            BinOp::Lt => write!(f, "<"),
            BinOp::Le => write!(f, "<="),
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    /// Variable: x, y
    Var(String),
    /// Integer literal: 42
    LitInt(i64),
    /// Lambda: \x -> e
    Lambda { param: String, body: Box<Expr> },
    /// Function application: f x
    App { func: Box<Expr>, arg: Box<Expr> },
    /// Conditional: if cond then e1 else e2
    If {
        cond: Box<Expr>,
        then_branch: Box<Expr>,
        else_branch: Box<Expr>,
    },
    /// Pattern matching: match e { p1 -> e1, p2 -> e2 }
    Match {
        expr: Box<Expr>,
        arms: Vec<MatchArm>,
    },
    /// Binary operation: e1 + e2
    BinOp {
        op: BinOp,
        left: Box<Expr>,
        right: Box<Expr>,
    },
}
}

The surface expression language supports:

Variables and Literals: Standard identifiers and numeric/boolean constants Function Application: f x applies function f to argument x Lambda Abstractions: \x -> e creates anonymous functions Pattern Matching: match e { p1 -> e1; p2 -> e2; } for case analysis Constructor Application: Just 42, Cons x xs for building data structures

The surface language omits explicit type abstractions (Λα. e) and type applications (e [τ]). These System F-ω constructs are handled automatically by the elaboration process.

Core Language Representation

The core language provides an explicit encoding of System F-ω with full type-level computation capabilities. This representation makes type checking tractable by exposing all implicit operations from the surface language.

Core Types and Kinds

#![allow(unused)]
fn main() {
use std::fmt;
/// Core language for System-F-ω with explicit type applications and
/// abstractions
#[derive(Debug, Clone, PartialEq)]
pub struct CoreModule {
    pub type_defs: Vec<TypeDef>,
    pub term_defs: Vec<TermDef>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TypeDef {
    pub name: String,
    pub kind: Kind,
    pub constructors: Vec<DataConstructor>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct DataConstructor {
    pub name: String,
    pub ty: CoreType,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TermDef {
    pub name: String,
    pub ty: CoreType,
    pub body: CoreTerm,
}
/// Core types (System-F-ω)
#[derive(Debug, Clone, PartialEq)]
pub enum CoreType {
    /// Type variable: a
    Var(String),
    /// Existential type variable: ^a
    ETVar(String),
    /// Type constructor: Int, Bool
    Con(String),
    /// Function type: T1 -> T2
    Arrow(Box<CoreType>, Box<CoreType>),
    /// Universal quantification: ∀a. T
    Forall(String, Box<CoreType>),
    /// Type application: F T
    App(Box<CoreType>, Box<CoreType>),
    /// Product type: T1 × T2
    Product(Vec<CoreType>),
}
/// Core terms (System-F)
#[derive(Debug, Clone, PartialEq)]
pub enum CoreTerm {
    /// Variable: x
    Var(String),
    /// Integer literal: 42
    LitInt(i64),
    /// Lambda abstraction: λx:T. e
    Lambda {
        param: String,
        param_ty: CoreType,
        body: Box<CoreTerm>,
    },
    /// Application: e1 e2
    App {
        func: Box<CoreTerm>,
        arg: Box<CoreTerm>,
    },
    /// Type abstraction: Λα. e
    TypeLambda { param: String, body: Box<CoreTerm> },
    /// Constructor: C e1 ... en
    Constructor { name: String, args: Vec<CoreTerm> },
    /// Pattern matching: case e of { p1 -> e1; ... }
    Case {
        scrutinee: Box<CoreTerm>,
        arms: Vec<CaseArm>,
    },
    /// Built-in operations
    BinOp {
        op: CoreBinOp,
        left: Box<CoreTerm>,
        right: Box<CoreTerm>,
    },
    /// Conditional: if e1 then e2 else e3
    If {
        cond: Box<CoreTerm>,
        then_branch: Box<CoreTerm>,
        else_branch: Box<CoreTerm>,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct CaseArm {
    pub pattern: CorePattern,
    pub body: CoreTerm,
}
#[derive(Debug, Clone, PartialEq)]
pub enum CorePattern {
    /// Wildcard: _
    Wildcard,
    /// Variable: x
    Var(String),
    /// Constructor: C p1 ... pn
    Constructor {
        name: String,
        args: Vec<CorePattern>,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub enum CoreBinOp {
    Add,
    Sub,
    Mul,
    Div,
    Lt,
    Le,
}
impl fmt::Display for Kind {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Kind::Star => write!(f, "*"),
            Kind::Arrow(k1, k2) => write!(f, "{} -> {}", k1, k2),
        }
    }
}
impl fmt::Display for CoreType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CoreType::Var(name) => write!(f, "{}", name),
            CoreType::ETVar(name) => write!(f, "^{}", name),
            CoreType::Con(name) => write!(f, "{}", name),
            CoreType::Arrow(t1, t2) => write!(f, "{} -> {}", t1, t2),
            CoreType::Forall(var, ty) => write!(f, "∀{}. {}", var, ty),
            CoreType::App(t1, t2) => write!(f, "({} {})", t1, t2),
            CoreType::Product(types) => {
                write!(f, "(")?;
                for (i, ty) in types.iter().enumerate() {
                    if i > 0 {
                        write!(f, " × ")?;
                    }
                    write!(f, "{}", ty)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl fmt::Display for CoreBinOp {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CoreBinOp::Add => write!(f, "+"),
            CoreBinOp::Sub => write!(f, "-"),
            CoreBinOp::Mul => write!(f, "*"),
            CoreBinOp::Div => write!(f, "/"),
            CoreBinOp::Lt => write!(f, "<"),
            CoreBinOp::Le => write!(f, "<="),
        }
    }
}
impl fmt::Display for CoreTerm {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CoreTerm::Var(name) => write!(f, "{}", name),
            CoreTerm::LitInt(n) => write!(f, "{}", n),
            CoreTerm::Lambda {
                param,
                param_ty,
                body,
            } => {
                write!(f, "λ{} : {}. {}", param, param_ty, body)
            }
            CoreTerm::App { func, arg } => write!(f, "{} {}", func, arg),
            CoreTerm::Constructor { name, args } => {
                if args.is_empty() {
                    write!(f, "{}", name)
                } else {
                    write!(f, "{}", name)?;
                    for arg in args {
                        write!(f, " {}", arg)?;
                    }
                    Ok(())
                }
            }
            CoreTerm::If {
                cond,
                then_branch,
                else_branch,
            } => {
                write!(f, "if {} then {} else {}", cond, then_branch, else_branch)
            }
            CoreTerm::Case { scrutinee, arms: _ } => {
                write!(f, "match {} {{ ... }}", scrutinee)
            }
            CoreTerm::BinOp { op, left, right } => {
                write!(f, "{} {} {}", left, op, right)
            }
            _ => write!(f, "<term>"),
        }
    }
}
/// Kinds for types
#[derive(Debug, Clone, PartialEq)]
pub enum Kind {
    /// Base kind: *
    Star,
    /// Function kind: k1 -> k2
    Arrow(Box<Kind>, Box<Kind>),
}
}

The kind system classifies types hierarchically:

• Star (*) for ordinary types like Int, Bool, List Int
• Arrow(k1, k2) for type constructors like Maybe : * -> * or Either : * -> * -> *

#![allow(unused)]
fn main() {
use std::fmt;
/// Core language for System-F-ω with explicit type applications and
/// abstractions
#[derive(Debug, Clone, PartialEq)]
pub struct CoreModule {
    pub type_defs: Vec<TypeDef>,
    pub term_defs: Vec<TermDef>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TypeDef {
    pub name: String,
    pub kind: Kind,
    pub constructors: Vec<DataConstructor>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct DataConstructor {
    pub name: String,
    pub ty: CoreType,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TermDef {
    pub name: String,
    pub ty: CoreType,
    pub body: CoreTerm,
}
/// Kinds for types
#[derive(Debug, Clone, PartialEq)]
pub enum Kind {
    /// Base kind: *
    Star,
    /// Function kind: k1 -> k2
    Arrow(Box<Kind>, Box<Kind>),
}
/// Core terms (System-F)
#[derive(Debug, Clone, PartialEq)]
pub enum CoreTerm {
    /// Variable: x
    Var(String),
    /// Integer literal: 42
    LitInt(i64),
    /// Lambda abstraction: λx:T. e
    Lambda {
        param: String,
        param_ty: CoreType,
        body: Box<CoreTerm>,
    },
    /// Application: e1 e2
    App {
        func: Box<CoreTerm>,
        arg: Box<CoreTerm>,
    },
    /// Type abstraction: Λα. e
    TypeLambda { param: String, body: Box<CoreTerm> },
    /// Constructor: C e1 ... en
    Constructor { name: String, args: Vec<CoreTerm> },
    /// Pattern matching: case e of { p1 -> e1; ... }
    Case {
        scrutinee: Box<CoreTerm>,
        arms: Vec<CaseArm>,
    },
    /// Built-in operations
    BinOp {
        op: CoreBinOp,
        left: Box<CoreTerm>,
        right: Box<CoreTerm>,
    },
    /// Conditional: if e1 then e2 else e3
    If {
        cond: Box<CoreTerm>,
        then_branch: Box<CoreTerm>,
        else_branch: Box<CoreTerm>,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct CaseArm {
    pub pattern: CorePattern,
    pub body: CoreTerm,
}
#[derive(Debug, Clone, PartialEq)]
pub enum CorePattern {
    /// Wildcard: _
    Wildcard,
    /// Variable: x
    Var(String),
    /// Constructor: C p1 ... pn
    Constructor {
        name: String,
        args: Vec<CorePattern>,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub enum CoreBinOp {
    Add,
    Sub,
    Mul,
    Div,
    Lt,
    Le,
}
impl fmt::Display for Kind {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Kind::Star => write!(f, "*"),
            Kind::Arrow(k1, k2) => write!(f, "{} -> {}", k1, k2),
        }
    }
}
impl fmt::Display for CoreType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CoreType::Var(name) => write!(f, "{}", name),
            CoreType::ETVar(name) => write!(f, "^{}", name),
            CoreType::Con(name) => write!(f, "{}", name),
            CoreType::Arrow(t1, t2) => write!(f, "{} -> {}", t1, t2),
            CoreType::Forall(var, ty) => write!(f, "∀{}. {}", var, ty),
            CoreType::App(t1, t2) => write!(f, "({} {})", t1, t2),
            CoreType::Product(types) => {
                write!(f, "(")?;
                for (i, ty) in types.iter().enumerate() {
                    if i > 0 {
                        write!(f, " × ")?;
                    }
                    write!(f, "{}", ty)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl fmt::Display for CoreBinOp {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CoreBinOp::Add => write!(f, "+"),
            CoreBinOp::Sub => write!(f, "-"),
            CoreBinOp::Mul => write!(f, "*"),
            CoreBinOp::Div => write!(f, "/"),
            CoreBinOp::Lt => write!(f, "<"),
            CoreBinOp::Le => write!(f, "<="),
        }
    }
}
impl fmt::Display for CoreTerm {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CoreTerm::Var(name) => write!(f, "{}", name),
            CoreTerm::LitInt(n) => write!(f, "{}", n),
            CoreTerm::Lambda {
                param,
                param_ty,
                body,
            } => {
                write!(f, "λ{} : {}. {}", param, param_ty, body)
            }
            CoreTerm::App { func, arg } => write!(f, "{} {}", func, arg),
            CoreTerm::Constructor { name, args } => {
                if args.is_empty() {
                    write!(f, "{}", name)
                } else {
                    write!(f, "{}", name)?;
                    for arg in args {
                        write!(f, " {}", arg)?;
                    }
                    Ok(())
                }
            }
            CoreTerm::If {
                cond,
                then_branch,
                else_branch,
            } => {
                write!(f, "if {} then {} else {}", cond, then_branch, else_branch)
            }
            CoreTerm::Case { scrutinee, arms: _ } => {
                write!(f, "match {} {{ ... }}", scrutinee)
            }
            CoreTerm::BinOp { op, left, right } => {
                write!(f, "{} {} {}", left, op, right)
            }
            _ => write!(f, "<term>"),
        }
    }
}
/// Core types (System-F-ω)
#[derive(Debug, Clone, PartialEq)]
pub enum CoreType {
    /// Type variable: a
    Var(String),
    /// Existential type variable: ^a
    ETVar(String),
    /// Type constructor: Int, Bool
    Con(String),
    /// Function type: T1 -> T2
    Arrow(Box<CoreType>, Box<CoreType>),
    /// Universal quantification: ∀a. T
    Forall(String, Box<CoreType>),
    /// Type application: F T
    App(Box<CoreType>, Box<CoreType>),
    /// Product type: T1 × T2
    Product(Vec<CoreType>),
}
}

Core types include all the expressive power of System F-ω:

Type Variables: Both ordinary (Var) and existential (ETVar) variables for unification Type Constructors: Built-in types (Con) like Int, Bool Function Types: Explicit arrow types (Arrow) Universal Quantification: Forall types with kind-annotated bound variables Type Application: App for applying type constructors to arguments Type Abstraction: TAbs for creating type-level functions

The core representation makes explicit all type-level computation that remains implicit in the surface language.

Core Terms

#![allow(unused)]
fn main() {
use std::fmt;
/// Core language for System-F-ω with explicit type applications and
/// abstractions
#[derive(Debug, Clone, PartialEq)]
pub struct CoreModule {
    pub type_defs: Vec<TypeDef>,
    pub term_defs: Vec<TermDef>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TypeDef {
    pub name: String,
    pub kind: Kind,
    pub constructors: Vec<DataConstructor>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct DataConstructor {
    pub name: String,
    pub ty: CoreType,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TermDef {
    pub name: String,
    pub ty: CoreType,
    pub body: CoreTerm,
}
/// Kinds for types
#[derive(Debug, Clone, PartialEq)]
pub enum Kind {
    /// Base kind: *
    Star,
    /// Function kind: k1 -> k2
    Arrow(Box<Kind>, Box<Kind>),
}
/// Core types (System-F-ω)
#[derive(Debug, Clone, PartialEq)]
pub enum CoreType {
    /// Type variable: a
    Var(String),
    /// Existential type variable: ^a
    ETVar(String),
    /// Type constructor: Int, Bool
    Con(String),
    /// Function type: T1 -> T2
    Arrow(Box<CoreType>, Box<CoreType>),
    /// Universal quantification: ∀a. T
    Forall(String, Box<CoreType>),
    /// Type application: F T
    App(Box<CoreType>, Box<CoreType>),
    /// Product type: T1 × T2
    Product(Vec<CoreType>),
}
#[derive(Debug, Clone, PartialEq)]
pub struct CaseArm {
    pub pattern: CorePattern,
    pub body: CoreTerm,
}
#[derive(Debug, Clone, PartialEq)]
pub enum CorePattern {
    /// Wildcard: _
    Wildcard,
    /// Variable: x
    Var(String),
    /// Constructor: C p1 ... pn
    Constructor {
        name: String,
        args: Vec<CorePattern>,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub enum CoreBinOp {
    Add,
    Sub,
    Mul,
    Div,
    Lt,
    Le,
}
impl fmt::Display for Kind {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Kind::Star => write!(f, "*"),
            Kind::Arrow(k1, k2) => write!(f, "{} -> {}", k1, k2),
        }
    }
}
impl fmt::Display for CoreType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CoreType::Var(name) => write!(f, "{}", name),
            CoreType::ETVar(name) => write!(f, "^{}", name),
            CoreType::Con(name) => write!(f, "{}", name),
            CoreType::Arrow(t1, t2) => write!(f, "{} -> {}", t1, t2),
            CoreType::Forall(var, ty) => write!(f, "∀{}. {}", var, ty),
            CoreType::App(t1, t2) => write!(f, "({} {})", t1, t2),
            CoreType::Product(types) => {
                write!(f, "(")?;
                for (i, ty) in types.iter().enumerate() {
                    if i > 0 {
                        write!(f, " × ")?;
                    }
                    write!(f, "{}", ty)?;
                }
                write!(f, ")")
            }
        }
    }
}
impl fmt::Display for CoreBinOp {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CoreBinOp::Add => write!(f, "+"),
            CoreBinOp::Sub => write!(f, "-"),
            CoreBinOp::Mul => write!(f, "*"),
            CoreBinOp::Div => write!(f, "/"),
            CoreBinOp::Lt => write!(f, "<"),
            CoreBinOp::Le => write!(f, "<="),
        }
    }
}
impl fmt::Display for CoreTerm {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CoreTerm::Var(name) => write!(f, "{}", name),
            CoreTerm::LitInt(n) => write!(f, "{}", n),
            CoreTerm::Lambda {
                param,
                param_ty,
                body,
            } => {
                write!(f, "λ{} : {}. {}", param, param_ty, body)
            }
            CoreTerm::App { func, arg } => write!(f, "{} {}", func, arg),
            CoreTerm::Constructor { name, args } => {
                if args.is_empty() {
                    write!(f, "{}", name)
                } else {
                    write!(f, "{}", name)?;
                    for arg in args {
                        write!(f, " {}", arg)?;
                    }
                    Ok(())
                }
            }
            CoreTerm::If {
                cond,
                then_branch,
                else_branch,
            } => {
                write!(f, "if {} then {} else {}", cond, then_branch, else_branch)
            }
            CoreTerm::Case { scrutinee, arms: _ } => {
                write!(f, "match {} {{ ... }}", scrutinee)
            }
            CoreTerm::BinOp { op, left, right } => {
                write!(f, "{} {} {}", left, op, right)
            }
            _ => write!(f, "<term>"),
        }
    }
}
/// Core terms (System-F)
#[derive(Debug, Clone, PartialEq)]
pub enum CoreTerm {
    /// Variable: x
    Var(String),
    /// Integer literal: 42
    LitInt(i64),
    /// Lambda abstraction: λx:T. e
    Lambda {
        param: String,
        param_ty: CoreType,
        body: Box<CoreTerm>,
    },
    /// Application: e1 e2
    App {
        func: Box<CoreTerm>,
        arg: Box<CoreTerm>,
    },
    /// Type abstraction: Λα. e
    TypeLambda { param: String, body: Box<CoreTerm> },
    /// Constructor: C e1 ... en
    Constructor { name: String, args: Vec<CoreTerm> },
    /// Pattern matching: case e of { p1 -> e1; ... }
    Case {
        scrutinee: Box<CoreTerm>,
        arms: Vec<CaseArm>,
    },
    /// Built-in operations
    BinOp {
        op: CoreBinOp,
        left: Box<CoreTerm>,
        right: Box<CoreTerm>,
    },
    /// Conditional: if e1 then e2 else e3
    If {
        cond: Box<CoreTerm>,
        then_branch: Box<CoreTerm>,
        else_branch: Box<CoreTerm>,
    },
}
}

Core terms expose the full System F-ω term language:

Variables and Literals: Direct correspondence with surface language Function Abstraction/Application: Explicitly typed lambda calculus Type Abstraction: TyAbs for creating polymorphic functions Type Application: TyApp for instantiating polymorphic functions Data Constructors: Con for algebraic data type constructors Pattern Matching: Match with explicit constructor patterns

Every implicit type operation from the surface language becomes explicit in the core representation.

Pattern Matching and Case Analysis

Pattern matching provides the fundamental mechanism for deconstructing algebraic data types and extracting their constituent values. Our implementation supports comprehensive pattern matching that integrates seamlessly with the type system, ensuring that all pattern analyses are both exhaustive and type-safe.

The pattern matching construct match e { p1 -> e1; p2 -> e2; } performs case analysis on the scrutinee expression e, attempting to match it against each pattern in sequence. When a pattern matches, the corresponding branch expression executes with pattern variables bound to the extracted values.

Constructor Patterns and Variable Binding

#![allow(unused)]
fn main() {
use std::fmt;
#[derive(Debug, Clone, PartialEq)]
pub struct Module {
    pub declarations: Vec<Declaration>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Declaration {
    /// Data type declaration: data Color = Red | Blue
    Data {
        name: String,
        type_params: Vec<String>,
        constructors: Vec<Constructor>,
    },
    /// Function type signature: fib :: Int -> Int
    TypeSig {
        name: String,
        type_scheme: TypeScheme,
    },
    /// Function definition: fib n = match n { 0 -> 0, ... }
    FunDef {
        name: String,
        params: Vec<String>,
        body: Expr,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct Constructor {
    pub name: String,
    pub fields: Vec<Type>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TypeScheme {
    pub quantified: Vec<String>, // forall a b. ...
    pub ty: Type,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    /// Type variable: a, b
    Var(String),
    /// Type constructor: Int, Bool, Color
    Con(String),
    /// Function type: a -> b
    Arrow(Box<Type>, Box<Type>),
    /// Type application: List a, Maybe Int
    App(Box<Type>, Box<Type>),
    /// Record type: { x :: Int, y :: Int }
    Record(Vec<(String, Type)>),
    /// Forall type: forall a b. a -> b
    Forall(Vec<String>, Box<Type>),
}
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    /// Variable: x, y
    Var(String),
    /// Integer literal: 42
    LitInt(i64),
    /// Lambda: \x -> e
    Lambda { param: String, body: Box<Expr> },
    /// Function application: f x
    App { func: Box<Expr>, arg: Box<Expr> },
    /// Conditional: if cond then e1 else e2
    If {
        cond: Box<Expr>,
        then_branch: Box<Expr>,
        else_branch: Box<Expr>,
    },
    /// Pattern matching: match e { p1 -> e1, p2 -> e2 }
    Match {
        expr: Box<Expr>,
        arms: Vec<MatchArm>,
    },
    /// Binary operation: e1 + e2
    BinOp {
        op: BinOp,
        left: Box<Expr>,
        right: Box<Expr>,
    },
}
#[derive(Debug, Clone, PartialEq)]
pub struct MatchArm {
    pub pattern: Pattern,
    pub body: Expr,
}
#[derive(Debug, Clone, PartialEq)]
pub enum BinOp {
    Add,
    Sub,
    Mul,
    Div,
    // Comparison operations
    Lt,
    Le,
}
impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Type::Var(name) | Type::Con(name) => write!(f, "{}", name),
            Type::Arrow(t1, t2) => write!(f, "{} -> {}", t1, t2),
            Type::App(t1, t2) => write!(f, "{} {}", t1, t2),
            Type::Record(fields) => {
                write!(f, "{{ ")?;
                for (i, (name, ty)) in fields.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{} :: {}", name, ty)?;
                }
                write!(f, " }}")
            }
            Type::Forall(vars, ty) => {
                write!(f, "forall ")?;
                for (i, var) in vars.iter().enumerate() {
                    if i > 0 {
                        write!(f, " ")?;
                    }
                    write!(f, "{}", var)?;
                }
                write!(f, ". {}", ty)
            }
        }
    }
}
impl fmt::Display for TypeScheme {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.quantified.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "forall {}. {}", self.quantified.join(" "), self.ty)
        }
    }
}
impl fmt::Display for BinOp {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            BinOp::Add => write!(f, "+"),
            BinOp::Sub => write!(f, "-"),
            BinOp::Mul => write!(f, "*"),
            BinOp::Div => write!(f, "/"),
            BinOp::Lt => write!(f, "<"),
            BinOp::Le => write!(f, "<="),
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum Pattern {
    /// Wildcard pattern: _
    Wildcard,
    /// Variable pattern: x
    Var(String),
    /// Constructor pattern: Red, Cons x xs
    Constructor { name: String, args: Vec<Pattern> },
}
}

Constructor patterns decompose algebraic data types by matching against specific constructors and binding their arguments to pattern variables. The pattern Cons x xs matches values constructed with Cons, binding the first argument to x and the second to xs. This binding mechanism provides type-safe access to the components of structured data.

Variable patterns like x match any value and bind the entire matched value to the variable name. The pattern matching compiler ensures that each variable binding receives the appropriate type based on the context in which the pattern appears.

Wildcard patterns represented by _ match any value without creating bindings, useful for ignoring components that are not needed in the branch expression. The type checker verifies that wildcard patterns are used consistently with the expected type structure.

Exhaustiveness and Type Safety

The pattern matching implementation enforces exhaustiveness, requiring that pattern sets cover all possible values of the matched type. For algebraic data types, this means providing patterns for every constructor defined in the type declaration. The compiler rejects programs with non-exhaustive patterns, preventing runtime errors that could occur when unhandled cases arise.

Type safety extends through pattern matching by ensuring that pattern variables receive types consistent with the constructor arguments they represent. When matching Just x against a Maybe Int, the variable x automatically receives type Int, eliminating the need for explicit type annotations or runtime type checks.

Nested Patterns and Deep Matching

Patterns can nest arbitrarily deeply, enabling decomposition of complex data structures in single pattern matches. The pattern Cons (Just x) xs simultaneously matches the outer list structure and the inner Maybe type, binding both the unwrapped value x and the remaining list xs in a single operation.

Nested pattern matching interacts correctly with polymorphism, maintaining type relationships across multiple levels of structure. The type checker propagates type information through nested patterns, ensuring that all bindings receive their most general types while maintaining compatibility with the overall pattern context.

Algebraic Data Types

Algebraic data types provide the foundation for structured data in our System F-ω implementation, combining sum types (disjoint unions) and product types (tuples and records) into a unified framework that supports both data abstraction and generic programming.

The data declaration syntax enables concise specification of complex type structures while automatically deriving the associated constructors, destructors, and type information needed for pattern matching and type checking.

Sum Types and Tagged Unions

Sum types represent choices between alternative data representations, with each alternative identified by a unique constructor tag. The declaration data Either a b = Left a | Right b creates a sum type with two alternatives, each carrying a value of a different type.

Sum types enable type-safe representation of optional values, error conditions, and other scenarios where data can take one of several forms. The type Maybe a = Nothing | Just a encapsulates the common pattern of values that might be absent, replacing null pointer patterns with type-safe alternatives.

Each constructor in a sum type creates values that are distinguishable through pattern matching, enabling exhaustive case analysis that the type checker can verify statically. The tag information embedded in sum type values allows pattern matching to dispatch correctly to the appropriate branch without runtime type inspection.

Product Types and Data Aggregation

Product types combine multiple values into single composite structures, with each component accessible through pattern matching or projection operations. Constructor syntax like Cons a (List a) creates product types where each constructor argument represents a field in the resulting structure.

Product types support both named and positional field access through pattern matching, providing flexibility in how composite data gets decomposed. The pattern Cons head tail extracts both components of a list cell, binding them to appropriately typed variables for use in the branch expression.

Tuples represent anonymous product types where component ordering determines access patterns. While our surface language does not include explicit tuple syntax, the pattern matching mechanism supports tuple-like destructuring of constructor arguments.

Recursive Types and Inductive Data Structures

Recursive type definitions enable the construction of arbitrarily large data structures through self-reference in constructor arguments. The declaration data List a = Nil | Cons a (List a) defines lists inductively, with the base case Nil and the recursive case Cons that references the type being defined.

Recursive types interact correctly with polymorphism, enabling the definition of generic container types that work uniformly across all element types. The list type List a demonstrates how recursive structure combines with parametric polymorphism to create flexible, reusable data abstractions.

Inductive types support well-founded recursion through pattern matching, enabling terminating recursive functions that process arbitrarily large data structures. The pattern matching compiler ensures that recursive calls operate on structurally smaller arguments, supporting termination analysis and optimization.

Kind Inference and Type Constructor Classification

Data type declarations automatically infer appropriate kinds for the defined type constructors based on their parameter structure and usage patterns. Simple types like Bool receive kind *, while parameterized types like Maybe receive kinds of the form * -> *.

Higher-kinded types emerge naturally from data declarations with multiple parameters or higher-order structure. The type Either receives kind * -> * -> *, indicating a type constructor that requires two type arguments to produce a complete type.

Kind inference propagates through type expressions, ensuring that type applications in data constructors receive appropriate kind annotations for use in the core language representation. This automatic kind inference eliminates the need for explicit kind annotations while maintaining the precision required for System F-ω’s type-level computation.

Elaboration Process

Elaboration forms the critical bridge between the user-friendly surface language and the theoretically precise core language, transforming high-level programming constructs into explicit System F-ω representations that enable sound type checking and compilation. This translation process handles the complex task of inferring and inserting all implicit type-level operations that the surface language deliberately omits for conciseness and programmer convenience.

The elaboration algorithm operates through multiple interdependent phases that work together to produce well-typed core language programs. Each phase builds upon the results of previous phases, creating a pipeline that systematically transforms surface constructs into their explicit core equivalents while maintaining type safety and semantic preservation.

Kind Inference and Type Constructor Analysis

Kind inference represents one of the most aspects of elaboration, requiring analysis of type constructor usage patterns to determine their proper classification in the kind hierarchy. The process begins by analyzing type parameters in data declarations to determine the kinds of type constructors being defined. A simple data type like Bool with no parameters receives kind \( \star \), indicating it represents a complete type that can classify terms.

Parameterized data types require more complex analysis to determine their proper kinds. The declaration data Maybe a = Nothing | Just a reveals that Maybe is a type constructor that takes one type argument, yielding kind \( \star \to \star \). Multi-parameter types like Either a b receive kinds of the form \( \star \to \star \to \star \), reflecting their need for multiple type arguments before producing complete types.

The kind inference algorithm propagates kind constraints through type applications and signatures, ensuring consistency across the entire program. When a type constructor appears in a type application like Maybe Int, the algorithm verifies that the argument Int has kind \( \star \) to match the expected parameter kind of Maybe. This constraint propagation catches kind errors early in the elaboration process, preventing malformed types from reaching the core language.

The final step inserts implicit kind abstractions in the core representation, making explicit the kind-level lambda abstractions that correspond to parameterized type constructors. The surface type Maybe becomes a kind-level function \( \lambda\alpha :: \star. \text{Maybe}\, \alpha \) in the core language, exposing the type-level computation that remains hidden in the surface syntax.

Type Inference and Polymorphic Instantiation

Surface programs deliberately omit many type-level operations to maintain readability and reduce annotation burden, requiring the elaboration process to infer and insert these operations automatically. The most complex aspect involves handling polymorphic functions, which use implicit universal quantification in the surface language but require explicit type abstractions and applications in the core.

Type abstraction insertion analyzes function definitions to identify polymorphic variables that must be abstracted at the type level. A surface function like identity x = x with inferred type a -> a becomes a core expression \( \Lambda\alpha :: \star. \lambda x : \alpha. x \), making explicit the type-level abstraction over the polymorphic variable a. This transformation ensures that the core representation captures the full generality of polymorphic functions.

Type application insertion occurs at call sites of polymorphic functions, where the surface language relies on type inference to determine appropriate instantiations. When a polymorphic function like identity is applied to a specific argument like 42, the elaboration process inserts the type application \( \text{identity}[\text{Int}]\, 42 \), making explicit the instantiation of the type parameter with Int.

The generation of existential variables handles unknown types that arise during inference, creating placeholders that the constraint solver can resolve later. When the elaboration process encounters expressions whose types cannot be determined immediately, it generates fresh existential variables and creates constraints that guide the solver toward appropriate instantiations.

Pattern Compilation and Constructor Analysis

Pattern matching in the surface language undergoes substantial transformation during elaboration, converting high-level pattern matching constructs into explicit case analysis that the core language can represent directly. This compilation process must handle the complex interactions between pattern structure, type information, and variable binding that make pattern matching both powerful and type-safe.

Constructor pattern analysis examines each pattern to determine the type of the scrutinee expression and the types of bound variables. A pattern like Cons x xs matching against List Int reveals that x has type Int and xs has type List Int. This type information gets embedded in the core representation to guide type checking and code generation.

The generation of core match expressions creates explicit case analysis constructs that enumerate all possible constructor alternatives. Surface pattern matches become core expressions of the form \( \text{match}\, e\, \{ \text{C}_1\, x_1 \to e_1; \ldots; \text{C}_n\, x_n \to e_n \} \), with each alternative explicitly typed and all constructor possibilities accounted for.

Pattern variable binding requires careful scope management to ensure that variables bound in patterns are available with the correct types in branch expressions. The elaboration process maintains binding contexts that track the types of pattern variables and ensures that core expressions correctly reference these bindings with appropriate type information.

Exhaustiveness checking during pattern compilation verifies that pattern sets cover all possible constructor alternatives, preventing runtime errors from unhandled cases. The elaboration process analyzes constructor declarations to determine the complete set of alternatives and reports errors when patterns are non-exhaustive, maintaining the type safety guarantees that System F-ω provides.

Design Rationale

The two-layer architecture of our System F-ω implementation represents a carefully considered approach to balancing theoretical precision with practical usability, addressing fundamental tensions that arise when building advanced type systems for real-world programming. This design philosophy emerges from the recognition that pure System F-ω, while theoretically elegant, imposes significant annotation burden that makes it impractical for everyday programming tasks.

Programmer Productivity and Cognitive Load Management

The surface language prioritizes programmer productivity by eliminating the explicit type-level operations that System F-ω requires while preserving all the expressive power of the underlying type system. Programmers can write natural, intuitive code using familiar algebraic data types and pattern matching without being forced to understand or manipulate higher-kinded types, type applications, or kind annotations directly.

This approach recognizes that cognitive load represents a scarce resource in software development. By hiding the complexity of type-level computation behind a clean surface syntax, we enable programmers to focus on problem-solving rather than wrestling with the mechanical details of type system operation. The surface language provides enough abstraction that polymorphic programming feels natural and obvious, even though the underlying elaboration process involves type inference and constraint solving.

The implicit quantification system exemplifies this philosophy. Where pure System F-ω requires explicit type abstractions like \( \Lambda\alpha :: \star. \lambda x : \alpha. x \), our surface language allows the simple definition identity x = x with automatic inference of the polymorphic type a -> a. This transformation eliminates tedious annotation while preserving full generality and type safety.

Type Safety and Theoretical Soundness

The core language ensures that all the theoretical guarantees of System F-ω remain intact by making every type operation explicit and checkable. This explicit representation enables rigorous type checking algorithms that can verify program correctness with mathematical precision, preventing the subtle errors that can arise when implicit operations are handled incorrectly.

By elaborating surface programs to core representations, we gain access to the full theoretical machinery of System F-ω, including decidable type checking, principal types, and strong normalization. The core language serves as a certificate of correctness, providing concrete evidence that surface programs satisfy all the constraints of the type system.

The explicit nature of core representations also enables advanced optimizations and program transformations that rely on precise type information. Code generators can exploit type-level information to produce efficient implementations, while program analysis tools can reason about program behavior with greater precision than would be possible with surface-level representations alone.

Modularity and Language Evolution

The separation between surface and core languages provides crucial modularity that enables independent evolution of user-facing features and theoretical foundations. Surface language features can be added, modified, or removed without affecting the core type checking algorithms, provided the elaboration process can translate them to appropriate core representations.

This modularity proves essential for language experimentation and extension. New surface constructs like syntactic sugar, alternative pattern matching styles, or domain-specific language features can be implemented entirely through elaboration without requiring changes to the type checker or core language semantics. This flexibility enables rapid prototyping of language features while maintaining implementation stability.

The core language can similarly evolve independently, with improvements to type checking algorithms, optimization passes, or code generation strategies affecting all surface programs automatically through the existing elaboration interface. This separation enables focused development where surface language designers can concentrate on usability while type theorists can focus on correctness and efficiency.

Debugging and Developer Understanding

The explicit core representations provide invaluable insight into type checker behavior, enabling developers to understand how their programs are interpreted by the type system. When type errors occur, the core representation shows exactly which type-level operations failed and why, providing much more precise diagnostic information than surface-level error reporting alone could offer.

This debugging capability extends to performance analysis, where core representations reveal the type-level computation overhead imposed by different programming patterns. Developers can examine core representations to understand which surface constructs generate expensive type-level operations and adjust their programming style accordingly.

The elaboration process itself serves as an educational tool, demonstrating how high-level programming constructs decompose into fundamental type theoretic operations. Students learning advanced type systems can examine the core representations of their programs to develop intuition about how polymorphism, higher-kinded types, and type-level computation actually work.

Balancing Expressiveness with Accessibility

Our design demonstrates that advanced type systems can be made accessible without sacrificing their theoretical foundations. The surface language captures the essential patterns that programmers want to express - polymorphic functions, generic data structures, type-safe pattern matching - while the core language ensures that these patterns have rigorous foundations in type theory.

This balance addresses a fundamental challenge in programming language design: how to provide the benefits of advanced type systems without requiring all programmers to become experts in type theory. By carefully choosing which aspects to make implicit and which to keep explicit, we create a system that scales from beginning programmers writing simple functions to expert developers implementing complex generic libraries.

The success of this approach validates the broader principle that programming language design should prioritize human factors alongside theoretical considerations. Technical elegance and mathematical precision remain essential, but they must be balanced against the practical realities of software development where programmer time, cognitive load, and learning curves represent real constraints that affect the ultimate utility of programming languages.

Through this two-layer architecture, we achieve a System F-ω implementation that provides the full power of higher-kinded polymorphism while remaining approachable for programmers who simply want to write correct, generic, type-safe code without becoming experts in the theoretical foundations that make such code possible.

Lexer and Parser

Our System F-ω implementation employs a carefully designed parsing strategy that contrasts sharply with the organically evolved complexity found in production Haskell implementations. While Haskell’s grammar has grown over decades through endless language extensions and “pragmatic” compromises, resulting in context-dependent parsing, whitespace sensitivity, and intricate layout rules with virtual braces and semicolons, and a general clusterf*** our implementation opts for a much smaller surface langauge.

The Haskell language specification presents significant parsing challenges that have led to numerous implementation variations across different compilers. The language’s sensitivity to indentation creates a complex interaction between lexical analysis and parsing phases, where layout rules must insert implicit structure markers. Context-dependent syntax means that identical token sequences can parse differently depending on surrounding declarations, while the proliferation of language extensions has created a patchwork of parsing rules that interact in unexpected ways.

Our toy System F-ω implementation deliberately sidesteps these complexities by adopting explicit syntax with clear delimiters, enabling straightforward LALR(1) parsing that produces predictable results independent of context or whitespace variations.

Lexical Analysis Strategy

The lexical analysis phase transforms source text into a stream of tokens that the parser can process deterministically. Unlike Haskell’s context-sensitive lexer that must track indentation levels and insert layout tokens, our lexer operates as a pure finite state machine.

#![allow(unused)]
fn main() {
use std::fmt;
use logos::Logos;
pub struct Lexer<'input> {
    token_stream: logos::SpannedIter<'input, Token>,
    last_token: Option<Token>,
    pending_token: Option<(usize, Token, usize)>, // Token to emit after semicolon
    just_inserted_semicolon: bool,                // Track if we just inserted a semicolon
}
impl<'input> Lexer<'input> {
    pub fn new(input: &'input str) -> Self {
        Self {
            token_stream: Token::lexer(input).spanned(),
            last_token: None,
            pending_token: None,
            just_inserted_semicolon: false,
        }
    }
}
impl fmt::Display for Token {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Token::Data => write!(f, "data"),
            Token::Forall => write!(f, "forall"),
            Token::Match => write!(f, "match"),
            Token::If => write!(f, "if"),
            Token::Then => write!(f, "then"),
            Token::Else => write!(f, "else"),
            Token::Arrow => write!(f, "->"),
            Token::DoubleColon => write!(f, "::"),
            Token::Equal => write!(f, "="),
            Token::Plus => write!(f, "+"),
            Token::Minus => write!(f, "-"),
            Token::Star => write!(f, "*"),
            Token::Slash => write!(f, "/"),
            Token::LessEqual => write!(f, "<="),
            Token::Less => write!(f, "<"),
            Token::Pipe => write!(f, "|"),
            Token::Underscore => write!(f, "_"),
            Token::Dot => write!(f, "."),
            Token::Comma => write!(f, ","),
            Token::Backslash => write!(f, "\\"),
            Token::LeftParen => write!(f, "("),
            Token::RightParen => write!(f, ")"),
            Token::LeftBrace => write!(f, "{{"),
            Token::RightBrace => write!(f, "}}"),
            Token::Ident(s) => write!(f, "{}", s),
            Token::Integer(n) => write!(f, "{}", n),
            Token::Semicolon => write!(f, ";"),
            Token::Newline => write!(f, "\\n"),
            Token::Error => write!(f, "ERROR"),
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum LexicalError {
    InvalidToken,
}
impl fmt::Display for LexicalError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            LexicalError::InvalidToken => write!(f, "Invalid token"),
        }
    }
}
pub type Spanned<Tok, Loc, Error> = Result<(Loc, Tok, Loc), Error>;
impl<'input> Iterator for Lexer<'input> {
    type Item = Spanned<Token, usize, LexicalError>;

    fn next(&mut self) -> Option<Self::Item> {
        // If we have a pending token, return it
        if let Some((start, token, end)) = self.pending_token.take() {
            self.last_token = Some(token.clone());
            self.just_inserted_semicolon = false;
            return Some(Ok((start, token, end)));
        }

        match self.token_stream.next()? {
            (Ok(Token::Newline), span) => {
                // Check if we should insert a semicolon before this newline
                // But don't insert if we just inserted one
                if !self.just_inserted_semicolon {
                    if let Some(ref last) = self.last_token {
                        if should_insert_semicolon(last) {
                            self.just_inserted_semicolon = true;
                            return Some(Ok((span.start, Token::Semicolon, span.start)));
                        }
                    }
                }
                // Skip newlines and continue
                self.just_inserted_semicolon = false;
                self.next()
            }
            (Ok(token), span) => {
                // Check for semicolon insertion at statement boundaries
                // But don't insert if we just inserted one
                if !self.just_inserted_semicolon {
                    if let Some(ref last) = self.last_token {
                        if should_insert_semicolon(last) && needs_semicolon_before(&token) {
                            // Store the current token for the next call, and return a semicolon
                            self.pending_token = Some((span.start, token, span.end));
                            self.just_inserted_semicolon = true;
                            return Some(Ok((span.start, Token::Semicolon, span.start)));
                        }
                    }
                }

                self.last_token = Some(token.clone());
                self.just_inserted_semicolon = false;
                Some(Ok((span.start, token, span.end)))
            }
            (Err(_), _span) => Some(Err(LexicalError::InvalidToken)),
        }
    }
}
fn should_insert_semicolon(token: &Token) -> bool {
    matches!(
        token,
        Token::Ident(_) | Token::Integer(_) | Token::RightParen | Token::RightBrace
    )
}
fn needs_semicolon_before(token: &Token) -> bool {
    matches!(token, Token::Data | Token::If | Token::Match)
    // Note: Removed Token::Ident(_) to prevent semicolon insertion in function
    // parameters
}
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_simple_tokens() {
        let lexer = Lexer::new("data Bool = True | False");
        let tokens: Vec<_> = lexer.collect();

        // Check that we get the expected tokens without errors
        assert!(tokens.iter().all(|t| t.is_ok()));

        let token_values: Vec<_> = tokens.into_iter().map(|t| t.unwrap().1).collect();
        assert_eq!(
            token_values,
            vec![
                Token::Data,
                Token::Ident("Bool".to_string()),
                Token::Equal,
                Token::Ident("True".to_string()),
                Token::Pipe,
                Token::Ident("False".to_string()),
            ]
        );
    }

    #[test]
    fn test_function_definition() {
        let lexer = Lexer::new("fib :: Int -> Int");
        let tokens: Vec<_> = lexer.map(|t| t.unwrap().1).collect();

        assert_eq!(
            tokens,
            vec![
                Token::Ident("fib".to_string()),
                Token::DoubleColon,
                Token::Ident("Int".to_string()),
                Token::Arrow,
                Token::Ident("Int".to_string()),
            ]
        );
    }

    #[test]
    fn test_arithmetic() {
        let lexer = Lexer::new("1 + 2 * 3 <= 42");
        let tokens: Vec<_> = lexer.map(|t| t.unwrap().1).collect();

        assert_eq!(
            tokens,
            vec![
                Token::Integer(1),
                Token::Plus,
                Token::Integer(2),
                Token::Star,
                Token::Integer(3),
                Token::LessEqual,
                Token::Integer(42),
            ]
        );
    }

    #[test]
    fn test_comments_ignored() {
        let lexer = Lexer::new("fib -- This is a comment\n:: Int");
        let tokens: Vec<_> = lexer.map(|t| t.unwrap().1).collect();

        assert_eq!(
            tokens,
            vec![
                Token::Ident("fib".to_string()),
                Token::Semicolon, // Auto-inserted at newline
                Token::DoubleColon,
                Token::Ident("Int".to_string()),
            ]
        );
    }

    #[test]
    fn test_auto_semicolon_insertion() {
        let lexer = Lexer::new("data Bool = True | False\nfib n = 42");
        let tokens: Vec<_> = lexer.map(|t| t.unwrap().1).collect();

        // Should insert semicolon after "False" before "fib"
        assert!(tokens.contains(&Token::Semicolon));

        // Check that we have the right tokens in order
        assert_eq!(tokens[0], Token::Data);
        assert_eq!(tokens[1], Token::Ident("Bool".to_string()));
        assert_eq!(tokens[2], Token::Equal);
        assert_eq!(tokens[3], Token::Ident("True".to_string()));
        assert_eq!(tokens[4], Token::Pipe);
        assert_eq!(tokens[5], Token::Ident("False".to_string()));
        assert_eq!(tokens[6], Token::Semicolon); // Auto-inserted
        assert_eq!(tokens[7], Token::Ident("fib".to_string()));
        assert_eq!(tokens[8], Token::Ident("n".to_string()));
        assert_eq!(tokens[9], Token::Equal);
        assert_eq!(tokens[10], Token::Integer(42));
    }

    #[test]
    fn test_function_parameter_tokens() {
        let lexer = Lexer::new("id x = x;");
        let tokens: Vec<_> = lexer.map(|t| t.unwrap().1).collect();

        println!("Function parameter tokens: {:?}", tokens);

        assert_eq!(tokens[0], Token::Ident("id".to_string()));
        assert_eq!(tokens[1], Token::Ident("x".to_string()));
        assert_eq!(tokens[2], Token::Equal);
        assert_eq!(tokens[3], Token::Ident("x".to_string()));
        assert_eq!(tokens[4], Token::Semicolon);
    }
}
#[derive(Logos, Debug, Clone, PartialEq)]
pub enum Token {
    // Keywords
    #[token("data")]
    Data,
    #[token("forall")]
    Forall,
    #[token("match")]
    Match,
    #[token("if")]
    If,
    #[token("then")]
    Then,
    #[token("else")]
    Else,

    // Operators and delimiters
    #[token("->")]
    Arrow,
    #[token("::")]
    DoubleColon,
    #[token("=")]
    Equal,
    #[token("+")]
    Plus,
    #[token("-")]
    Minus,
    #[token("*")]
    Star,
    #[token("/")]
    Slash,
    #[token("<=")]
    LessEqual,
    #[token("<")]
    Less,
    #[token("|")]
    Pipe,
    #[token("_", priority = 2)]
    Underscore,
    #[token(".")]
    Dot,
    #[token(",")]
    Comma,
    #[token("\\")]
    Backslash,

    // Delimiters
    #[token("(")]
    LeftParen,
    #[token(")")]
    RightParen,
    #[token("{")]
    LeftBrace,
    #[token("}")]
    RightBrace,

    // Identifiers and literals
    #[regex(r"[a-zA-Z][a-zA-Z0-9_']*", |lex| lex.slice().to_string(), priority = 1)]
    Ident(String),

    #[regex(r"[0-9]+", |lex| lex.slice().parse::<i64>().unwrap_or(0))]
    Integer(i64),

    // Automatic semicolon insertion
    #[token(";")]
    Semicolon,

    // Special tokens
    Newline,

    // Skip whitespace and comments
    #[regex(r"[ \t\f]+", logos::skip)]
    #[regex(r"--[^\r\n]*", logos::skip)]
    #[token("\r\n", |_| Token::Newline)]
    #[token("\n", |_| Token::Newline)]
    Error,
}
}

The lexer recognizes several categories of tokens that capture the essential elements of our surface language. Keywords like data, match, and forall establish the syntactic framework for declarations and expressions. Identifiers distinguish between type variables, term variables, and constructor names through naming conventions that the lexer enforces consistently.

Operators receive special treatment to handle the function arrow (->) and type signature marker (::), both crucial for expressing types and function signatures. Delimiters including parentheses, braces, and semicolons provide explicit structure that eliminates the ambiguity inherent in layout-based syntax.

Numeric and boolean literals complete the token vocabulary, providing the primitive values that serve as the foundation for more complex expressions.

This is not ALL of Haskell, but it’s enough to do non-trivial typechecking over. And that’s all we really need. The full langauge is, to put it mildly, a bit too large and “organic”.

Bidirectional Type Checking

The type inference engine represents the most component of our System F-ω implementation. Built around the DK (Dunfield-Krishnaswami) worklist algorithm, it handles the complex interactions between higher-rank polymorphism, existential type variables, and bidirectional type checking that make System F-ω both powerful and challenging to implement efficiently.

The algorithm operates on a worklist of judgments and constraints, systematically reducing complex type checking problems into simpler ones until all constraints are resolved. This approach provides decidable type inference for System F-ω while maintaining principal types and supporting polymorphic programming patterns.

Notably this approach is different than our direct tree-walking approach before. Now we walk the syntax tree and emit constraints, which are then solved by the worklist algorithm, and then back-propagate the results to the original syntax tree to give the final types. This system allows more copmlex rules to be introduced, but it also makes the entire process “non-local” in the sense that we are looking at a constraint generated by many different expressions, and thus do error reporting becomes a bit more challenging because tracing back a constraint problem to its original syntax requires quite a bit of bookkeeping

Typing Rules

Now we’re going to bodly go into the next generation of type checking. We have some new symbols in our type system, but they’re mostly the same as before except for the addition of new kind-level operations.

• \( \Gamma \) (Gamma) - The typing context that stores information about variables and their types.

• \( \vdash \) (Turnstile) - The “proves” relation, meaning “given this context, we can conclude this judgment.”

• \( \Rightarrow \) (Double Right Arrow) - Synthesis mode, where the type checker figures out what type an expression has.

• \( \Leftarrow \) (Double Left Arrow) - Checking mode, where we verify that an expression has the expected type.

• \( \forall\alpha \) (Forall Alpha) - Universal quantification, meaning “for any type α.”

• \( :: \) (Double Colon) - The “has kind” relation, telling us what category a type constructor belongs to.

• \( \star \) (Star) - The kind of concrete types like Int, Bool, and String.

• \( \square \) (Box) - The kind of kinds, the kind that kinds themselves have.

• \( \to \) (Arrow) - Function types at the term level, or kind arrows at the type level.

• \( [B/\alpha]A \) - Type substitution that replaces all occurrences of type variable α with type B in type A.

• \( \Lambda\alpha \) (Big Lambda) - Type abstraction that creates polymorphic functions at the type level.

• \( \kappa \) (Kappa) - Kind variables representing unknown categories in our type system.

• \( \overline{x} \) - Overline notation indicating a sequence or list of elements, such as multiple variables or arguments.

Before examining the DK worklist algorithm implementation, let’s establish the formal typing rules that govern System F-ω. These rules extend System F with higher-kinded types and more polymorphism.

Basic Expression Rules

Variable lookup finds types in the context:

\[ \frac{x : A \in \Gamma}{\Gamma \vdash x \Rightarrow A} \text{(T-Var)} \]

Function application with bidirectional checking:

\[ \frac{\Gamma \vdash e_1 \Rightarrow A \to B \quad \Gamma \vdash e_2 \Leftarrow A}{\Gamma \vdash e_1 ; e_2 \Rightarrow B} \text{(T-App)} \]

Lambda abstraction in checking mode:

\[ \frac{\Gamma, x : A \vdash e \Leftarrow B}{\Gamma \vdash \lambda x. e \Leftarrow A \to B} \text{(T-Lam)} \]

Type Application and Abstraction

Type application instantiates polymorphic expressions:

\[ \frac{\Gamma \vdash e \Rightarrow \forall \alpha :: \kappa. A \quad \Gamma \vdash B :: \kappa}{\Gamma \vdash e ; [B] \Rightarrow [B/\alpha]A} \text{(T-TApp)} \]

Type abstraction introduces universal quantification:

\[ \frac{\Gamma, \alpha :: \kappa \vdash e \Leftarrow A}{\Gamma \vdash \Lambda \alpha :: \kappa. e \Leftarrow \forall \alpha :: \kappa. A} \text{(T-TAbs)} \]

Higher-Kinded Type Rules

Kind checking ensures type constructors are well-formed:

\[ \frac{}{\Gamma \vdash \star :: \square} \text{(K-Type)} \]

\[ \frac{\Gamma \vdash \kappa_1 :: \square \quad \Gamma \vdash \kappa_2 :: \square}{\Gamma \vdash \kappa_1 \to \kappa_2 :: \square} \text{(K-Arrow)} \]

Type constructor application:

\[ \frac{\Gamma \vdash F :: \kappa_1 \to \kappa_2 \quad \Gamma \vdash A :: \kappa_1}{\Gamma \vdash F ; A :: \kappa_2} \text{(K-App)} \]

Pattern Matching Rules

Constructor patterns check against data types:

\[ \frac{\Gamma \vdash e \Rightarrow T \quad \overline{\Gamma \vdash p_i \Leftarrow T \leadsto \Delta_i \quad \Gamma, \Delta_i \vdash e_i \Leftarrow A}}{\Gamma \vdash \text{match } e \text{ with } \overline{p_i \to e_i} \Rightarrow A} \text{(T-Match)} \]

Bidirectional Control

Mode switching from inference to checking:

\[ \frac{\Gamma \vdash e \Rightarrow A}{\Gamma \vdash e \Leftarrow A} \text{(T-Sub)} \]

Existential variable instantiation:

\[ \frac{\Gamma, \hat{\alpha} :: \kappa \vdash e \Rightarrow A}{\Gamma \vdash e \Rightarrow [\hat{\alpha}/\alpha]A} \text{(T-InstL)} \]

In these rules, \( \Rightarrow \) denotes inference mode (synthesizing types), \( \Leftarrow \) denotes checking mode (verifying against expected types), \( :: \) relates types to their kinds, and \( \hat{\alpha} \) represents existential type variables solved during inference.

The DK Worklist Algorithm

The DK algorithm represents type inference as a constraint solving problem. Instead of recursively traversing expressions and immediately making type checking decisions, it accumulates constraints in a worklist and processes them systematically.

Worklist Structure

#![allow(unused)]
fn main() {
use std::collections::HashMap;
use crate::core::{CoreBinOp, CorePattern, CoreTerm, CoreType};
use crate::errors::{TypeError, TypeResult};
/// DK Worklist Algorithm for System-F-ω
/// Based on "A Mechanical Formalization of Higher-Ranked Polymorphic Type
/// Inference"
pub type TyVar = String;
pub type TmVar = String;
#[derive(Debug, Clone, PartialEq)]
pub enum TyVarKind {
    /// Universal type variable: α
    Universal,
    /// Existential type variable: ^α
    Existential,
    /// Solved existential: ^α = τ
    Solved(CoreType),
    /// Marker: ►α (for scoping)
    Marker,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Judgment {
    /// Subtyping: A <: B
    Sub { left: CoreType, right: CoreType },
    /// Type inference: e ⊢ A
    Inf { term: CoreTerm, ty: CoreType },
    /// Type checking: e ⇐ A
    Chk { term: CoreTerm, ty: CoreType },
    /// Application inference helper: A • e ⊢ C
    InfApp {
        func_ty: CoreType,
        arg: CoreTerm,
        result_ty: CoreType,
    },
}
#[derive(Debug, Clone)]
pub struct Worklist {
    entries: Vec<WorklistEntry>,
    next_var: usize,
}
impl Default for Worklist {
    fn default() -> Self {
        Self::new()
    }
}
impl Worklist {
    pub fn new() -> Self {
        Worklist {
            entries: Vec::new(),
            next_var: 0,
        }
    }

    pub fn fresh_var(&mut self) -> TyVar {
        let var = format!("α{}", self.next_var);
        self.next_var += 1;
        var
    }

    pub fn fresh_evar(&mut self) -> TyVar {
        let var = format!("^α{}", self.next_var);
        self.next_var += 1;
        var
    }

    pub fn push(&mut self, entry: WorklistEntry) {
        self.entries.push(entry);
    }

    pub fn pop(&mut self) -> Option<WorklistEntry> {
        self.entries.pop()
    }

    pub fn find_var(&self, name: &str) -> Option<&CoreType> {
        for entry in self.entries.iter().rev() {
            if let WorklistEntry::Var(var_name, ty) = entry {
                if var_name == name {
                    return Some(ty);
                }
            }
        }
        None
    }

    pub fn solve_evar(&mut self, name: &str, ty: CoreType) -> TypeResult<()> {
        for entry in self.entries.iter_mut() {
            if let WorklistEntry::TVar(var_name, kind) = entry {
                if var_name == name {
                    match kind {
                        TyVarKind::Existential => {
                            *kind = TyVarKind::Solved(ty);
                            return Ok(());
                        }
                        TyVarKind::Solved(_) => {
                            // Variable already solved, that's OK
                            return Ok(());
                        }
                        _ => {
                            // Skip universal variables, markers, etc.
                            continue;
                        }
                    }
                }
            }
        }
        Err(TypeError::UnboundVariable {
            name: name.to_string(),
            span: None,
        })
    }

    pub fn before(&self, a: &str, b: &str) -> bool {
        let mut pos_a = None;
        let mut pos_b = None;

        for (i, entry) in self.entries.iter().enumerate() {
            if let WorklistEntry::TVar(name, _) = entry {
                if name == a {
                    pos_a = Some(i);
                }
                if name == b {
                    pos_b = Some(i);
                }
            }
        }

        match (pos_a, pos_b) {
            (Some(pa), Some(pb)) => pa < pb,
            _ => false,
        }
    }
}
pub struct DKInference {
    worklist: Worklist,
    trace: Vec<String>,
    data_constructors: HashMap<String, CoreType>,
    /// Variable typing context for pattern-bound variables
    var_context: HashMap<String, CoreType>,
}
impl DKInference {
    pub fn with_context(
        data_constructors: HashMap<String, CoreType>,
        var_context: HashMap<String, CoreType>,
    ) -> Self {
        DKInference {
            worklist: Worklist::new(),
            trace: Vec::new(),
            data_constructors,
            var_context,
        }
    }

    pub fn check_type(&mut self, term: &CoreTerm, expected_ty: &CoreType) -> TypeResult<()> {
        self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
            term: term.clone(),
            ty: expected_ty.clone(),
        }));

        self.solve()
    }

    fn solve(&mut self) -> TypeResult<()> {
        while let Some(entry) = self.worklist.pop() {
            match entry {
                WorklistEntry::TVar(_, _) => {
                    // Skip variable bindings during processing
                    continue;
                }
                WorklistEntry::Var(_, _) => {
                    // Skip term variable bindings during processing
                    continue;
                }
                WorklistEntry::Judgment(judgment) => {
                    self.solve_judgment(judgment)?;
                }
            }
        }
        Ok(())
    }

    fn solve_judgment(&mut self, judgment: Judgment) -> TypeResult<()> {
        match judgment {
            Judgment::Sub { left, right } => self.solve_subtype(left, right),
            Judgment::Inf { term, ty } => self.solve_inference(term, ty),
            Judgment::Chk { term, ty } => self.solve_checking(term, ty),
            Judgment::InfApp {
                func_ty,
                arg,
                result_ty,
            } => self.solve_inf_app(func_ty, arg, result_ty),
        }
    }

    fn solve_subtype(&mut self, left: CoreType, right: CoreType) -> TypeResult<()> {
        self.trace.push(format!("Sub {} <: {}", left, right));

        if left == right {
            return Ok(());
        }

        match (&left, &right) {
            // Reflexivity
            (CoreType::Con(a), CoreType::Con(b)) if a == b => Ok(()),
            (CoreType::Var(a), CoreType::Var(b)) if a == b => Ok(()),
            (CoreType::ETVar(a), CoreType::ETVar(b)) if a == b => Ok(()),

            // Function subtyping (contravariant in argument, covariant in result)
            (CoreType::Arrow(a1, a2), CoreType::Arrow(b1, b2)) => {
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: *b1.clone(),
                    right: *a1.clone(),
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: *a2.clone(),
                    right: *b2.clone(),
                }));
                Ok(())
            }

            // Application subtyping (covariant in both components)
            (CoreType::App(a1, a2), CoreType::App(b1, b2)) => {
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: *a1.clone(),
                    right: *b1.clone(),
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: *a2.clone(),
                    right: *b2.clone(),
                }));
                Ok(())
            }

            // Forall right
            (_, CoreType::Forall(var, ty)) => {
                let fresh_var = self.worklist.fresh_var();
                self.worklist
                    .push(WorklistEntry::TVar(fresh_var.clone(), TyVarKind::Universal));
                let substituted_ty = self.substitute_type(var, &CoreType::Var(fresh_var), ty);
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left,
                    right: substituted_ty,
                }));
                Ok(())
            }

            // Forall left
            (CoreType::Forall(var, ty), _) => {
                let fresh_evar = self.worklist.fresh_evar();
                self.worklist
                    .push(WorklistEntry::TVar(fresh_evar.clone(), TyVarKind::Marker));
                self.worklist.push(WorklistEntry::TVar(
                    fresh_evar.clone(),
                    TyVarKind::Existential,
                ));
                let substituted_ty = self.substitute_type(var, &CoreType::ETVar(fresh_evar), ty);
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: substituted_ty,
                    right,
                }));
                Ok(())
            }

            // Existential variable instantiation
            (CoreType::ETVar(a), _) if !self.occurs_check(a, &right) => {
                self.instantiate_left(a, &right)
            }
            (_, CoreType::ETVar(a)) if !self.occurs_check(a, &left) => {
                self.instantiate_right(&left, a)
            }

            _ => Err(TypeError::SubtypingError {
                left,
                right,
                span: None,
            }),
        }
    }

    fn solve_inference(&mut self, term: CoreTerm, ty: CoreType) -> TypeResult<()> {
        self.trace
            .push(format!("Inf {} ⊢ {}", self.term_to_string(&term), ty));

        match term {
            CoreTerm::Var(name) => {
                // Check pattern variable context first
                if let Some(var_ty) = self.var_context.get(&name).cloned() {
                    self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                        left: var_ty,
                        right: ty,
                    }));
                    Ok(())
                } else if let Some(var_ty) = self.worklist.find_var(&name).cloned() {
                    self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                        left: var_ty,
                        right: ty,
                    }));
                    Ok(())
                } else if let Some(var_ty) = self.data_constructors.get(&name).cloned() {
                    // Check data constructors for constructor variables
                    self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                        left: var_ty,
                        right: ty,
                    }));
                    Ok(())
                } else {
                    Err(TypeError::UnboundVariable {
                        name: name.to_string(),
                        span: None,
                    })
                }
            }

            CoreTerm::LitInt(_) => {
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: CoreType::Con("Int".to_string()),
                    right: ty,
                }));
                Ok(())
            }

            CoreTerm::Lambda {
                param,
                param_ty,
                body,
            } => {
                let result_ty = CoreType::ETVar(self.worklist.fresh_evar());
                if let CoreType::ETVar(var_name) = &result_ty {
                    self.worklist.push(WorklistEntry::TVar(
                        var_name.clone(),
                        TyVarKind::Existential,
                    ));
                }

                let arrow_ty =
                    CoreType::Arrow(Box::new(param_ty.clone()), Box::new(result_ty.clone()));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: arrow_ty,
                    right: ty,
                }));

                self.worklist.push(WorklistEntry::Var(param, param_ty));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Inf {
                    term: *body,
                    ty: result_ty,
                }));

                Ok(())
            }

            CoreTerm::App { func, arg } => {
                let func_ty = CoreType::ETVar(self.worklist.fresh_evar());
                if let CoreType::ETVar(var_name) = &func_ty {
                    self.worklist.push(WorklistEntry::TVar(
                        var_name.clone(),
                        TyVarKind::Existential,
                    ));
                }

                self.worklist
                    .push(WorklistEntry::Judgment(Judgment::InfApp {
                        func_ty: func_ty.clone(),
                        arg: *arg,
                        result_ty: ty,
                    }));

                self.worklist.push(WorklistEntry::Judgment(Judgment::Inf {
                    term: *func,
                    ty: func_ty,
                }));

                Ok(())
            }

            CoreTerm::TypeLambda { param, body } => {
                let body_ty = CoreType::ETVar(self.worklist.fresh_evar());
                if let CoreType::ETVar(var_name) = &body_ty {
                    self.worklist.push(WorklistEntry::TVar(
                        var_name.clone(),
                        TyVarKind::Existential,
                    ));
                }

                let forall_ty = CoreType::Forall(param.clone(), Box::new(body_ty.clone()));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: forall_ty,
                    right: ty,
                }));

                self.worklist
                    .push(WorklistEntry::TVar(param, TyVarKind::Universal));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Inf {
                    term: *body,
                    ty: body_ty,
                }));

                Ok(())
            }

            CoreTerm::Constructor { name, args: _ } => {
                if let Some(constructor_ty) = self.data_constructors.get(&name) {
                    // Instantiate the constructor type with fresh existential variables
                    let instantiated_ty = self.instantiate_constructor_type(constructor_ty.clone());

                    self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                        left: instantiated_ty,
                        right: ty,
                    }));
                    Ok(())
                } else {
                    Err(TypeError::UnboundDataConstructor {
                        name: name.clone(),
                        span: None,
                    })
                }
            }

            CoreTerm::BinOp { op, left, right } => {
                let (left_ty, right_ty, result_ty) = self.infer_binop_types(&op);

                // Add any existential variables to the worklist
                self.add_etvars_to_worklist(&left_ty);
                self.add_etvars_to_worklist(&right_ty);
                self.add_etvars_to_worklist(&result_ty);

                // Check that the operands have the expected types
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term: *left,
                    ty: left_ty,
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term: *right,
                    ty: right_ty,
                }));

                // Check that the result type matches the expected type
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: result_ty,
                    right: ty,
                }));

                Ok(())
            }

            CoreTerm::If {
                cond,
                then_branch,
                else_branch,
            } => {
                // The condition must be Bool
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term: *cond,
                    ty: CoreType::Con("Bool".to_string()),
                }));

                // Both branches must have the same type as the expected result
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term: *then_branch,
                    ty: ty.clone(),
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term: *else_branch,
                    ty,
                }));

                Ok(())
            }

            CoreTerm::Case { scrutinee, arms } => {
                // Create a fresh type variable for the scrutinee
                let scrutinee_ty = CoreType::ETVar(self.worklist.fresh_evar());
                self.add_etvars_to_worklist(&scrutinee_ty);

                // Check that the scrutinee has the inferred type
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term: *scrutinee,
                    ty: scrutinee_ty.clone(),
                }));

                // Process each pattern arm
                for arm in arms {
                    // Check pattern constraints and get pattern variable bindings
                    let pattern_bindings = self.check_pattern(&arm.pattern, &scrutinee_ty)?;

                    // Add pattern variable bindings to worklist as regular variables
                    for (var_name, var_type) in pattern_bindings {
                        self.worklist.push(WorklistEntry::Var(var_name, var_type));
                    }

                    // Check the body with pattern variables in the worklist
                    self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                        term: arm.body.clone(),
                        ty: ty.clone(),
                    }));
                }

                Ok(())
            }
        }
    }

    fn solve_checking(&mut self, term: CoreTerm, ty: CoreType) -> TypeResult<()> {
        self.trace
            .push(format!("Chk {} ⇐ {}", self.term_to_string(&term), ty));

        match (&term, &ty) {
            (
                CoreTerm::Lambda {
                    param,
                    param_ty,
                    body,
                },
                CoreType::Arrow(expected_param, result_ty),
            ) => {
                // Check that parameter types match
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: param_ty.clone(),
                    right: *expected_param.clone(),
                }));

                self.worklist
                    .push(WorklistEntry::Var(param.clone(), param_ty.clone()));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term: *body.clone(),
                    ty: *result_ty.clone(),
                }));

                Ok(())
            }

            (_, CoreType::Forall(var, body_ty)) => {
                let fresh_var = self.worklist.fresh_var();
                self.worklist
                    .push(WorklistEntry::TVar(fresh_var.clone(), TyVarKind::Universal));
                let substituted_ty = self.substitute_type(var, &CoreType::Var(fresh_var), body_ty);
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term,
                    ty: substituted_ty,
                }));
                Ok(())
            }

            _ => {
                // Fallback: infer type and check subtyping
                let inferred_ty = CoreType::ETVar(self.worklist.fresh_evar());
                if let CoreType::ETVar(var_name) = &inferred_ty {
                    self.worklist.push(WorklistEntry::TVar(
                        var_name.clone(),
                        TyVarKind::Existential,
                    ));
                }

                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: inferred_ty.clone(),
                    right: ty,
                }));

                self.worklist.push(WorklistEntry::Judgment(Judgment::Inf {
                    term,
                    ty: inferred_ty,
                }));

                Ok(())
            }
        }
    }

    fn solve_inf_app(
        &mut self,
        func_ty: CoreType,
        arg: CoreTerm,
        result_ty: CoreType,
    ) -> TypeResult<()> {
        match func_ty {
            CoreType::Arrow(param_ty, ret_ty) => {
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: *ret_ty,
                    right: result_ty,
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term: arg,
                    ty: *param_ty,
                }));
                Ok(())
            }

            CoreType::Forall(var, body) => {
                let fresh_evar = self.worklist.fresh_evar();
                self.worklist.push(WorklistEntry::TVar(
                    fresh_evar.clone(),
                    TyVarKind::Existential,
                ));
                let substituted = self.substitute_type(&var, &CoreType::ETVar(fresh_evar), &body);
                self.worklist
                    .push(WorklistEntry::Judgment(Judgment::InfApp {
                        func_ty: substituted,
                        arg,
                        result_ty,
                    }));
                Ok(())
            }

            CoreType::ETVar(a) => {
                let param_ty_name = self.worklist.fresh_evar();
                let ret_ty_name = self.worklist.fresh_evar();
                let param_ty = CoreType::ETVar(param_ty_name.clone());
                let ret_ty = CoreType::ETVar(ret_ty_name.clone());

                // Add the fresh existential variables to the worklist
                self.worklist
                    .push(WorklistEntry::TVar(param_ty_name, TyVarKind::Existential));
                self.worklist
                    .push(WorklistEntry::TVar(ret_ty_name, TyVarKind::Existential));

                let arrow_ty =
                    CoreType::Arrow(Box::new(param_ty.clone()), Box::new(ret_ty.clone()));
                self.worklist.solve_evar(&a, arrow_ty)?;

                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: ret_ty,
                    right: result_ty,
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Chk {
                    term: arg,
                    ty: param_ty,
                }));

                Ok(())
            }

            _ => Err(TypeError::NotAFunction {
                ty: func_ty,
                span: None,
            }),
        }
    }

    fn instantiate_left(&mut self, var: &str, ty: &CoreType) -> TypeResult<()> {
        match ty {
            CoreType::ETVar(b) if self.worklist.before(var, b) => {
                self.worklist
                    .solve_evar(b, CoreType::ETVar(var.to_string()))?;
                Ok(())
            }
            CoreType::Arrow(t1, t2) => {
                let a1 = self.worklist.fresh_evar();
                let a2 = self.worklist.fresh_evar();
                let arrow_ty = CoreType::Arrow(
                    Box::new(CoreType::ETVar(a1.clone())),
                    Box::new(CoreType::ETVar(a2.clone())),
                );
                self.worklist.solve_evar(var, arrow_ty)?;

                self.worklist
                    .push(WorklistEntry::TVar(a1.clone(), TyVarKind::Existential));
                self.worklist
                    .push(WorklistEntry::TVar(a2.clone(), TyVarKind::Existential));

                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: *t1.clone(),
                    right: CoreType::ETVar(a1),
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: CoreType::ETVar(a2),
                    right: *t2.clone(),
                }));

                Ok(())
            }
            CoreType::App(t1, t2) => {
                let a1 = self.worklist.fresh_evar();
                let a2 = self.worklist.fresh_evar();
                let app_ty = CoreType::App(
                    Box::new(CoreType::ETVar(a1.clone())),
                    Box::new(CoreType::ETVar(a2.clone())),
                );
                self.worklist.solve_evar(var, app_ty)?;

                self.worklist
                    .push(WorklistEntry::TVar(a1.clone(), TyVarKind::Existential));
                self.worklist
                    .push(WorklistEntry::TVar(a2.clone(), TyVarKind::Existential));

                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: CoreType::ETVar(a1),
                    right: *t1.clone(),
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: CoreType::ETVar(a2),
                    right: *t2.clone(),
                }));

                Ok(())
            }
            CoreType::Forall(b, t) => {
                let fresh_var = self.worklist.fresh_var();
                self.worklist
                    .push(WorklistEntry::TVar(fresh_var.clone(), TyVarKind::Universal));
                let substituted = self.substitute_type(b, &CoreType::Var(fresh_var), t);
                self.instantiate_left(var, &substituted)
            }
            _ if self.is_monotype(ty) => {
                self.worklist.solve_evar(var, ty.clone())?;
                Ok(())
            }
            _ => Err(TypeError::InstantiationError {
                var: var.to_string(),
                ty: ty.clone(),
                span: None,
            }),
        }
    }

    fn instantiate_right(&mut self, ty: &CoreType, var: &str) -> TypeResult<()> {
        match ty {
            CoreType::ETVar(a) if self.worklist.before(var, a) => {
                self.worklist
                    .solve_evar(a, CoreType::ETVar(var.to_string()))?;
                Ok(())
            }
            CoreType::Arrow(t1, t2) => {
                let a1 = self.worklist.fresh_evar();
                let a2 = self.worklist.fresh_evar();
                let arrow_ty = CoreType::Arrow(
                    Box::new(CoreType::ETVar(a1.clone())),
                    Box::new(CoreType::ETVar(a2.clone())),
                );
                self.worklist.solve_evar(var, arrow_ty)?;

                self.worklist
                    .push(WorklistEntry::TVar(a1.clone(), TyVarKind::Existential));
                self.worklist
                    .push(WorklistEntry::TVar(a2.clone(), TyVarKind::Existential));

                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: CoreType::ETVar(a1),
                    right: *t1.clone(),
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: *t2.clone(),
                    right: CoreType::ETVar(a2),
                }));

                Ok(())
            }
            CoreType::App(t1, t2) => {
                let a1 = self.worklist.fresh_evar();
                let a2 = self.worklist.fresh_evar();
                let app_ty = CoreType::App(
                    Box::new(CoreType::ETVar(a1.clone())),
                    Box::new(CoreType::ETVar(a2.clone())),
                );
                self.worklist.solve_evar(var, app_ty)?;

                self.worklist
                    .push(WorklistEntry::TVar(a1.clone(), TyVarKind::Existential));
                self.worklist
                    .push(WorklistEntry::TVar(a2.clone(), TyVarKind::Existential));

                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: *t1.clone(),
                    right: CoreType::ETVar(a1),
                }));
                self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                    left: *t2.clone(),
                    right: CoreType::ETVar(a2),
                }));

                Ok(())
            }
            CoreType::Forall(a, t) => {
                let fresh_evar = self.worklist.fresh_evar();
                self.worklist
                    .push(WorklistEntry::TVar(fresh_evar.clone(), TyVarKind::Marker));
                self.worklist.push(WorklistEntry::TVar(
                    fresh_evar.clone(),
                    TyVarKind::Existential,
                ));
                let substituted = self.substitute_type(a, &CoreType::ETVar(fresh_evar), t);
                self.instantiate_right(&substituted, var)
            }
            _ if self.is_monotype(ty) => {
                self.worklist.solve_evar(var, ty.clone())?;
                Ok(())
            }
            _ => Err(TypeError::InstantiationError {
                var: var.to_string(),
                ty: ty.clone(),
                span: None,
            }),
        }
    }

    fn occurs_check(&self, var: &str, ty: &CoreType) -> bool {
        match ty {
            CoreType::ETVar(name) | CoreType::Var(name) => name == var,
            CoreType::Arrow(t1, t2) | CoreType::App(t1, t2) => {
                self.occurs_check(var, t1) || self.occurs_check(var, t2)
            }
            CoreType::Forall(_, t) => self.occurs_check(var, t),
            CoreType::Product(types) => types.iter().any(|t| self.occurs_check(var, t)),
            _ => false,
        }
    }

    fn is_monotype(&self, ty: &CoreType) -> bool {
        match ty {
            CoreType::Con(_) | CoreType::Var(_) | CoreType::ETVar(_) => true,
            CoreType::Arrow(t1, t2) | CoreType::App(t1, t2) => {
                self.is_monotype(t1) && self.is_monotype(t2)
            }
            CoreType::Product(types) => types.iter().all(|t| self.is_monotype(t)),
            CoreType::Forall(_, _) => false,
        }
    }

    fn substitute_type(&self, var: &str, replacement: &CoreType, ty: &CoreType) -> CoreType {
        match ty {
            CoreType::Var(name) if name == var => replacement.clone(),
            CoreType::Arrow(t1, t2) => CoreType::Arrow(
                Box::new(self.substitute_type(var, replacement, t1)),
                Box::new(self.substitute_type(var, replacement, t2)),
            ),
            CoreType::Forall(bound_var, body) if bound_var != var => CoreType::Forall(
                bound_var.clone(),
                Box::new(self.substitute_type(var, replacement, body)),
            ),
            CoreType::App(t1, t2) => CoreType::App(
                Box::new(self.substitute_type(var, replacement, t1)),
                Box::new(self.substitute_type(var, replacement, t2)),
            ),
            CoreType::Product(types) => CoreType::Product(
                types
                    .iter()
                    .map(|t| self.substitute_type(var, replacement, t))
                    .collect(),
            ),
            _ => ty.clone(),
        }
    }

    fn term_to_string(&self, term: &CoreTerm) -> String {
        match term {
            CoreTerm::Var(name) => name.clone(),
            CoreTerm::LitInt(n) => n.to_string(),
            _ => format!("{:?}", term), // Simplified for now
        }
    }

    pub fn get_trace(&self) -> &[String] {
        &self.trace
    }

    fn infer_binop_types(&mut self, op: &CoreBinOp) -> (CoreType, CoreType, CoreType) {
        match op {
            // Arithmetic operations: Int -> Int -> Int
            CoreBinOp::Add | CoreBinOp::Sub | CoreBinOp::Mul | CoreBinOp::Div => (
                CoreType::Con("Int".to_string()),
                CoreType::Con("Int".to_string()),
                CoreType::Con("Int".to_string()),
            ),
            // Comparison operations: Int -> Int -> Bool
            CoreBinOp::Lt | CoreBinOp::Le => (
                CoreType::Con("Int".to_string()),
                CoreType::Con("Int".to_string()),
                CoreType::Con("Bool".to_string()),
            ),
        }
    }

    fn add_etvars_to_worklist(&mut self, ty: &CoreType) {
        match ty {
            CoreType::ETVar(var_name) => {
                self.worklist.push(WorklistEntry::TVar(
                    var_name.clone(),
                    TyVarKind::Existential,
                ));
            }
            CoreType::Arrow(t1, t2) | CoreType::App(t1, t2) => {
                self.add_etvars_to_worklist(t1);
                self.add_etvars_to_worklist(t2);
            }
            CoreType::Forall(_, t) => {
                self.add_etvars_to_worklist(t);
            }
            CoreType::Product(types) => {
                for t in types {
                    self.add_etvars_to_worklist(t);
                }
            }
            CoreType::Var(_) | CoreType::Con(_) => {
                // No existential variables to add
            }
        }
    }

    /// Check a pattern against a type and return variable bindings
    /// This implements pattern type checking with unification
    fn check_pattern(
        &mut self,
        pattern: &CorePattern,
        expected_ty: &CoreType,
    ) -> TypeResult<HashMap<String, CoreType>> {
        let mut bindings = HashMap::new();

        match pattern {
            CorePattern::Wildcard => {
                // Wildcard matches anything, no bindings
                Ok(bindings)
            }

            CorePattern::Var(var_name) => {
                // Variable patterns bind the scrutinee type
                bindings.insert(var_name.clone(), expected_ty.clone());
                Ok(bindings)
            }

            CorePattern::Constructor { name, args } => {
                // Look up constructor type from data constructor environment
                if let Some(constructor_ty) = self.data_constructors.get(name).cloned() {
                    // Constructor type should be of the form: T1 -> T2 -> ... -> DataType
                    // We need to unify the result type with expected_ty and
                    // check argument patterns against parameter types

                    let (param_types, result_type) =
                        self.extract_constructor_signature(&constructor_ty)?;

                    // The result type should unify with the expected type
                    self.worklist.push(WorklistEntry::Judgment(Judgment::Sub {
                        left: result_type,
                        right: expected_ty.clone(),
                    }));

                    // Check that we have the right number of pattern arguments
                    if args.len() != param_types.len() {
                        return Err(TypeError::ArityMismatch {
                            expected: param_types.len(),
                            actual: args.len(),
                            span: None,
                        });
                    }

                    // Recursively check argument patterns
                    for (arg_pattern, param_ty) in args.iter().zip(param_types.iter()) {
                        let arg_bindings = self.check_pattern(arg_pattern, param_ty)?;
                        // Merge bindings (should check for conflicts in a real implementation)
                        for (var, ty) in arg_bindings {
                            bindings.insert(var, ty);
                        }
                    }

                    Ok(bindings)
                } else {
                    Err(TypeError::UnboundDataConstructor {
                        name: name.clone(),
                        span: None,
                    })
                }
            }
        }
    }

    /// Extract parameter types and result type from a constructor type
    /// signature e.g., Int -> String -> Bool -> MyData => ([Int, String,
    /// Bool], MyData)
    fn extract_constructor_signature(
        &mut self,
        constructor_ty: &CoreType,
    ) -> TypeResult<(Vec<CoreType>, CoreType)> {
        let mut param_types = Vec::new();
        let mut current_ty = constructor_ty.clone();
        let mut substitutions = HashMap::new();

        // Instantiate forall quantifiers with fresh existential variables
        while let CoreType::Forall(var, body) = current_ty {
            let fresh_var_name = self.worklist.fresh_evar();
            let fresh_evar = CoreType::ETVar(fresh_var_name.clone());

            // Add the existential variable to the worklist
            self.worklist.push(WorklistEntry::TVar(
                fresh_var_name.clone(),
                TyVarKind::Existential,
            ));

            substitutions.insert(var.clone(), fresh_evar);
            current_ty = *body;
        }

        // Apply substitutions to the type
        current_ty = self.apply_type_substitutions(&current_ty, &substitutions);

        // Extract parameter types from arrows
        while let CoreType::Arrow(param_ty, result_ty) = current_ty {
            param_types.push(*param_ty.clone());
            current_ty = *result_ty;
        }

        // The final type is the result type
        Ok((param_types, current_ty))
    }

    fn instantiate_constructor_type(&mut self, constructor_ty: CoreType) -> CoreType {
        let mut current_ty = constructor_ty;
        let mut substitutions = HashMap::new();

        // Instantiate forall quantifiers with fresh existential variables
        while let CoreType::Forall(var, body) = current_ty {
            let fresh_var_name = self.worklist.fresh_evar();
            let fresh_evar = CoreType::ETVar(fresh_var_name.clone());

            // Add the existential variable to the worklist
            self.worklist.push(WorklistEntry::TVar(
                fresh_var_name.clone(),
                TyVarKind::Existential,
            ));

            substitutions.insert(var.clone(), fresh_evar);
            current_ty = *body;
        }

        // Apply substitutions to the type
        self.apply_type_substitutions(&current_ty, &substitutions)
    }

    fn apply_type_substitutions(
        &self,
        ty: &CoreType,
        substitutions: &HashMap<String, CoreType>,
    ) -> CoreType {
        match ty {
            CoreType::Var(name) => substitutions
                .get(name)
                .cloned()
                .unwrap_or_else(|| ty.clone()),
            CoreType::Con(name) => CoreType::Con(name.clone()),
            CoreType::ETVar(name) => CoreType::ETVar(name.clone()),
            CoreType::Arrow(left, right) => CoreType::Arrow(
                Box::new(self.apply_type_substitutions(left, substitutions)),
                Box::new(self.apply_type_substitutions(right, substitutions)),
            ),
            CoreType::App(left, right) => CoreType::App(
                Box::new(self.apply_type_substitutions(left, substitutions)),
                Box::new(self.apply_type_substitutions(right, substitutions)),
            ),
            CoreType::Forall(var, body) => {
                // Don't substitute under forall bindings - this is a simplification
                CoreType::Forall(
                    var.clone(),
                    Box::new(self.apply_type_substitutions(body, substitutions)),
                )
            }
            CoreType::Product(types) => CoreType::Product(
                types
                    .iter()
                    .map(|t| self.apply_type_substitutions(t, substitutions))
                    .collect(),
            ),
        }
    }
}
#[derive(Debug, Clone, PartialEq)]
pub enum WorklistEntry {
    /// Type variable binding: α
    TVar(TyVar, TyVarKind),
    /// Term variable binding: x : T
    Var(TmVar, CoreType),
    /// Judgment: Sub A B | Inf e ⊢ A | Chk e ⇐ A
    Judgment(Judgment),
}
}