lean-pl-tutorials/tutorial-02-semantics/14-algorithm-w.org

* Unit 14 --- Algorithm W (Hindley-Milner Type Inference)
:PROPERTIES:
:CUSTOM_ID: unit-14-algorithm-w-hindley-milner-type-inference
:END:
*Tutorial 2: PL Semantics in Lean* · [[../README.org][← Back to README]]

** Goals
:PROPERTIES:
:CUSTOM_ID: goals
:END:
- Understand type inference as solving unification constraints
- Implement *Algorithm W*: the classic HM inference algorithm
- Write =unify= (Robinson's unification)
- Write =infer= (the main inference loop)

** Sources
:PROPERTIES:
:CUSTOM_ID: sources
:END:
- [[https://www.jeffvaughan.net/docs/hmproof.pdf][Vaughan 2008, §3]]
- [[https://github.com/sdemos/type-inference][sdemos/type-inference]]
- [[https://en.wikipedia.org/wiki/Hindley%E2%80%93Milner_type_system#Algorithm_W][Wikipedia: Algorithm W]]

** Background
:PROPERTIES:
:CUSTOM_ID: background
:END:
Algorithm W takes an expression and returns a substitution + type:

#+begin_example
infer(Γ, e) = (S, τ)
#+end_example

- =Γ= maps variables to type schemes
- =e= is the expression to type
- =S= is a type substitution (unifies constraints found during inference)
- =τ= is the inferred monotype

The algorithm uses a supply of /fresh type variables/ to build constraints,
then unifies them.

** Exercises
:PROPERTIES:
:CUSTOM_ID: exercises
:END:
#+begin_src lean
open MonoType
open HMExpr
open TypeScheme

-- 14.1 — Substitutions
-- A substitution maps type variables to monotypes
abbrev Subst := Nat → MonoType

-- Identity substitution
def idSubst : Subst := fun α => MonoType.tvar α

-- Apply a substitution to a monotype
def applySubst (S : Subst) : MonoType → MonoType
  | MonoType.tvar α => S α
  | MonoType.fn τ₁ τ₂ => MonoType.fn (applySubst S τ₁) (applySubst S τ₂)

-- Compose substitutions: (S₁ ∘ S₂)(α) = S₁(S₂(α))
def compose (S₁ S₂ : Subst) : Subst :=
  fun α => applySubst S₁ (S₂ α)

-- 14.2 — Unification (Robinson's algorithm)
-- Returns a substitution that makes τ₁ and τ₂ equal, if possible.
-- Fails if there's a type mismatch (e.g., unifying α → β with α is impossible).
def unify (τ₁ τ₂ : MonoType) : Option Subst :=
  match τ₁, τ₂ with
  | MonoType.tvar α, MonoType.tvar β =>
      if α == β then some idSubst
      else some (fun γ => if γ == α then MonoType.tvar β else MonoType.tvar γ)
  | MonoType.tvar α, τ =>
      if occurs α τ then none  -- occurs check: α ∉ ftv(τ)
      else some (fun γ => if γ == α then τ else MonoType.tvar γ)
  | τ, MonoType.tvar α =>
      if occurs α τ then none
      else some (fun γ => if γ == α then τ else MonoType.tvar γ)
  | MonoType.fn τ₁ᵃ τ₁ᵇ, MonoType.fn τ₂ᵃ τ₂ᵇ =>
      match unify τ₁ᵃ τ₂ᵃ with
      | none => none
      | some S₁ =>
          match unify (applySubst S₁ τ₁ᵇ) (applySubst S₁ τ₂ᵇ) with
          | none => none
          | some S₂ => some (compose S₂ S₁)

-- Occurs check: does α appear in τ?
def occurs (α : Nat) : MonoType → Bool
  | MonoType.tvar β => α == β
  | MonoType.fn τ₁ τ₂ => occurs α τ₁ || occurs α τ₂

-- 14.3 — Fresh variable supply
-- We use a counter to generate fresh type variables
def freshVar (counter : Nat) : Nat × Nat := (counter, counter + 1)

-- 14.4 — Generalization: close a type under the environment
-- `generalize(Γ, τ)` produces `∀αs. τ` where αs = ftv(τ) \ ftv(Γ)
def generalize (Γ : HMEnv) (τ : MonoType) : TypeScheme :=
  { vars := (ftv τ).filter (fun α => α ∉ ftv_env Γ)
  , body := τ }

-- 14.5 — Algorithm W (the core inference algorithm)
-- Returns `(S, τ)` where S is a substitution and τ the inferred type.
-- Uses a state monad for the fresh variable counter (simplified here).
def inferW (Γ : HMEnv) (e : HMExpr) (counter : Nat) : Option (Subst × MonoType × Nat) :=
  match e with
  | HMExpr.var i =>
      match lookup Γ i with
      | none => none
      | some σ => some (idSubst, instantiate σ counter, counter + length σ.vars)

  | HMExpr.lam body =>
      -- Create a fresh type variable for the parameter
      let (α, counter') := freshVar counter
      let τ_param := MonoType.tvar α
      -- Add x : α to the environment
      let Γ' := {vars := [], body := τ_param} :: Γ
      -- Infer the body type
      match inferW Γ' body counter' with
      | none => none
      | some (S, τ_body, counter'') =>
          some (S, MonoType.fn (applySubst S τ_param) τ_body, counter'')

  | HMExpr.app f a =>
      match inferW Γ f counter with
      | none => none
      | some (S₁, τ_f, counter₁) =>
          match inferW (applySubstEnv S₁ Γ) a counter₁ with
          | none => none
          | some (S₂, τ_a, counter₂) =>
              let β := freshVar counter₂
              let α_counter₃ := β.2
              match unify (applySubst S₂ τ_f) (MonoType.fn τ_a (MonoType.tvar β.1)) with
              | none => none
              | some S₃ =>
                  let S := compose S₃ (compose S₂ S₁)
                  some (S, applySubst S₃ (MonoType.tvar β.1), α_counter₃)

  | HMExpr.lett e₁ e₂ =>
      match inferW Γ e₁ counter with
      | none => none
      | some (S₁, τ₁, counter₁) =>
          let σ₁ := generalize (applySubstEnv S₁ Γ) τ₁
          let Γ' := σ₁ :: applySubstEnv S₁ Γ
          match inferW Γ' e₂ counter₁ with
          | none => none
          | some (S₂, τ₂, counter₂) =>
              some (compose S₂ S₁, τ₂, counter₂)

-- 14.6 — Exercise: infer the type of λx. x
def infer_id : Option (Subst × MonoType × Nat) :=
  inferW [] (HMExpr.lam (HMExpr.var 0)) 0

-- Should return a substitution mapping α₀ to α₀ and type α₀ → α₀
-- #eval infer_id

-- 14.7 — Exercise: infer the type of let x = λy. y in x x
def infer_self_app : Option (Subst × MonoType × Nat) :=
  inferW [] self_app_id 0
#+end_src

--------------

← [[file:13-hm-declarative.org][Previous: Unit 13]] · Next: [[file:15-soundness-completeness.org][Unit 15 --- Soundness and Completeness]]