A type checker for CoC (really any PTS)

CoC.hs

{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable, LambdaCase #-}
module CoC where
import           Bound
import           Control.Applicative
import           Control.Monad
import           Control.Monad.Gen
import           Control.Monad.Trans
import           Data.Foldable       hiding (elem)
import qualified Data.Map            as M
import           Data.Maybe
import           Data.Traversable
import           Prelude.Extras

data Exp a = V a
           | Lam (Exp a) (Scope () Exp a)
           | Pi  (Exp a) (Scope () Exp a)
           | App (Exp a) (Exp a)
           | Star
           | Box
           deriving(Eq, Show, Functor, Foldable, Traversable)

instance Eq1 Exp where
instance Show1 Exp where
instance Applicative Exp where
  pure = return
  (<*>) = ap
instance Monad Exp where
  return = V
  V a >>= f = f a
  Lam t b >>= f = Lam (t >>= f) (b >>>= f)
  Pi t b  >>= f = Pi (t >>= f) (b >>>= f)
  App a b >>= f = App (a >>= f) (b >>= f)
  Star >>= _ = Star
  Box >>= _ = Box

whnf :: Exp a -> Maybe (Exp a)
whnf (App f a) = whnf f >>= \case
  Lam t b -> whnf (instantiate1 a b)
  _       -> Nothing
whnf a = Just a

equal :: Eq a => Exp a -> Exp a -> GenT a Maybe Bool
equal a b = do
  a' <- lift . whnf $ a
  b' <- lift . whnf $ b
  case (a', b') of
   (Star, Star) -> return True
   (Box, Box)   -> return True
   (Pi t b, Pi t' b') -> do
     v <- V <$> gen
     (&&) <$> equal t t' <*> equal (instantiate1 v b) (instantiate1 v b')
   (Lam t b, Lam t' b') -> do
     v <- V <$> gen
     (&&) <$> equal t t' <*> equal (instantiate1 v b) (instantiate1 v b')
   (V v, V v') -> return (v == v')
   _ -> return False

check :: Ord a => M.Map a (Exp a) -> Exp a -> GenT a Maybe (Exp a)
check _ Box = mzero -- Never can type Box
check _ Star = return Box
check env (V v) = lift $ M.lookup v env
check env (App f a) = do
  fTy <- check env f >>= lift . whnf
  aTy <- check env a
  case fTy of
   Pi t retTy -> do
     equal t aTy >>= guard
     return (instantiate1 a retTy)
   _ -> mzero
check env (Lam t b) = do
  _ <- check env t
  t' <- lift (whnf t)
  v <- gen
  retTy <- check (M.insert v t' env) (instantiate1 (V v) b)
  let pi = Pi t' (abstract1 v retTy)
  check env pi >> return pi
check env (Pi t b) = do
  s1 <- check env t
  t' <- lift (whnf t)
  v <- gen
  s2 <- check (M.insert v t' env) (instantiate1 (V v) b)
  guard (s1 `elem` sorts && s2 `elem` sorts) >> return s2
  where sorts = [Star, Box]

typeOf :: (Ord a, Enum a) => Exp a -> Maybe (Exp a)
typeOf = runGenT . check M.empty