Skip to content

Instantly share code, notes, and snippets.

@paf31
Last active September 5, 2018 03:54
Show Gist options
  • Save paf31/d1cd25c7527efbe896da to your computer and use it in GitHub Desktop.
Save paf31/d1cd25c7527efbe896da to your computer and use it in GitHub Desktop.
A simple type checker with Rank N types
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment