litTypeOf,

-- | compute type of literal
litTypeOf :: Literal -> PrimTy
litTypeOf (LNat _) = PtUnsInt PIsize
litTypeOf (LIntSz _) = PtIntSz
litTypeOf (LPrimTy _) = PtType

  unify,

  typeInf,
  typeChk,

import GardGround.Syntax.Literal (Literal, lookupBinder)

import GardGround.Syntax.Literal (Literal(..), PrimTy(..), lookupBinder, litTypeOf)

  | VEInvalidLamTy Value

  | VEInferGlobal Ident
  | VEInferLam Expr
  | VEUnifyFailed Value Value

unify :: Natural -> Value -> Value -> Either ValueError ()
unify i l r =
  let vf = vfree $ Quote i in
  case (l, r) of
    (VLam f, VLam f') -> do
      v  <- f  vf
      v' <- f' vf
      unify (i + 1) v v'
    (VPi a f, VPi a' f') -> do
      unify i a a'
      v  <- f  vf
      v' <- f' vf
      unify (i + 1) v v'
    (VNeutral (Neutral nr xs), VNeutral (Neutral nr' xs')) -> do
      lres <- unifyList xs xs'
      if (nr /= nr') || (not lres) then (Left $ VEUnifyFailed (VNeutral $ Neutral nr xs) (VNeutral $ Neutral nr' xs')) else pure ()
    _ -> Left $ VEUnifyFailed l r
  where
    unifyList :: [Value] -> [Value] -> Either ValueError Bool
    unifyList [] [] = Right True
    unifyList [] (_:_) = Right False
    unifyList (_:_) [] = Right False
    unifyList (x:xs) (y:ys) = unify i x y >> unifyList xs ys
evalTenv :: [(Value, Value)] -> Expr -> EvalResult
evalTenv tenv e = eval e $ fmap (\(x, _) -> x) tenv
{-# INLINE evalTenv #-}
tenvLen :: [(Value, Value)] -> Natural
tenvLen tenv = toEnum $ length tenv
{-# INLINE tenvLen #-}
-- | Try to infer the type of an expression with a given context of (values, types) on stack
typeInf :: [(Value, Value)] -> Expr -> EvalResult
typeInf tenv kind = case kind of
  ENeutralRoot (Local x) -> case lookupBinder tenv x of
    Nothing -> Left . VEOutOfRangeLocal $ x - (tenvLen tenv)
    Just (_, y) -> Right y
  ENeutralRoot (NrLiteral lit) -> Right . VNeutral $ Neutral (NrLiteral . LPrimTy $ litTypeOf lit) mempty
  ENeutralRoot (Global g) -> Left $ VEInferGlobal g
  ENeutralRoot (Quote q) -> Left $ VEOutOfRangeQuote q
  EAnnot e e' -> do
    v' <- evalTenv tenv e'
    typeChk tenv e v'
    Right v'
  EApply (ELam e) e' -> do
    t' <- typeInf tenv e'
    v' <- evalTenv tenv e'
    -- here, we don't have to unify the argument
    typeInf ((v', t'):tenv) e
  EApply e e' -> do
    t  <- typeInf tenv e
    case t of
      VPi arg resf -> do
        t' <- typeInf tenv e'
        unify 0 arg t'
        v' <- evalTenv tenv e'
        resf v'
      _ -> Left $ VEInvalidLamTy t
  EPi e e' -> do
    let vstar = VNeutral $ Neutral (NrLiteral $ LPrimTy PtType) mempty
    typeChk tenv e vstar
    v  <- evalTenv tenv  e
    let tenv' = (VNeutral $ Neutral (Local $ tenvLen tenv) mempty, v):tenv
    v' <- evalTenv tenv' e'
    q' <- quote 0 v'
    typeChk tenv' q' vstar
    pure vstar
  ELam f -> Left $ VEInferLam f
-- | Check if an expression matches the expected type
typeChk :: [(Value, Value)] -> Expr -> Value -> Either ValueError ()
typeChk tenv kind xty = case (kind, xty) of
  (ELam f, VPi arg res) -> do
    let llte = VNeutral $ Neutral (Local $ tenvLen tenv) mempty
    let tenv' = (llte, arg):tenv
    v <- evalTenv tenv' f
    rv <- res llte
    q <- quote 0 v
    typeChk tenv' q rv
  _ -> do
    k' <- typeInf tenv kind
    unify 0 k' xty