A simple type checker with Rank N types

Main.hs

module Main where
  
import Data.Maybe (fromMaybe)

import Control.Applicative
import Control.Arrow (first)
import Control.Monad (ap) 
    
import Debug.Trace
  
data Nat = Z | S Nat deriving (Eq, Ord)
    
instance Show Nat where
  show = show . fromNat
    where
    fromNat Z = 0
    fromNat (S n) = 1 + fromNat n
    
data Tm 
  = Var Nat
  | App Tm Tm
  | Lam Nat Tm
  | Let Nat Tm Tm
  | Typed Tm Ty
  deriving (Show)
  
type Skolem = Nat
  
data Ty
  = TyVar Nat
  | TySko Skolem (Nat, Nat)
  | TyArr Ty Ty
  | TyForall Nat Ty

instance Show Ty where
  show (TyVar v) = "t" ++ show v
  show (TySko v _) = "s" ++ show v
  show (TyArr t1 t2) = "(" ++ show t1 ++ " ~> " ++ show t2 ++ ")"
  show (TyForall x t) = "forall t" ++ show x ++ ". " ++ show t
  
lam :: (Tm -> Tm) -> Tm
lam f = Lam b e
  where
  b = S $ maxVar e
  e = f (Var b)

  maxVar (Lam b _) = b
  maxVar (Let b _ _) = b
  maxVar (Var _) = Z
  maxVar (App e1 e2) = max (maxVar e1) (maxVar e2)  
  maxVar (Typed e _) = maxVar e
  
forall_ :: (Ty -> Ty) -> Ty
forall_ f = TyForall b t
  where
  b = S $ maxVar t
  t = f (TyVar b)

  maxVar (TyForall b _) = b
  maxVar (TyArr t1 t2) = max (maxVar t1) (maxVar t2)  
  maxVar _ = Z
  
infixl 4 $$
  
($$) = App
   
data Subst = Subst 
  { substVars :: [(Nat, Ty)] 
  , substSkolemChecks :: [(Nat, Nat, Nat)]
  } deriving Show

mkSubst :: Nat -> Ty -> Subst
mkSubst v t = Subst [(v, t)] []

emptySubst :: Subst
emptySubst = Subst [] []

infixr 4 ~>

(~>) = TyArr

subst :: Subst -> Ty -> Ty
subst s = go
  where
  go (TyVar b) = TyVar b `fromMaybe` lookup b (substVars s)
  go (TyArr t1 t2) = TyArr (go t1) (go t2)
  go other = other
  
combine :: Subst -> Subst -> Subst
combine s1@(Subst sv1 sk1) s2@(Subst sv2 sk2) = Subst sv sk 
  where
  sv = [ (b, subst s1 ty) | (b, ty) <- sv2 ] ++ 
       [ (b, subst s2 ty) | (b, ty) <- sv1 ]
  sk = sk1 ++ sk2

occursIn :: Nat -> Ty -> Bool
b `occursIn` TyVar b1       = b == b1
b `occursIn` TyArr t1 t2    = (b `occursIn` t1) || (b `occursIn` t2)
b `occursIn` TyForall bs t  = (b `occursIn` t) && b /= bs
_ `occursIn` _              = False

escapeCheck :: Subst -> Bool
escapeCheck = all (\(lo, s, hi) -> lo <= s && s <= hi) . substSkolemChecks

newtype Fresh a = Fresh { unFresh :: (Nat, Skolem) -> (a, (Nat, Skolem)) }

runFresh :: Fresh a -> a
runFresh = fst . flip unFresh (Z, Z)

