TypeChecker.hs

{-# LANGUAGE TupleSections #-}
module TypeChecker(
  runInfer
  , verboseEnv
  , runDeclss
  , env
  , TEnv) where

import Control.Applicative hiding (empty)
import Control.Monad
import Control.Monad.Except
import Control.Monad.Reader
import Data.Map (Map,(!),mapWithKey,assocs,filterWithKey,elems,keys
                ,intersection,intersectionWith,intersectionWithKey
                ,toList,fromList)
import qualified Data.Map as Map
import qualified Data.Traversable as T

import CTT
import qualified Connections as C
import qualified Eval as E

type Typing a = ReaderT TEnv (ExceptT String IO) a

data TEnv =
  TEnv { names   :: [String] -- generated names
       , indent  :: Int
       , env     :: Env
       , verbose :: Bool  -- Should it be verbose and print what it typechecks?
       } deriving (Eq)

verboseEnv, silentEnv :: TEnv
verboseEnv = TEnv [] 0 emptyEnv True
silentEnv  = TEnv [] 0 emptyEnv False

trace :: String -> Typing ()
trace s = do
  b <- asks verbose
  when b $ liftIO (putStrLn s)

runTyping :: TEnv -> Typing a -> IO (Either String a)
runTyping env t = runExceptT $ runReaderT t env

runDecls :: TEnv -> Decls -> IO (Either String TEnv)
runDecls tenv d = runTyping tenv $ do
  checkDecls d
  return $ addDecls d tenv

runDeclss :: TEnv -> [Decls] -> IO (Maybe String,TEnv)
runDeclss tenv []     = return (Nothing, tenv)
runDeclss tenv (d:ds) = do
  x <- runDecls tenv d
  case x of
    Right tenv' -> runDeclss tenv' ds
    Left s      -> return (Just s, tenv)

runInfer :: TEnv -> Ter -> IO (Either String Val)
runInfer lenv e = runTyping lenv (infer e)

addTypeVal :: (Ident,Val) -> TEnv -> TEnv
addTypeVal (x,a) (TEnv ns ind rho v) =
  let w@(VVar n _) = mkVarNice ns x a
  in TEnv (n:ns) ind (upd (x,w) rho) v

addSub :: (C.Name,C.Formula) -> TEnv -> TEnv
addSub iphi (TEnv ns ind rho v) = TEnv ns ind (sub iphi rho) v

addSubs :: [(C.Name,C.Formula)] -> TEnv -> TEnv
addSubs = flip $ foldr addSub

addType :: (Ident,Ter) -> TEnv -> TEnv
addType (x,a) tenv@(TEnv _ _ rho _) = addTypeVal (x,E.eval rho a) tenv

addBranch :: [(Ident,Val)] -> Env -> TEnv -> TEnv
addBranch nvs env (TEnv ns ind rho v) =
  TEnv ([n | (_,VVar n _) <- nvs] ++ ns) ind (upds nvs rho) v

addDecls :: Decls -> TEnv -> TEnv
addDecls d (TEnv ns ind rho v) = TEnv ns ind (def d rho) v

addTele :: Tele -> TEnv -> TEnv
addTele xas lenv = foldl (flip addType) lenv xas

faceEnv :: C.Face -> TEnv -> TEnv
faceEnv alpha tenv = tenv{env=env tenv `C.face` alpha}

getLblType :: LIdent -> Val -> Typing (Tele, Env)
getLblType c (Ter (Sum _ _ cas) r) = case lookupLabel c cas of
  Just as -> return (as,r)
  Nothing -> throwError ("getLblType: " ++ show c ++ " in " ++ show cas)
getLblType c (Ter (HSum _ _ cas) r) = case lookupLabel c cas of
  Just as -> return (as,r)
  Nothing -> throwError ("getLblType: " ++ show c ++ " in " ++ show cas)
getLblType c u = throwError ("expected a data type for the constructor "
                             ++ c ++ " but got " ++ show u)

unlessM :: Monad m => m Bool -> m () -> m ()
unlessM mb x = mb >>= flip unless x

mkVars :: [String] -> Tele -> Env -> [(Ident,Val)]
mkVars _ [] _           = []
mkVars ns ((x,a):xas) nu =
  let w@(VVar n _) = mkVarNice ns x (E.eval nu a)
  in (x,w) : mkVars (n:ns) xas (upd (x,w) nu)

(===) :: E.Convertible a => a -> a -> Typing Bool
u === v = E.conv <$> asks names <*> pure u <*> pure v

evalTyping :: Ter -> Typing Val
evalTyping t = E.eval <$> asks env <*> pure t

check :: Val -> Ter -> Typing ()
check a t = case (a,t) of
  (_,Undef{}) -> return ()
  (_,Hole l)  -> do
      rho <- asks env
      let e = unlines (reverse (contextOfEnv rho))
      ns <- asks names
      trace $ "\nHole at " ++ show l ++ ":\n\n" ++
              e ++ replicate 80 '-' ++ "\n" ++ show (E.normal ns a)  ++ "\n"
  (_,Con c es) -> do
    (bs,nu) <- getLblType c a
    checks (bs,nu) es
  (VU,Pi f)       -> checkFam f
  (VU,Sigma f)    -> checkFam f
  (VU,Sum _ _ bs) -> forM_ bs $ \lbl -> case lbl of
    OLabel _ tele -> checkTele tele
    PLabel _ tele is ts ->
      throwError $ "check: no path constructor allowed in " ++ show t
  (VU,HSum _ _ bs) -> forM_ bs $ \lbl -> case lbl of
    OLabel _ tele -> checkTele tele
    PLabel _ tele is ts -> do
      checkTele tele
      rho <- asks env
      unless (all (`elem` is) (C.domain ts)) $
        throwError "names in path label C.System" -- TODO
      mapM_ checkFresh is
      let iis = zip is (map C.Atom is)
      local (addSubs iis . addTele tele) $ do
        checkSystemWith ts $ \alpha talpha ->
          local (faceEnv alpha) $

            check (Ter t rho) talpha
        rho' <- asks env
        checkCompSystem (E.evalSystem rho' ts)
  (VPi va@(Ter (Sum _ _ cas) nu) f,Split _ _ ty ces) -> do
    check VU ty
    rho <- asks env
    unlessM (a === E.eval rho ty) $ throwError "check: split annotations"
    if map labelName cas == map branchName ces
       then sequence_ [ checkBranch (lbl,nu) f brc (Ter t rho) va
                      | (brc, lbl) <- zip ces cas ]
       else throwError "case branches does not match the data type"
  (VPi va@(Ter (HSum _ _ cas) nu) f,Split _ _ ty ces) -> do
    check VU ty
    rho <- asks env
    unlessM (a === E.eval rho ty) $ throwError "check: split annotations"
    if map labelName cas == map branchName ces
       then sequence_ [ checkBranch (lbl,nu) f brc (Ter t rho) va
                      | (brc, lbl) <- zip ces cas ]
       else throwError "case branches does not match the data type"
  (VPi a f,Lam x a' t)  -> do
    check VU a'
    ns <- asks names
    rho <- asks env
    unlessM (a === E.eval rho a') $
      throwError $ "check: lam types don't match"
        ++ "\nlambda type annotation: " ++ show a'
        ++ "\ndomain of Pi: " ++ show a
        ++ "\nnormal form of type: " ++ show (E.normal ns a)
    let var = mkVarNice ns x a

    local (addTypeVal (x,a)) $ check (E.app f var) t
  (VSigma a f, Pair t1 t2) -> do
    check a t1
    v <- evalTyping t1
    check (E.app f v) t2
  (_,Where e d) -> do
    local (\tenv@TEnv{indent=i} -> tenv{indent=i + 2}) $ checkDecls d
    local (addDecls d) $ check a e
  (VU,PathP a e0 e1) -> do
    (a0,a1) <- checkPLam (constPath VU) a
    check a0 e0
    check a1 e1
  (VPathP p a0 a1,PLam _ e) -> do
    (u0,u1) <- checkPLam p t
    ns <- asks names
    unless (E.conv ns a0 u0 && E.conv ns a1 u1) $
      throwError $ "path endpoints don't match for " ++ show e ++ ", got " ++
                   show (u0,u1) ++ ", but expected " ++ show (a0,a1)
  (VU,Glue a ts) -> do
    check VU a
    rho <- asks env
    checkGlue (E.eval rho a) ts
  (VGlue va ts,GlueElem u us) -> do
    check va u
    vu <- evalTyping u
    checkGlueElem vu ts us
  (VCompU va ves,GlueElem u us) -> do
    check va u
    vu <- evalTyping u
    checkGlueElemU vu ves us
  (VU,Id a a0 a1) -> do
    check VU a
    va <- evalTyping a
    check va a0
    check va a1
  (VId va va0 va1,IdPair w ts) -> do
    check (VPathP (constPath va) va0 va1) w
    vw <- evalTyping w
    checkSystemWith ts $ \alpha tAlpha ->
      local (faceEnv alpha) $ do
        check (va `C.face` alpha) tAlpha
        vtAlpha <- evalTyping tAlpha
        unlessM (vw `C.face` alpha === constPath vtAlpha) $
          throwError "malformed eqC"
    rho <- asks env
    checkCompSystem (E.evalSystem rho ts) -- Not needed
  _ -> do
    v <- infer t
    unlessM (v === a) $
      throwError $ "check conv:\n" ++ show v ++ "\n/=\n" ++ show a

checkDecls :: Decls -> Typing ()
checkDecls (MutualDecls _ []) = return ()
checkDecls (MutualDecls l d)  = do
  a <- asks env
  let (idents,tele,ters) = (declIdents d,declTele d,declTers d)
  ind <- asks indent
  trace (replicate ind ' ' ++ "Checking: " ++ unwords idents)
  checkTele tele
  local (addDecls (MutualDecls l d)) $ do
    rho <- asks env
    checks (tele,rho) ters
checkDecls (OpaqueDecl _)      = return ()
checkDecls (TransparentDecl _) = return ()
checkDecls TransparentAllDecl  = return ()

checkTele :: Tele -> Typing ()
checkTele []          = return ()
checkTele ((x,a):xas) = do
  check VU a
  local (addType (x,a)) $ checkTele xas

checkFam :: Ter -> Typing ()
checkFam (Lam x a b) = do
  check VU a
  local (addType (x,a)) $ check VU b
checkFam x = throwError $ "checkFam: " ++ show x

checkCompSystem :: C.System Val -> Typing ()
checkCompSystem vus = do
  ns <- asks names
  unless (E.isCompSystem ns vus)
    (throwError $ "Incompatible System " ++ C.showSystem vus)

checkSystemsWith :: C.System a -> C.System b -> (C.Face -> a -> b -> Typing c) ->
                    Typing ()
checkSystemsWith us vs f = sequence_ $ elems $ intersectionWithKey f us vs

checkSystemWith :: C.System a -> (C.Face -> a -> Typing b) -> Typing ()
checkSystemWith us f = sequence_ $ elems $ mapWithKey f us

checkGlueElem :: Val -> C.System Val -> C.System Ter -> Typing ()
checkGlueElem vu ts us = do
  unless (keys ts == keys us)
    (throwError ("Keys don't match in " ++ show ts ++ " and " ++ show us))
  rho <- asks env
  checkSystemsWith ts us
    (\alpha vt u -> local (faceEnv alpha) $ check (E.equivDom vt) u)
  let vus = E.evalSystem rho us
  checkSystemsWith ts vus (\alpha vt vAlpha ->
    unlessM (E.app (E.equivFun vt) vAlpha === (vu `C.face` alpha)) $
      throwError $ "Image of glue component " ++ show vAlpha ++
                   " doesn't match " ++ show vu)
  checkCompSystem vus

checkGlueElemU :: Val -> C.System Val -> C.System Ter -> Typing ()
checkGlueElemU vu ves us = do
  unless (keys ves == keys us)
    (throwError ("Keys don't match in " ++ show ves ++ " and " ++ show us))
  rho <- asks env
  checkSystemsWith ves us
    (\alpha ve u -> local (faceEnv alpha) $ check (ve E.@@ C.One) u)
  let vus = E.evalSystem rho us
  checkSystemsWith ves vus (\alpha ve vAlpha ->
    unlessM (E.eqFun ve vAlpha === (vu `C.face` alpha)) $
      throwError $ "Transport of glueElem (for compU) component " ++ show vAlpha ++
                   " doesn't match " ++ show vu)
  checkCompSystem vus

checkGlue :: Val -> C.System Ter -> Typing ()
checkGlue va ts = do
  checkSystemWith ts (\alpha tAlpha -> checkEquiv (va `C.face` alpha) tAlpha)
  rho <- asks env
  checkCompSystem (E.evalSystem rho ts)

mkIso :: Val -> Val
mkIso vb = E.eval rho $
  Sigma $ Lam "a" U $
  Sigma $ Lam "f" (Pi (Lam "_" a b)) $
  Sigma $ Lam "g" (Pi (Lam "_" b a)) $
  Sigma $ Lam "s" (Pi (Lam "y" b $ PathP (PLam (C.Name "_") b) (App f (App g y)) y)) $
    Pi (Lam "x" a $ PathP (PLam (C.Name "_") a) (App g (App f x)) x)
  where [a,b,f,g,x,y] = map Var ["a","b","f","g","x","y"]
        rho = upd ("b",vb) emptyEnv

mkEquiv :: Val -> Val
mkEquiv va = E.eval rho $
  Sigma $ Lam "t" U $
  Sigma $ Lam "f" (Pi (Lam "_" t a)) $
  Pi (Lam "x" a $ iscontrfib)
  where [a,b,f,x,y,s,t,z] = map Var ["a","b","f","x","y","s","t","z"]
        rho = upd ("a",va) emptyEnv
        fib = Sigma $ Lam "y" t (PathP (PLam (C.Name "_") a) x (App f y))
        iscontrfib = Sigma $ Lam "s" fib $
                     Pi $ Lam "z" fib $ PathP (PLam (C.Name "_") fib) s z

checkEquiv :: Val -> Ter -> Typing ()
checkEquiv va equiv = check (mkEquiv va) equiv

checkIso :: Val -> Ter -> Typing ()
checkIso vb iso = check (mkIso vb) iso

checkBranch :: (Label,Env) -> Val -> Branch -> Val -> Val -> Typing ()
checkBranch (OLabel _ tele,nu) f (OBranch c ns e) _ _ = do
  ns' <- asks names
  let us = map snd $ mkVars ns' tele nu
  local (addBranch (zip ns us) nu) $ check (E.app f (VCon c us)) e
checkBranch (PLabel _ tele is ts,nu) f (PBranch c ns js e) g va = do
  ns' <- asks names

  let us   = mkVars ns' tele nu
      vus  = map snd us
      js'  = map C.Atom js
      vts  = E.evalSystem (subs (zip is js') (upds us nu)) ts
      vgts = intersectionWith E.app (C.border g vts) vts
  local (addSubs (zip js js') . addBranch (zip ns vus) nu) $ do
    check (E.app f (VPCon c va vus js')) e
    ve  <- evalTyping e -- TODO: combine with next two lines?
    let veborder = C.border ve vts
    unlessM (veborder === vgts) $
      throwError $ "Faces in branch for " ++ show c ++ " don't match:"
                   ++ "\ngot\n" ++ C.showSystem veborder ++ "\nbut expected\n"
                   ++ C.showSystem vgts

checkFormula :: C.Formula -> Typing ()
checkFormula phi = do
  rho <- asks env
  let dom = domainEnv rho
  unless (all (`elem` dom) (C.support phi)) $
    throwError $ "checkFormula: " ++ show phi

checkFresh :: C.Name -> Typing ()
checkFresh i = do
  rho <- asks env
  when (i `elem` C.support rho)
    (throwError $ show i ++ " is already declared")

checkPLam :: Val -> Ter -> Typing (Val,Val)
checkPLam v (PLam i a) = do
  rho <- asks env

  local (addSub (i,C.Atom i)) $ check (v E.@@ i) a
  return (E.eval (sub (i,C.Dir 0) rho) a, E.eval (sub (i,C.Dir 1) rho) a)
checkPLam v t = do
  vt <- infer t
  case vt of
    VPathP a a0 a1 -> do
      unlessM (a === v) $ throwError (
        "checkPLam\n" ++ show v ++ "\n/=\n" ++ show a)
      return (a0,a1)
    _ -> throwError $ show vt ++ " is not a path"

checkPLamSystem :: Ter -> Val -> C.System Ter -> Typing (C.System Val)
checkPLamSystem t0 va ps = do
  rho <- asks env
  v <- T.sequence $ mapWithKey (\alpha pAlpha ->
    local (faceEnv alpha) $ do
      rhoAlpha <- asks env
      (a0,a1)  <- checkPLam (va `C.face` alpha) pAlpha
      unlessM (a0 === E.eval rhoAlpha t0) $
        throwError $ "Incompatible C.System " ++ C.showSystem ps ++
                     ", component\n " ++ show pAlpha ++
                     "\nincompatible with\n " ++ show t0 ++
                     "\na0 = " ++ show a0 ++
                     "\nt0alpha = " ++ show (E.eval rhoAlpha t0) ++
                     "\nva = " ++ show va
      return a1) ps
  checkCompSystem (E.evalSystem rho ps)
  return v

checks :: (Tele,Env) -> [Ter] -> Typing ()
checks ([],_)         []     = return ()
checks ((x,a):xas,nu) (e:es) = do
  check (E.eval nu a) e
  v' <- evalTyping e
  checks (xas,upd (x,v') nu) es
checks _              _      =
  throwError "checks: incorrect number of arguments"

infer :: Ter -> Typing Val
infer e = case e of
  U           -> return VU  -- U : U
  Var n       -> E.lookType n <$> asks env
  App t u -> do
    c <- infer t
    case c of
      VPi a f -> do
        check a u
        v <- evalTyping u
        return $ E.app f v
      _       -> throwError $ show c ++ " is not a product"
  Fst t -> do
    c <- infer t
    case c of
      VSigma a f -> return a
      _          -> throwError $ show c ++ " is not a sigma-type"
  Snd t -> do
    c <- infer t
    case c of
      VSigma a f -> do
        v <- evalTyping t
        return $ E.app f (E.fstVal v)
      _          -> throwError $ show c ++ " is not a sigma-type"
  Where t d -> do
    checkDecls d
    local (addDecls d) $ infer t
  UnGlueElem e _ -> do
    t <- infer e
    case t of
     VGlue a _ -> return a
     _ -> throwError (show t ++ " is not a Glue")
  AppFormula e phi -> do
    checkFormula phi
    t <- infer e
    case t of
      VPathP a _ _ -> return $ a E.@@ phi
      _ -> throwError (show e ++ " is not a path")
  Comp a t0 ps -> do
    (va0, va1) <- checkPLam (constPath VU) a
    va <- evalTyping a
    check va0 t0
    checkPLamSystem t0 va ps
    return va1
  HComp a u0 us -> do
    check VU a
    va <- evalTyping a
    check va u0
    checkPLamSystem u0 (constPath va) us
    return va
  Fill a t0 ps -> do
    (va0, va1) <- checkPLam (constPath VU) a
    va <- evalTyping a
    check va0 t0
    checkPLamSystem t0 va ps
    vt  <- evalTyping t0
    rho <- asks env
    let vps = E.evalSystem rho ps
    return (VPathP va vt (E.compLine va vt vps))
  PCon c a es phis -> do
    check VU a
    va <- evalTyping a
    (bs,nu) <- getLblType c va
    checks (bs,nu) es
    mapM_ checkFormula phis
    return va
  IdJ a u c d x p -> do
    check VU a
    va <- evalTyping a
    check va u
    vu <- evalTyping u
    let refu = VIdPair (constPath vu) $ C.mkSystem [(C.eps,vu)]
    rho <- asks env
    let z = Var "z"
        ctype = E.eval rho $ Pi $ Lam "z" a $ Pi $ Lam "_" (Id a u z) U
    check ctype c
    vc <- evalTyping c
    check (E.app (E.app vc vu) refu) d
    check va x
    vx <- evalTyping x
    check (VId va vu vx) p
    vp <- evalTyping p
    return (E.app (E.app vc vx) vp)
  _ -> throwError ("infer " ++ show e)