Algorithm W

W.lhs

## 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)