newt/src/Lib/Eval.newt

module Lib.Eval

import Prelude
import Lib.Common
import Lib.Prettier
import Lib.Types
import Lib.TopContext

import Data.IORef
import Data.SnocList
import Data.SortedMap

eval : Env -> Tm -> M Val

-- REVIEW everything is evalutated whether it's needed or not
-- It would be nice if the environment were lazy.
-- e.g. case is getting evaluated when passed to a function because
-- of dependencies in pi-types, even if the dependency isn't used

infixl 8 _$$_

_$$_ : Closure -> Val -> M Val
_$$_ (MkClosure env tm) u = eval (u :: env) tm

vapp : Val -> Val -> M Val
vapp (VLam _ _ _ _ t) u = t $$ u
vapp (VVar fc k sp) u = pure $ VVar fc k (sp :< u)
vapp (VRef fc nm sp) u = pure $ VRef fc nm (sp :< u)
vapp (VMeta fc k sp) u = pure $ VMeta fc k (sp :< u)
vapp t u = error' "impossible in vapp \{show t} to \{show u}\n"

vappSpine : Val -> SnocList Val -> M Val
vappSpine t Lin = pure t
vappSpine t (xs :< x) = do
  rest <- vappSpine t xs
  vapp rest x

lookupVar : Env -> Int -> Maybe Val
lookupVar env k = let l = cast $ length env in
  if k > l
    then Nothing
    else case getAt (cast $ lvl2ix l k) env of
      Just v@(VVar fc k' sp) => if k == k' then Nothing else Just v
      Just v => Just v
      Nothing => Nothing

-- hoping to apply what we got via pattern matching
-- TODO see if we can drop this after updating pattern matching
unlet : Env -> Val -> M Val
unlet env t@(VVar fc k sp) = case lookupVar env k of
      Just tt@(VVar fc' k' sp') => do
        debug $ \ _ =>  "lookup \{show k} is \{show tt}"
        if k' == k then pure t else (vappSpine (VVar fc' k' sp') sp >>= unlet env)
      Just t => vappSpine t sp >>= unlet env
      Nothing => do
        debug $ \ _ =>  "lookup \{show k} is Nothing in env \{show env}"
        pure t
unlet env x = pure x

-- Try applying VRef to spine, back out if it is stuck
tryEval : Env -> Val -> M (Maybe Val)
tryEval env (VRef fc k sp) = do
  top <- getTop
  case lookup k top of
      Just (MkEntry _ name ty (Fn tm) _) =>
        catchError (
            do
            debug $ \ _ =>  "app \{show name} to \{show sp}"
            vtm <- eval env tm
            debug $ \ _ =>  "tm is \{render 90 $ pprint Nil tm}"
            val <- vappSpine vtm sp
            case val of
              VCase _ _ _ => pure Nothing
              VLam _ _ _ _ _ => pure Nothing
              -- For now? There is a spot in Compile.newt that has
              -- two applications of fresh that is getting unfolded even
              -- though it has the same head and spine.  Possibly because it's
              -- coming out of a let and is instantly applied
              VLetRec _ _ _ _ _ => pure Nothing
              v => pure $ Just v)
            (\ _ => pure Nothing)
      _ => do
        debug $ \ _ => "tryEval blocked on undefined \{show k}"
        pure Nothing
tryEval _ _ = pure Nothing

-- Force far enough to compare types

forceType : Env -> Val -> M Val
forceType env (VMeta fc ix sp) = do
  meta <- lookupMeta ix
  case meta of
    (Solved _ k t) => vappSpine t sp >>= forceType env
    _ => pure (VMeta fc ix sp)

forceType env x = do
  Just x' <- tryEval env x
    | _ => pure x
  forceType env x'