bounded :: ((Nat, Nat) -> Fresh a) -> Fresh a
bounded f = Fresh $ \b@(lo, _) ->
  let (a, b'@(hi, _)) = unFresh (f (lo, hi)) b
  in (a, b')

fresh :: Fresh Nat
fresh = Fresh $ \(n, m) -> (n, (S n, m))

skolem :: Fresh Skolem
skolem = Fresh $ \(n, m) -> (m, (n, S m))

instance Functor Fresh where
  fmap f fr = Fresh (first f . unFresh fr)

instance Applicative Fresh where
  pure = return
  (<*>) = ap

instance Monad Fresh where
  return a = Fresh $ \i -> (a, i)
  fr >>= f = Fresh $ \i -> let (a, j) = unFresh fr i
                           in unFresh (f a) j
   
type Context = [(Nat, Ty)]

unify :: Ty -> Ty -> Fresh Subst 
unify t1 t2 = trace ("unify: " ++ show (t1, t2)) $ unify' t1 t2
  where
  unify' t@(TySko s (lo, hi)) (TyVar b) = return (Subst [(b, t)] [(lo, b, hi)])
  unify' (TySko s1 _) (TySko s2 _) = return emptySubst
  unify' (TyVar b1) (TyVar b2) | b1 == b2 = return emptySubst
  unify' (TyVar b) t 
    | b `occursIn` t = error "Occurs check failed"
    | otherwise = return (mkSubst b t)
  unify' t (TyVar b) = unify (TyVar b) t
  unify' (TyArr t1 t2) (TyArr t3 t4) = combine <$> unify t1 t3 <*> unify t2 t4
  unify' s t = error "Unification failed"

infer :: Context -> Subst -> Tm -> Fresh (Subst, Ty)
infer ctx s e = trace ("infer: " ++ show e) $ infer' ctx s e
  where
  infer' ctx s0 (Var b) = do
    case lookup b ctx of
      Just ty -> return (s0, ty)
      Nothing -> error "Unbound name"
  infer' ctx s0 (App e1 e2) = do
    (s1, t1) <- infer ctx s0 e1
    apply ctx s1 t1 e2
  infer' ctx s0 (Lam x e1) = do
    b <- fmap TyVar fresh
    (s1, t1) <- infer ((x, b) : ctx) s0 e1
    return (s1, subst s1 (TyArr b t1))
  infer' ctx s0 (Let x e1 e2) = do
    (s1, t1) <- infer ctx s0 e1
    (s2, t2) <- infer ((x, t1) : ctx) s1 e2
    return (s2, t2)
  infer' ctx s0 (Typed e ty) = do
    s1 <- check ctx s0 e ty
    return (s1, ty)

infer' :: Tm -> (Subst, Ty)
infer' = runFresh . infer [] emptySubst

check :: Context -> Subst -> Tm -> Ty -> Fresh Subst
check ctx s e t = trace ("check: " ++ show (e, t)) $ check' ctx s e t
  where
  check' ctx s0 (Lam x e1) (TyArr t1 t2) = do
    s1 <- check ((x, t1) : ctx) s0 e1 t2
    return s1
  check' ctx s0 e (TyForall x ty) = bounded $ \bounds -> do
    s <- skolem
    check ctx s0 e (subst (mkSubst x (TySko s bounds)) ty)
  check' ctx s0 e ty = do
    (s1, ty1) <- infer ctx s0 e
    subsumes ctx s1 ty1 ty

check' :: Tm -> Ty -> Subst
check' e = runFresh . check [] emptySubst e

apply :: Context -> Subst -> Ty -> Tm -> Fresh (Subst, Ty)
apply ctx s t e = trace ("apply: " ++ show (t, e)) $ apply' ctx s t e
  where
  apply' ctx s0 (TyForall x t) e = do
    f <- fresh 
    apply ctx s0 (subst (mkSubst x (TyVar f)) t) e
  apply' ctx s0 (TyArr t1 t2) e = do
    s1 <- check ctx s0 e t1
    return (s1, subst s1 t2)
  apply' ctx s0 ty e = do
    f <- fresh
    (s1, t1) <- infer ctx s0 e
    s2 <- ty `unify` TyArr t1 (TyVar f)
    let s3 = combine s2 s1
    return (s3, subst s3 (TyVar f))

subsumes :: Context -> Subst -> Ty -> Ty -> Fresh Subst
subsumes ctx s t1 t2 = trace ("subsumes: " ++ show (t1, t2)) $ subsumes' ctx s t1 t2
  where
  subsumes' ctx s0 (TyForall x t1) t2 = do
    f <- fresh 
    subsumes ctx s0 (subst (mkSubst x (TyVar f)) t1) t2
  subsumes' ctx s0 t1 (TyForall x t2) = bounded $ \bounds -> do
    s <- skolem
    subsumes ctx s0 (subst (mkSubst x (TySko s bounds)) t1) t2
  subsumes' ctx s0 (TyArr t1 t2) (TyArr t3 t4) = do
    s1 <- subsumes ctx s0 t2 t4
    subsumes ctx s1 (subst s1 t3) (subst s1 t1)
  subsumes' ctx s0 t1 t2 = do
    s1 <- unify t1 t2
    return (combine s1 s0)
  
church :: Ty
church = forall_ $ \t -> (t ~> t) ~> t ~> t

zero :: Tm
zero = lam $ \_ -> lam $ \x -> x

suck :: Tm
suck = lam $ \f -> lam $ \g -> lam $ \x -> f $$ g $$ (g $$ x)

test = escapeCheck $ check' (suck $$ zero) $ church