Evaluation.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module PiSigma.Evaluation where

import Control.Monad.Error
import Control.Monad.State
import Control.Monad.Identity

import PiSigma.Syntax

-- * Monad used for evaluation

newtype Eval e a =
  Eval {unEval :: StateT e (ErrorT String Identity) a}
  deriving (Monad,MonadError String,MonadState e)

run :: e -> Eval e a -> Either String (a,e)
run e (Eval p) = runIdentity $ runErrorT $ runStateT p e 

catchE :: Eval e a -> (String -> Eval e a) -> Eval e a
catchE  = catchError

getEnv :: Eval e e
getEnv = get

setEnv :: e -> Eval e ()
setEnv = put

getId :: Name -> Scope -> Eval e Id
getId x s = case lookupScope s x of 
              Just i -> return i
              Nothing -> fail ("Undefined Variable: "++x++"\n")

-- | Takes a name, a scope, and potentially a type.
-- It extends the environment and the scope with that name.
decl :: Env e =>
        Name -> PrtInfo -> Scope -> Maybe (Clos Type) ->
        Eval e (Id,Scope)
decl x x' s a =
    do e <- getEnv
       let (i,e') = extE e x'
       setEnv e'
       return (i,extendScope s x (i,a))

decl' :: Env e => Name -> Scope -> Eval e (Id,Scope)
decl' x s   = decl x (PrtInfo {name = x, expand = True}) s Nothing

tdecl :: Env e => Name -> Scope -> (Clos Type) -> Eval e (Id,Scope)
tdecl x g a = decl x (PrtInfo {name = x, expand = True}) g (Just a)

-- | Updates the environment.
defn :: Env e => Id -> EnvEntry -> Eval e ()
defn i ei =
    do e <- getEnv
       setEnv (setE e i ei)

defn' :: Env e => Id -> Clos Type -> Eval e ()
defn' i t = defn i (Closure t)

-- | Locally updates the environment.
letn :: Env e => Id -> EnvEntry -> Eval e a -> Eval e a
letn i ei p =
    do eo <- lookupId i
       defn i ei
       a <- p
       defn i eo
       return a

letn' :: Env e => Id -> Clos Type -> Eval e a -> Eval e a
letn' i t p = letn i (Closure t) p

subst :: Env e => Bind (Clos Term) -> (Clos Term) -> Eval e (Clos Term)
subst (x,(t,s)) u =
    do (i,s') <- decl' x s
       defn' i u    
       return (t,s')

evalApp :: Env e => Val -> (Clos Term) -> Eval e Val
evalApp (VLam xt) u = subst xt u >>= eval
evalApp (Ne t)    u = return (Ne (t :.. u))
evalApp _         _ = fail "function expected"

lookupId :: Env e => Id -> Eval e EnvEntry
lookupId i = liftM (flip getE i) getEnv

evalId :: Env e => Id -> Eval e Val
evalId i =
    do ei <- lookupId i
       case ei of
         Closure t -> eval t
         Id j      -> return (Ne (NVar j))

-- use subst!
evalSplit :: Env e => Val -> (Bind (Bind (Clos Term))) ->
             Eval e Val
evalSplit (VPair ((l,r),s)) (x,(y,(t,s'))) = 
    do (x',s2) <- decl' x s'
       (y',s3) <- decl' y s2
       defn' x' (l,s)
       defn' y' (r,s)
       eval (t,s3)

evalSplit (Ne n) (xy,t) = return (Ne (NSplit n (xy,t)))
evalSplit _ _ = fail "Pair expected"

evalCase :: Env e => Val -> Clos [(Label,Term)] -> Eval e Val
evalCase (VLabel l) (lts,s) = 
    case lookup l lts of
      Nothing -> fail "case not matched"
      Just t  -> eval (t,s)
evalCase (Ne n) lts = return (Ne (NCase n lts))
evalCase _ _ = fail "Label expected"

force :: Env e => Val -> Eval e Val
force (VBox (Boxed c)) = eval c
force (Ne n) = return (Ne (NForce n))
force _ = fail "Box expected"

eval :: Env e => (Clos Term) -> Eval e Val
eval (Var x, s)     = getId x s >>= evalId
eval (Let g t, s)   = evalProg (g,s) >>= curry eval t
eval (Type,s)       = return VType
eval (Q ps axb,s)   = return (VQ ps (axb,s))
eval (Lift t,s)     = return (VLift (t,s))
eval (Lam (x,t), s) = return (VLam (x,(t,s)))
eval (t :. u ,s)    = eval (t,s) >>= flip evalApp (u,s)
eval (Pair tu,s)    = return (VPair (tu,s))
eval (Split t (x,(y,u)),s)
                    = eval (t,s) >>= flip evalSplit (x,(y,(u,s)))
eval (Enum ls,s)    = return (VEnum ls)
eval (Label l,s)    = return (VLabel l)
eval (Case t lts,s) = eval (t,s) >>= flip evalCase (lts,s)
eval (Box t,s)      = return (VBox (Boxed (t,s)))
eval (Force t,s)    = eval (t,s) >>= force


-- forced eval (should be integrated into eval?)
feval :: Env e => (Clos Term) -> Eval e Val
feval t = do t' <- eval t
             case t' of 
               VBox (Boxed u) -> feval u
               v -> return v

               

evalProg :: Env e => Clos Prog -> Eval e Scope
evalProg ([],s) = return s
evalProg ((Decl x _):tel, s) = do (_,s') <- decl x (PrtInfo x False) s Nothing
                                  evalProg (tel,s')
evalProg ((Defn x t):tel, s) = do i <- getId x s
                                  defn' i (t,s)
                                  evalProg (tel,s)