-- for cases applied to a value
-- TODO this does not handle CaseLit
evalCase : Env -> Val -> List CaseAlt -> M (Maybe Val)
evalCase env sc@(VRef _ nm sp) (cc@(CaseCons name nms t) :: xs) = do
  top <- getTop
  if nm == name
    then do
      debug $ \ _ =>  "ECase \{show nm} \{show sp} \{show nms} \{showTm t}"
      pushArgs env (sp <>> Nil) nms
    else case lookup nm top of
      (Just (MkEntry _ str type (DCon _ k str1) _)) => evalCase env sc xs
      -- bail for a stuck function
      _ => pure Nothing
  where
    pushArgs : Env -> List Val -> List String -> M (Maybe Val)
    pushArgs env (arg :: args) (nm :: nms) = pushArgs (arg :: env) args nms
    pushArgs env args Nil = do
      t' <- eval env t
      Just <$> vappSpine t' (Lin <>< args)
    pushArgs env Nil rest = pure Nothing
-- REVIEW - this is handled in the caller already
evalCase env sc@(VVar fc k sp) alts = case lookupVar env k of
    Just tt@(VVar fc' k' sp') => do
        debug $ \ _ =>  "lookup \{show k} is \{show tt}"
        if k' == k
          then pure Nothing
          else do
            val <- vappSpine (VVar fc' k' sp') sp
            evalCase env val alts
    Just t => do
      val <- vappSpine t sp
      evalCase env val alts
    Nothing => do
        debug $ \ _ =>  "lookup \{show k} is Nothing in env \{show env}"
        pure Nothing
evalCase env sc (CaseDefault u :: xs) = Just <$> eval env u
evalCase env sc cc = do
  debug $ \ _ =>  "CASE BAIL sc \{show sc} vs " -- \{show cc}"
  debug $ \ _ =>  "env is \{show env}"
  pure Nothing

-- neutral alts
evalAlt : Env → CaseAlt → M VCaseAlt
evalAlt env (CaseDefault tm) = VCaseDefault <$> eval env tm
evalAlt env (CaseLit lit tm) = VCaseLit lit <$> eval env tm
-- in the cons case, we're binding args
evalAlt env (CaseCons nm args tm) = pure $ VCaseCons nm args env tm

-- So smalltt says:
--   Smalltt has the following approach:
--   - Top-level and local definitions are lazy.
--   - We instantiate Pi types during elaboration with lazy values.
--   - Applications headed by top-level variables are lazy.
--   - Any other function application is call-by-value during evaluation.

-- TODO maybe add glueing

eval env (Ref fc x) = pure $ VRef fc x Lin
eval env (App _ t u) = do
  t' <- eval env t
  u' <- eval env u
  vapp t' u'
eval env (UU fc) = pure (VU fc)
eval env (Erased fc) = pure (VErased fc)
eval env (Meta fc i) = do
  meta <- lookupMeta i
  case meta of
        (Solved _ k t) => pure $ t
        _ => pure $ VMeta fc i Lin
eval env (Lam fc x icit rig t) = pure $ VLam fc x icit rig (MkClosure env t)
eval env (Pi fc x icit rig a b) = do
  a' <- eval env a
  pure $ VPi fc x icit rig a' (MkClosure env b)
eval env (Let fc nm t u) = do
  t' <- eval env t
  u' <- eval (VVar fc (length' env) Lin :: env) u
  pure $ VLet fc nm t' u'
eval env (LetRec fc nm ty t u) = do
  ty' <- eval env ty
  t' <- eval (VVar fc (length' env) Lin :: env) t
  u' <- eval (VVar fc (length' env) Lin :: env) u
  pure $ VLetRec fc nm ty' t' u'
-- Here, we assume env has everything. We push levels onto it during type checking.
-- I think we could pass in an l and assume everything outside env is free and
-- translate to a level
eval env (Bnd fc i) = case getAt' i env of
  Just rval => pure rval
  Nothing => error fc "Bad deBruin index \{show i}"
eval env (Lit fc lit) = pure $ VLit fc lit

eval env tm@(Case fc sc alts) = do
  -- TODO we need to be able to tell eval to expand aggressively here.
  sc' <- eval env sc
  sc' <- unlet env sc' -- try to expand lets from pattern matching
  -- possibly too aggressive?
  sc' <- forceType env sc'
  Nothing <- evalCase env sc' alts
    | Just v => pure v

  vsc <- eval env sc
  alts' <- traverse (evalAlt env) alts
  pure $ VCase fc vsc alts'

quote : (lvl : Int) -> Val -> M Tm

