TypeChecker.hs

{-# LANGUAGE PackageImports, TypeSynonymInstances, FlexibleInstances, GADTs, TupleSections, GeneralizedNewtypeDeriving #-}

module TypeChecker where

import Prelude hiding (length)
import Basics
import Display
import Control.Monad.Error
import Data.Maybe
import Control.Monad.Error (ErrorT, runErrorT)
import Control.Monad.Writer
import Control.Monad.Reader
import Data.Functor.Identity
import Data.Sequence hiding (replicate,zip,reverse)
import Data.Foldable (toList)
import Terms
import Options

instance Error (Term,Doc) where
  strMsg s = (hole "strMsg: panic!",text s)

newtype Result a = Result (
  (ReaderT (Term,Type)) -- the module being type-checked and its type (in normal form)
  ((ErrorT (Term,Doc)) -- term is used for position information
  (Writer [Doc])) a)
 deriving (Functor, Applicative, Monad, MonadReader (Term,Type), MonadError (Term,Doc), MonadWriter [Doc])

report :: Doc -> Result ()
report x = tell [x]

runChecker :: (Term,Type) -> Result a -> (Either (Term,Doc) a,[Doc])
runChecker x (Result m) = runIdentity $ runWriterT $ runErrorT $ runReaderT m x

type Value    = NF
type Type     = Value
data Bind     = Bind {entryIdent :: Ident, 
                      entryType :: Type   -- ^ Note: the context of the type does not contain the variable bound here.
                     }
type Context  = Seq Bind

display :: Context -> NF -> Doc
display c = prettyTerm (fmap entryIdent c)

dispContext :: Context -> Doc
dispContext ctx = case viewl ctx of
  EmptyL -> mempty
  Bind x typ :< ctx' -> (pretty x <+> ":" <+> display ctx' typ) $$ dispContext ctx'

-- FIXME: flag an error if impredicativity disabled and we use it anyway.

