content: {-# LANGUAGE TupleSections #-}
  module TypeChecker 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 Connections
  import CTT
  import Eval
  
  -- Type checking monad
  type Typing a = ReaderT TEnv (ExceptT String IO) a
  
  -- Environment for type checker
  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 function that depends on the verbosity flag
  trace :: String -> Typing ()
  trace s = do
    b <- asks verbose
    when b $ liftIO (putStrLn s)
  
  -------------------------------------------------------------------------------
  -- | Functions for running computations in the type checker monad
  
  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)
  
  -------------------------------------------------------------------------------
  -- | Modifiers for the environment
  
  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 :: (Name,Formula) -> TEnv -> TEnv
  addSub iphi (TEnv ns ind rho v) = TEnv ns ind (sub iphi rho) v
  
  addSubs :: [(Name,Formula)] -> TEnv -> TEnv
  addSubs = flip $ foldr addSub
  
  addType :: (Ident,Ter) -> TEnv -> TEnv
  addType (x,a) tenv@(TEnv _ _ rho _) = addTypeVal (x,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 :: Face -> TEnv -> TEnv
  faceEnv alpha tenv = tenv{env=env tenv `face` alpha}
  
  -------------------------------------------------------------------------------
  -- | Various useful functions
  
  -- Extract the type of a label as a closure
  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)
  
  -- Monadic version of unless
  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 (eval nu a)
    in (x,w) : mkVars (n:ns) xas (upd (x,w) nu)
  
  -- Test if two values are convertible
  (===) :: Convertible a => a -> a -> Typing Bool
  u === v = conv <$> asks names <*> pure u <*> pure v
  
  -- eval in the typing monad
  evalTyping :: Ter -> Typing Val
  evalTyping t = eval <$> asks env <*> pure t
  
  -------------------------------------------------------------------------------
  -- | The bidirectional type checker
  
  -- Check that t has type a
  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 (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) (domain ts)) $
          throwError "names in path label system" -- TODO
        mapM_ checkFresh is
        let iis = zip is (map Atom is)
        local (addSubs iis . addTele tele) $ do
          checkSystemWith ts $ \alpha talpha ->
            local (faceEnv alpha) $
              -- NB: the type doesn't depend on is
              check (Ter t rho) talpha
          rho' <- asks env
          checkCompSystem (evalSystem rho' ts)
    (VPi va@(Ter (Sum _ _ cas) nu) f,Split _ _ ty ces) -> do
      check VU ty
      rho <- asks env
      unlessM (a === 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 === 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 === 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 (normal ns a)
      let var = mkVarNice ns x a
  
      local (addTypeVal (x,a)) $ check (app f var) t
    (VSigma a f, Pair t1 t2) -> do
      check a t1
      v <- evalTyping t1
      check (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 (conv ns a0 u0 && 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 (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 `face` alpha) tAlpha
          vtAlpha <- evalTyping tAlpha
          unlessM (vw `face` alpha === constPath vtAlpha) $
            throwError "malformed eqC"
      rho <- asks env
      checkCompSystem (evalSystem rho ts) -- Not needed
    _ -> do
      v <- infer t
      unlessM (v === a) $
        throwError $ "check conv:\n" ++ show v ++ "\n/=\n" ++ show a
  
  -- Check a list of declarations
  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 ()
  
  -- Check a telescope
  checkTele :: Tele -> Typing ()
  checkTele []          = return ()
  checkTele ((x,a):xas) = do
    check VU a
    local (addType (x,a)) $ checkTele xas
  
  -- Check a family
  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
  
  -- Check that a system is compatible
  checkCompSystem :: System Val -> Typing ()
  checkCompSystem vus = do
    ns <- asks names
    unless (isCompSystem ns vus)
      (throwError $ "Incompatible system " ++ showSystem vus)
  
  -- Check the values at corresponding faces with a function, assumes
  -- systems have the same faces
  checkSystemsWith :: System a -> System b -> (Face -> a -> b -> Typing c) ->
                      Typing ()
  checkSystemsWith us vs f = sequence_ $ elems $ intersectionWithKey f us vs
  
  -- Check the faces of a system
  checkSystemWith :: System a -> (Face -> a -> Typing b) -> Typing ()
  checkSystemWith us f = sequence_ $ elems $ mapWithKey f us
  
  -- Check a glueElem
  checkGlueElem :: Val -> System Val -> 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 (equivDom vt) u)
    let vus = evalSystem rho us
    checkSystemsWith ts vus (\alpha vt vAlpha ->
      unlessM (app (equivFun vt) vAlpha === (vu `face` alpha)) $
        throwError $ "Image of glue component " ++ show vAlpha ++
                     " doesn't match " ++ show vu)
    checkCompSystem vus
  
  -- Check a glueElem against VComp _ ves
  checkGlueElemU :: Val -> System Val -> 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 @@ One) u)
    let vus = evalSystem rho us
    checkSystemsWith ves vus (\alpha ve vAlpha ->
      unlessM (eqFun ve vAlpha === (vu `face` alpha)) $
        throwError $ "Transport of glueElem (for compU) component " ++ show vAlpha ++
                     " doesn't match " ++ show vu)
    checkCompSystem vus
  
  checkGlue :: Val -> System Ter -> Typing ()
  checkGlue va ts = do
    checkSystemWith ts (\alpha tAlpha -> checkEquiv (va `face` alpha) tAlpha)
    rho <- asks env
    checkCompSystem (evalSystem rho ts)
  
  -- An iso for a type b is a five-tuple: (a,f,g,s,t)   where
  --  a : U
  --  f : a -> b
  --  g : b -> a
  --  s : forall (y : b), f (g y) = y
  --  t : forall (x : a), g (f x) = x
  mkIso :: Val -> Val
  mkIso vb = 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 (Name "_") b) (App f (App g y)) y)) $
      Pi (Lam "x" a $ PathP (PLam (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
  
  -- An equivalence for a type a is a triple (t,f,p) where
  -- t : U
  -- f : t -> a
  -- p : (x : a) -> isContr ((y:t) * Id a x (f y))
  -- with isContr c = (z : c) * ((z' : C) -> Id c z z')
  mkEquiv :: Val -> Val
  mkEquiv va = 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 (Name "_") a) x (App f y))
          iscontrfib = Sigma $ Lam "s" fib $
                       Pi $ Lam "z" fib $ PathP (PLam (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 (app f (VCon c us)) e
  checkBranch (PLabel _ tele is ts,nu) f (PBranch c ns js e) g va = do
    ns' <- asks names
    -- mapM_ checkFresh js
    let us   = mkVars ns' tele nu
        vus  = map snd us
        js'  = map Atom js
        vts  = evalSystem (subs (zip is js') (upds us nu)) ts
        vgts = intersectionWith app (border g vts) vts
    local (addSubs (zip js js') . addBranch (zip ns vus) nu) $ do
      check (app f (VPCon c va vus js')) e
      ve  <- evalTyping e -- TODO: combine with next two lines?
      let veborder = border ve vts
      unlessM (veborder === vgts) $
        throwError $ "Faces in branch for " ++ show c ++ " don't match:"
                     ++ "\ngot\n" ++ showSystem veborder ++ "\nbut expected\n"
                     ++ showSystem vgts
  
  checkFormula :: Formula -> Typing ()
  checkFormula phi = do
    rho <- asks env
    let dom = domainEnv rho
    unless (all (`elem` dom) (support phi)) $
      throwError $ "checkFormula: " ++ show phi
  
  checkFresh :: Name -> Typing ()
  checkFresh i = do
    rho <- asks env
    when (i `elem` support rho)
      (throwError $ show i ++ " is already declared")
  
  -- Check that a term is a PLam and output the source and target
  checkPLam :: Val -> Ter -> Typing (Val,Val)
  checkPLam v (PLam i a) = do
    rho <- asks env
    -- checkFresh i
    local (addSub (i,Atom i)) $ check (v @@ i) a
    return (eval (sub (i,Dir 0) rho) a,eval (sub (i,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"
  
  -- Return system such that:
  --   rhoalpha |- p_alpha : Id (va alpha) (t0 rhoalpha) ualpha
  -- Moreover, check that the system ps is compatible.
  checkPLamSystem :: Ter -> Val -> System Ter -> Typing (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 `face` alpha) pAlpha
        unlessM (a0 === eval rhoAlpha t0) $
          throwError $ "Incompatible system " ++ showSystem ps ++
                       ", component\n " ++ show pAlpha ++
                       "\nincompatible with\n " ++ show t0 ++
                       "\na0 = " ++ show a0 ++
                       "\nt0alpha = " ++ show (eval rhoAlpha t0) ++
                       "\nva = " ++ show va
        return a1) ps
    checkCompSystem (evalSystem rho ps)
    return v
  
  checks :: (Tele,Env) -> [Ter] -> Typing ()
  checks ([],_)         []     = return ()
  checks ((x,a):xas,nu) (e:es) = do
    check (eval nu a) e
    v' <- evalTyping e
    checks (xas,upd (x,v') nu) es
  checks _              _      =
    throwError "checks: incorrect number of arguments"
  
  -- infer the type of e
  infer :: Ter -> Typing Val
  infer e = case e of
    U           -> return VU  -- U : U
    Var n       -> 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 $ 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 $ app f (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 @@ 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 = evalSystem rho ps
      return (VPathP va vt (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) $ mkSystem [(eps,vu)]
      rho <- asks env
      let z = Var "z"
          ctype = eval rho $ Pi $ Lam "z" a $ Pi $ Lam "_" (Id a u z) U
      check ctype c
      vc <- evalTyping c
      check (app (app vc vu) refu) d
      check va x
      vx <- evalTyping x
      check (VId va vu vx) p
      vp <- evalTyping p
      return (app (app vc vx) vp)
    _ -> throwError ("infer " ++ show e)
  
  -- Not used since we have U : U
  --
  -- (=?=) :: Typing Ter -> Ter -> Typing ()
  -- m =?= s2 = do
  --   s1 <- m
  --   unless (s1 == s2) $ throwError (show s1 ++ " =/= " ++ show s2)
  --
  -- checkTs :: [(String,Ter)] -> Typing ()
  -- checkTs [] = return ()
  -- checkTs ((x,a):xas) = do
  --   checkType a
  --   local (addType (x,a)) (checkTs xas)
  --
  -- checkType :: Ter -> Typing ()
  -- checkType t = case t of
  --   U              -> return ()
  --   Pi a (Lam x b) -> do
  --     checkType a
  --     local (addType (x,a)) (checkType b)
  --   _ -> infer t =?= U
  