quoteSp : (lvl : Int) -> Tm -> SnocList Val -> M Tm
quoteSp lvl t Lin = pure t
quoteSp lvl t (xs :< x) = do
  t' <- quoteSp lvl t xs
  x' <- quote lvl x
  pure $ App emptyFC t' x'

quoteAlt : Int → VCaseAlt → M CaseAlt
quoteAlt l (VCaseDefault val) = CaseDefault <$> quote l val
quoteAlt l (VCaseLit lit val) = CaseLit lit <$> quote l val
quoteAlt l (VCaseCons nm args env tm) = do
  val <- eval (mkenv l env args) tm
  tm <- quote (length' args + l) val
  pure $ CaseCons nm args tm
  where
    mkenv : Int → Env → List String → Env
    mkenv l env Nil = env
    mkenv l env (n :: ns) = mkenv (l + 1) (VVar emptyFC l Lin :: env) ns

quote l (VVar fc k sp) = if k < l
  then quoteSp l (Bnd fc (lvl2ix l k )) sp -- level to index
  else error fc "Bad index in quote \{show k} depth \{show l}"
quote l (VMeta fc i sp) = do
  meta <- lookupMeta i
  case meta of
        (Solved _ k t) => vappSpine t sp >>= quote l
        _ => quoteSp l (Meta fc i) sp
quote l (VLam fc x icit rig t) = do
  val <- t $$ VVar emptyFC l Lin
  tm <- quote (1 + l) val
  pure $ Lam fc x icit rig tm
quote l (VPi fc x icit rig a b) = do
  a' <- quote l a
  val <- b $$ VVar emptyFC l Lin
  tm <- quote (1 + l) val
  pure $ Pi fc x icit rig a' tm
quote l (VLet fc nm t u) = do
  t' <- quote l t
  u' <- quote (1 + l) u
  pure $ Let fc nm t' u'
quote l (VLetRec fc nm ty t u) = do
  ty' <- quote l ty
  t' <- quote (1 + l) t
  u' <- quote (1 + l) u
  pure $ LetRec fc nm ty' t' u'
quote l (VU fc) = pure (UU fc)
quote l (VRef fc n sp) = quoteSp l (Ref fc n) sp
quote l (VCase fc sc valts) = do
  sc' <- quote l sc
  alts <- traverse (quoteAlt l) valts
  pure $ Case fc sc' alts
quote l (VLit fc lit) = pure $ Lit fc lit
quote l (VErased fc) = pure $ Erased fc

-- Can we assume closed terms?
-- ezoo only seems to use it at Nil, but essentially does this:

nf : Env -> Tm -> M Tm
nf env t = eval env t >>= quote (length' env)

prvalCtx : {{ctx : Context}} -> Val -> M String
prvalCtx {{ctx}} v = do
  tm <- quote ctx.lvl v
  pure $ render 90 $ pprint (map fst ctx.types) tm

-- REVIEW - might be easier if we inserted the meta without a bunch of explicit App
-- I believe Kovacs is doing that.

-- we need to walk the whole thing
-- meta in Tm have a bunch of args, which should be the relevant
-- parts of the scope. So, meta has a bunch of lambdas, we've got a bunch of
-- args and we need to beta reduce, which seems like a lot of work for nothing
-- Could we put the "good bits" of the Meta in there and write it to Bnd directly
-- off of scope? I guess this might get dicey when a meta is another meta applied
-- to something.

-- ok, so we're doing something that looks lot like eval, having to collect args,
-- pull the def, and apply spine. Eval is trying for WHNF, so it doesn't walk the
-- whole thing. (We'd like to insert metas inside lambdas.)

zonk : TopContext -> Int -> Env -> Tm -> M Tm

zonkBind : TopContext -> Int -> Env -> Tm -> M Tm
zonkBind top l env tm = zonk top (1 + l) (VVar (getFC tm) l Lin :: env) tm

-- I don't know if app needs an FC...

appSpine : Tm -> List Tm -> Tm
appSpine t Nil = t
appSpine t (x :: xs) = appSpine (App (getFC t) t x) xs

-- REVIEW When metas are subst in, the fc point elsewhere
-- We might want to update when it is solved and update recursively?
-- For errors, I think we want to pretend the code has been typed in place
tweakFC : FC -> Tm -> Tm
tweakFC fc (Bnd fc1 k) = Bnd fc k
tweakFC fc (Ref fc1 nm) = Ref fc nm
tweakFC fc (UU fc1) = UU fc
tweakFC fc (Meta fc1 k) = Meta fc k
tweakFC fc (Lam fc1 nm icit rig t) = Lam fc nm icit rig t
tweakFC fc (App fc1 t u) = App fc t u
tweakFC fc (Pi fc1 nm icit x t u) = Pi fc nm icit x t u
tweakFC fc (Case fc1 t xs) = Case fc t xs
tweakFC fc (Let fc1 nm t u) = Let fc nm t u
tweakFC fc (LetRec fc1 nm ty t u) = LetRec fc nm ty t u
tweakFC fc (Lit fc1 lit) = Lit fc lit
tweakFC fc (Erased fc1) = Erased fc

zonkApp : TopContext -> Int -> Env -> Tm -> List Tm -> M Tm
zonkApp top l env (App fc t u) sp = do
  u' <- zonk top l env u
  zonkApp top l env t (u' :: sp)
zonkApp top l env t@(Meta fc k) sp = do
  meta <- lookupMeta k
  case meta of
    (Solved _ j v) => do
      sp' <- traverse (eval env) sp
      debug $ \ _ =>  "zonk \{show k} -> \{show v} spine \{show sp'}"
      foo <- vappSpine v (Lin <>< sp')
      debug $ \ _ =>  "-> result is \{show foo}"
      tweakFC fc <$> quote l foo
    _ => pure $ appSpine t sp
zonkApp top l env t sp = do
  t' <- zonk top l env t
  -- inlining
  let (Just tm) = inlineDef t' | _ => pure $ appSpine t' sp
  sp' <- traverse (eval env) sp
  vtm <- eval env tm
  catchError (do
    foo <- vappSpine vtm (Lin <>< sp')
    t' <- quote l foo
    zonk top l env t')
    (\_ => pure $ appSpine t' sp)
  where
    -- lookup name and return Def if flagged inline
    inlineDef : Tm → Maybe Tm
    inlineDef (Ref _ nm) = case lookup nm top of
      Just (MkEntry _ _ ty (Fn tm) flags) => if elem Inline flags then Just tm else Nothing
      _ => Nothing
    inlineDef _ = Nothing

zonkAlt : TopContext -> Int -> Env -> CaseAlt -> M CaseAlt
zonkAlt top l env (CaseDefault t) = CaseDefault <$> zonkBind top l env t
zonkAlt top l env (CaseLit lit t) = CaseLit lit <$> zonkBind top l env t
zonkAlt top l env (CaseCons name args t) = CaseCons name args <$> go l env args t
  where
    go : Int -> Env -> List String -> Tm -> M Tm
    go l env Nil tm = zonk top l env t
    go l env (x :: xs) tm = go (1 + l) (VVar (getFC tm) l Lin :: env) xs tm

zonk top l env t =
  let env' = VVar emptyFC l Lin :: env in
  case t of
  (Meta fc k) => zonkApp top l env t Nil
  (Lam fc nm icit rig u) => Lam fc nm icit rig <$> (zonk top (1 + l) env' u)
  (App fc _ _) => zonkApp top l env t Nil
  (Pi fc nm icit rig a b) => Pi fc nm icit rig <$> zonk top l env a <*> zonkBind top (l + 1) env' b
  (Let fc nm t u) => Let fc nm <$> zonk top l env t <*> zonkBind top (l + 1) env' u
  (LetRec fc nm ty t u) => LetRec fc nm <$> zonk top l env ty <*> zonkBind top (l + 1) env' t <*> zonkBind top (l + 1) env' u
  (Case fc sc alts) => do
      sc' <- zonk top l env sc
      alts' <- traverse (zonkAlt top l env) alts
      pure $ Case fc sc' alts'

  UU fc => pure $ UU fc
  Lit fc lit => pure $ Lit fc lit
  Bnd fc ix => pure $ Bnd fc ix
  Ref fc ix => pure $ Ref fc ix
  Erased fc => pure $ Erased fc