-- | Infer a type and evaluate to normal form
iType :: Context -> Term -> Result (Value,Type)
iType g (Ann e et) = do
  (et',o) <- iSort g et 
  e' <- cType g e et'
  return (e',et') -- type annotations are removed once they're checked
iType g t@(Star p s) =  return (Star p s,Star p $ above s)  
iType g (Pi ident tyt tyt')  = do
    (ty ,s1) <- iSort g tyt 
    (ty',s2) <- iSort (Bind ident ty <| g) tyt'
    let o = s1 ⊔ s2
    return (Pi ident ty ty', Star (identPosition ident) o)
iType g (Sigma i []) = return (Sigma i [],Star i 0)          
iType g (Sigma i ((f,t):ts)) = do
   (t',s1)  <- iSort g t 
   (Sigma _ ts',s2) <- iSort (Bind (synthId f)  t' <| g) (Sigma i ts)
   let o = s1 ⊔ s2          
   return (Sigma i ((f,t'):ts'), Star i o)
iType g e@(V _ x) = do
  return $ (var x, wkn (x+1) ∙ typ)
  where Bind _ typ = g `index` x
        
iType g (Hole p x) = do
  report $ hang (text ("context of " ++ x ++ " is")) 2 (dispContext g)
  return (Hole p x, Hole p ("type of " ++ x))

iType g t@(Lam x (Hole _ _) e) = throwError (t,"cannot infer type for" <+> display g t)
iType g (Lam x ty e) = do
    (vty,vs) <- iSort g ty
    (ve,t) <- iType (Bind x  vty <| g) e
    return $ (Lam x vty ve, Pi x vty t)

iType g (Pair p fs) = do
  (fs,vs,ts) <- unzip3 <$> (forM fs $ \(f,x) -> do 
                            (x',xt) <- iType g x
                            return (f,x',xt))
  return (Pair p (zip fs vs),Sigma p (zip fs [ wkn n ∙ t | (n,t) <- zip [0..] ts])) 
  -- possible dependencies in the type are not discovered 
  
iType g (Fin p ts) = return (Fin p ts,Star p 0)
  
iType g (App e1 e2) = do
   (e1',et) <- iType g e1
   et' <- whnf g et
   case et' of
     Pi _ ty ty' -> do 
          v2 <- cType g e2 ty
          return (app e1' v2, subst0 v2 ∙ ty') 
     _             ->  throwError (e1,"type should be function:" <> display g et)

iType g (Proj e f) = do
  (e',et) <- iType g e
  et' <- whnf g et
  case et' of 
    Sigma _ fs -> case projType e f fs of
                    Just tt -> return (proj e f,tt)
                    _ -> throwError (e,"field" <+> text f <+> "not found in type" <+> display g et')
    _ -> throwError (e,hang "expected record:" 2 $ sep ["term:" <+> display g e,
                                                        "type:" <+> display g et,
                                                        "field:" <+> text f])
    
iType g This = do
  (_,thist) <- ask
  return (This, thist)

iType g t = throwError (t,hang "cannot infer type for" 2 $ display g t)
         
-- | Infer a sort, and normalise
iSort :: Context -> Term -> Result (Type,Sort)
iSort g e = do
  (val,v) <- iType g e
  v' <- whnf g v
  case v' of 
    Star _ i -> return (val,i)
    Hole _ h -> do 
         report $ text h <+> "must be a type"
         return $ (hole h, Sort 1)
    _ -> throwError (e,display g e <+> "is not a type")


-- | Check the type and normalize
cType0 :: Context -> Term -> Type -> Result Value
cType0 g (Lam name (Hole _ _) e) (Pi name' ty ty') = do
        e' <- cType (Bind name ty <| g) e ty'
        return (Lam name ty e') -- the type and binder is filled in.

cType0 g (Lam name ty0 e) (Pi name' ty ty')
  =     do (t,_o) <- iSort g ty0
           unify g (Hole (identPosition name) (show name)) ty t
           e' <- cType (Bind name ty <| g) e ty'
           return (Lam name ty e')

-- note: there is automatic coercion of values to more specific types.
cType0 g (Pair p _ignored) (Sigma _ []) = return $ Pair p []
cType0 g (Pair p ((f,x):xs)) (Sigma i fs@((f',xt):xts)) 
  | f /= f' = do  Pair _ fs' <- cType g (Pair p xs) (Sigma p fs)
                  return $ Pair p ((f,x):fs')
  | f == f' = do
    x' <- cType g x xt
    -- report $ hang "Remains to check: " 2 $ sep ["pair:" <+> display g (pair xs), "type:" <+>  display g (subst0 x' ∙ sigma xts)]
    Pair _ xs' <- cType g (Pair p xs) (subst0 x' ∙ Sigma p xts)
    return (Pair p ((f,x'):xs'))

cType0 g c@(Cas p cs) (Pi _ (Fin p' ct') vt) = do
  let ct = Fin p' $ map fst cs
  unify (Bind dummyId ct <| g) (var 0) (Fin p' ct') ct
  cs' <- forM cs $ \ (x,v) -> do
     unless (x `elem` ct') $ throwError (v,"tag not found: " <> text x)
     v' <- cType g v (subst0 (Tag p x) ∙ vt)
     return (x,v')
  return (Cas p cs')
  
cType0 g e@(Tag p t) ty@(Fin p' ts) 
  | t `elem` ts = return e
  | otherwise = throwError (e,vcat ["tag: " <> text t, "not in: " <> display g ty])

cType0 g e v = do 
  -- report $ hang "no checking rule for" 2 $ vcat ["term = " <> display g e, "type = " <>display g v]
  (e',v') <- iType g e
  unify g e v' v
  return e'
  
-- | Unification
unify :: Context -> Term -> Type -> Type -> Result ()
unify g0 e q0' q0 = unif g0 q0' q0 where 
   unif, unif' :: Context -> Type -> Type -> Result ()
   unif' g q' q = do x' <- whnf g q'
                     x <- whnf g q
                     unif g x' x
   unif g q' q = case (q',q) of
     (Star _ s , Star _ s') -> check $ s == s'
     (Pi i a b , Pi _ a' b') -> (a' <: a) >> unif (Bind i a <| g) b b'
     (Lam i a b , Lam _ a' b') -> (a' <: a) >> unif (Bind i a <| g) b b'
     (App a b , App a' b') -> unif g a a' >> (b <: b')
     (Sigma _ [], Sigma _ []) -> return ()    
     (Sigma p ((f,x):xs) , Sigma p' ((f',x'):xs')) -> do
           check $ f == f'
           x <: x'
           unif (Bind (synthId f) x' <| g) (Sigma p xs) (Sigma p' xs')
     (Pair _ xs , Pair _ xs') -> eqList (\(f,x) (f',x') -> check (f == f') >> x <: x') xs xs'
     (Proj x f , Proj x' f') -> check (f == f') >> unif g x x'
     (V _ x , V _ x') -> check (x == x')
     (Hole _ _, _) -> constraint
     (_, Hole _ _) -> constraint
     (This,This) -> return ()
     (Tag _ t, Tag _ t') -> check $ t == t'
     (Fin _ ts , Fin _ ts') -> check $ ts == ts'
     (Cas _ xs , Cas _ xs') -> eqList (\(f,x) (f',x') -> check (f == f') >> x <: x') xs xs'
     _ | isBox q || isBox q' -> unif' g q' q
     _  -> crash
     where (<:) :: Type -> Type -> Result ()
           (<:) = unif g 
           infix 4 <:
           constraint = report $ vcat (["constraint: " <> display g q',
                                        "equal to: " <> display g q] ++ chkCtx)
           check :: Bool -> Result () 
           check c = unless c crash 
           crash :: Result ()
           crash = throwError (e,hang "type mismatch: " 2 $ vcat 
                               (["type:         " <+> display g q',
                                 "is not sub of:" <+> display g q] ++ chkCtx))
           chkCtx = ["inferred:" <+> display g0 q0',
                     "expected:" <+> display g0 q0 ,
                     "for:" <+> display g0 e ,
                     "context:" <+> dispContext g0]
           eqList p [] [] = return ()
           eqList p (x:xs) (x':xs') = p x x' >> eqList p xs xs'
           eqList p _ _ = crash

cType g e t = cType0 g e =<< whnf g t

-- | Reduce to whnf. May loop if 
whnf g x = do
  -- report $ "whnf: " <> (display g x)
  m <- fst <$> ask
  case whnf' m x of 
    Nothing -> return x
    Just x' -> whnf g x' -- TODO: consume some fuel

-- | One step of reduction to whnf
whnf' :: Term -> Term -> Maybe Term
whnf' m This = Just m
whnf' m (Proj x f) = proj <$> whnf' m x <*> pure f
whnf' m (App  f x) = app  <$> whnf' m f <*> pure x
whnf' m (Ann x t) = Just x
whnf' _ x = Nothing

-- | Will the argument reduce at all if passed to whnf?
isBox x = isJust (whnf' (hole "isBox: dummy") x)

projType e f fs = case break ((==f) . fst) fs of
                    (hs,((_,t):_)) -> Just (([proj e h | (h,_) <- reverse hs] ++ map var [0..]) ∙ t)
                    _ -> Nothing