Algorithm W

W.lhs.md

Principal type-schemes for functional programs

Luis Damas and Robin Milner, POPL '82

module W where

import Data.List
import Data.Maybe
import Data.Function

Terms

Here is our term language. Binders are labelled by integers (innermost terms first) and introduced by lambda abstractions and let bindings.

For example:

Lam 1 (Var 1)

represents the identity function.

type Binder = Integer

data Expr 
  = Var Binder
  | App Expr   Expr
  | Lam Binder Expr
  | Let Binder Expr Expr
  deriving (Show, Eq)

Syntactic Sugar

Here are some convenience functions for building terms in a HOAS style. We generate binders lazily assuming the outermost abstraction binds the freshest name:

lam :: (Expr -> Expr) -> Expr
lam f = Lam b e
  where
  b = maxVar e + 1
  e = f (Var b)

  maxVar (Lam b _) = b
  maxVar (Let b _ _) = b
  maxVar (Var b) = 0
  maxVar (App e1 e2) = max (maxVar e1) (maxVar e2) 


($$) = App

Now we can introduce terms more naturally as follows:

W> lam $ \x -> lam $ \y -> x $$ y
Lam 2 (Lam 1 (App (Var 2) (Var 1)))

Types

Next, here is the type language, consisting of primitive types (named by strings), variables and functions.

Type schemes are rolled into the same representation for simplicity.

type Prim = String

data Ty
  = P Prim
  | V Binder
  | Ty :~> Ty
  | S [Binder] Ty
  deriving (Show, Eq)

Substitutions

Substitutions are just mappings from binders to types, represented here using association lists:

type Subst = [(Binder, Ty)]

apply :: Subst -> Ty -> Ty
apply s = go
  where
  go (V b) = V b `fromMaybe` lookup b s
  go (t1 :~> t2) = go t1 :~> go t2
  go other = other

To combine two substitutions, we need to make sure all variables are fully applied, by applying the first substitution to types in the second:

combine :: Subst -> Subst -> Subst
combine s1 s2 = [ (b, apply s1 ty) | (b, ty) <- s2 ] ++ s1

Unification

Unification is defined by zipping two types until we hit unification variables. At that point, we grow the substitution to assert the appropriate equalities. If we encounter two inconsistent types, we raise an error.

The occurs check is necessary to ensure we don't generate infinite types:

occursIn :: Binder -> Ty -> Bool
b `occursIn` V b1    = b == b1
b `occursIn` (t1 :~> t2) = (b `occursIn` t1) || (b `occursIn` t2)
b `occursIn` S bs t  = (b `occursIn` t) && (b `notElem` bs)
_ `occursIn` _           = False

unify :: Ty -> Ty -> Subst
unify (P p1) (P p2) | p1 == p2 = []
unify (V b) t | b `occursIn` t = error "Occurs check failed!"
              | otherwise = [(b, t)]
unify t (V b) = unify (V b) t 
unify (t1 :~> t2) (t3 :~> t4) = unify t1 t3 `combine` unify t2 t4
unify _ _ = error "Unification failed"

Instantiation

To instantiate a type scheme, we replace all of the abstracted type variables by fresh unification variables:

instantiate :: Fresh -> Ty -> (Ty, Fresh)
instantiate fr (S bs t) = (s `apply` t, fr1)
  where
  (s, fr1) = makeFresh fr bs

  makeFresh :: Fresh -> [Binder] -> (Subst, Fresh)
  makeFresh fr [] = ([], fr)
  makeFresh fr (b:bs) = ((b, V fr):s, fr1)
    where
    (s, fr1) = makeFresh (fr + 1) bs
instantiate fr ty = (ty, fr)

Generalization

To generalize a type to a type scheme, we collect the free variables occurring in a type and abstract over them using S:

freeVars :: Ty -> [Binder]
freeVars (V b) = [b]
freeVars (t1 :~> t2) = freeVars t1 ++ freeVars t2
freeVars (S bs t) = freeVars t \\ bs
freeVars _ = []

generalize :: Subst -> Ty -> Ty
generalize ctx ty = S free ty
  where
  free = nub $ freeVars ty \\ map fst ctx

Algorithm W

Now we can assemble algorithm W. We proceed bottom up over the syntax of terms, growing the input substitution by asserting the type equalities which are forced on us by the typing rules for the current term.

We keep a fresh name supply, so that we can instantiate type schemes and generate fresh unification variables for the types of function arguments and let bound names.

This implementation is a direct translation of the description in the paper.

type Fresh = Binder

w :: Expr -> Ty
w = uncurry generalize . fst . go 0 []
  where
  go :: Fresh -> Subst -> Expr -> ((Subst, Ty), Fresh)
  go fr a (Var b) = (([], t), fr1)
    where
    (t, fr1) = maybe (error "Unbound name") (instantiate fr) mt 
    mt = lookup b a
  go fr a (App e1 e2) = ((v `combine` s2 `combine` s1, apply v b), fr2)
    where
    ((s1, t1), fr1) = go (fr + 1) a e1
    ((s2, t2), fr2) = go fr1 (s1 `combine` a) e2
    b = V fr
    v = unify (apply s2 t1) (t2 :~> b)
  go fr a (Lam x e1) = ((s1, apply s1 b :~> t1), fr1)
    where
    b = V fr
    ((s1, t1), fr1) = go (fr + 1) (remove x a ++ [(x, b)]) e1
  go fr a (Let x e1 e2) = ((s2 `combine` s1 `combine` a, t2), fr2)
    where
    ((s1, t1), fr1) = go fr a e1
    ((s2, t2), fr2) = go fr1 ((s1 `combine` remove x a) ++ [(x, generalize (s1 `combine` a) t1)]) e2
  
  remove :: (Eq a) => a -> [(a, b)] -> [(a, b)]
  remove a = filter ((/= a) . fst)