{-|
Copyright   : (C) 2020-2021, QBayLogic B.V.
License     : BSD2 (see the file LICENSE)
Maintainer  : QBayLogic B.V. <devops@qbaylogic.com>

The monad for partial evaluation, and its API. This should contain all
auxiliary functions needed to define new evaluator implementations. This
module is only needed to define new evaluators, for calling an existing
evaluator see Clash.Core.PartialEval.
-}

{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}

module Clash.Core.PartialEval.Monad
  ( -- * Partial Evaluation Monad
    Eval
  , runEval
    -- * Local and Global Environments
  , getLocalEnv
  , setLocalEnv
  , modifyLocalEnv
  , getGlobalEnv
  , modifyGlobalEnv
    -- * Evaluation Context
  , getContext
  , withContext
    -- * Local Type Bindings
  , getTvSubst
  , findTyVar
  , withTyVar
  , withTyVars
    -- * Local Term Bindings
  , findId
  , withId
  , withIds
  , withoutId
    -- * Global Term Bindings
  , findBinding
  , replaceBinding
    -- * IO Heap Bindings
  , getRef
  , setRef
    -- * Lifted Data Constructors
  , isKeepingLifted
  , keepLifted
    -- * Fuel
  , getFuel
  , withFuel
  , preserveFuel
    -- * Accessing Global State
  , getTyConMap
  , getInScope
    -- * Fresh Variable Generation
  , getUniqueId
  , getUniqueTyVar
    -- * Work free check
  , workFreeValue
  ) where

import           Control.Applicative (Alternative)
import           Control.Concurrent.Supply (Supply)
import           Control.Monad.Catch (MonadThrow, MonadCatch, MonadMask)
import           Control.Monad.IO.Class (MonadIO)

#if !MIN_VERSION_base(4,13,0)
import           Control.Monad.Fail (MonadFail)
#endif

import           Control.Monad.RWS.Strict (RWST, MonadReader, MonadState)
import qualified Control.Monad.RWS.Strict as RWS
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Map.Strict as Map

import           Clash.Core.HasFreeVars
import           Clash.Core.Name (OccName)
import           Clash.Core.PartialEval.AsTerm
import           Clash.Core.PartialEval.NormalForm
import           Clash.Core.Subst (Subst, mkTvSubst)
import           Clash.Core.TyCon (TyConMap)
import           Clash.Core.Type (Kind, KindOrType, Type)
import           Clash.Core.Util (mkUniqSystemId, mkUniqSystemTyVar)
import           Clash.Core.Var (Id, TyVar, Var)
import           Clash.Core.VarEnv
import           Clash.Driver.Types (Binding(..))
import           Clash.Rewrite.WorkFree (isWorkFree)

{-
NOTE [RWS monad]
~~~~~~~~~~~~~~~~
Local bindings are kept in the Reader monad and global bindings in the State
monad. This ensures that global changes are propagated to later evaluation
actions whereas local changes only exist when evaluating a particular sub-term.
For example, consider the term

   (let ... in f) (let ... in x)

When evaluating this, the let bindings in the left sub-term should not be in
scope when evaluating the right sub-term. By using only the State monad for
local and global state, too much care needs to be given to ensuring that local
bindings are saved and restored when evaluating different sub-terms.

The MonadWriter instance is deliberately not derived here, as the Writer monad
functionality of RWST is not wanted.
-}

-- TODO The inner monad here could be changed to STM to allow the evaluator
-- to work on evaluating sub-terms concurrently. That would require slightly
-- different environment types, where data can be stored in STM types.

-- | The monad of partial evaluation. The inner monad is IO, as primitive
-- evaluation can attempt to evaluate IO actions.
--
newtype Eval a = Eval
  { forall a. Eval a -> RWST LocalEnv () GlobalEnv IO a
unEval :: RWST LocalEnv () GlobalEnv IO a }
  deriving
    ( (forall a b. (a -> b) -> Eval a -> Eval b)
-> (forall a b. a -> Eval b -> Eval a) -> Functor Eval
forall a b. a -> Eval b -> Eval a
forall a b. (a -> b) -> Eval a -> Eval b
forall (f :: Type -> Type).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Eval a -> Eval b
fmap :: forall a b. (a -> b) -> Eval a -> Eval b
$c<$ :: forall a b. a -> Eval b -> Eval a
<$ :: forall a b. a -> Eval b -> Eval a
Functor
    , Functor Eval
Functor Eval =>
(forall a. a -> Eval a)
-> (forall a b. Eval (a -> b) -> Eval a -> Eval b)
-> (forall a b c. (a -> b -> c) -> Eval a -> Eval b -> Eval c)
-> (forall a b. Eval a -> Eval b -> Eval b)
-> (forall a b. Eval a -> Eval b -> Eval a)
-> Applicative Eval
forall a. a -> Eval a
forall a b. Eval a -> Eval b -> Eval a
forall a b. Eval a -> Eval b -> Eval b
forall a b. Eval (a -> b) -> Eval a -> Eval b
forall a b c. (a -> b -> c) -> Eval a -> Eval b -> Eval c
forall (f :: Type -> Type).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> Eval a
pure :: forall a. a -> Eval a
$c<*> :: forall a b. Eval (a -> b) -> Eval a -> Eval b
<*> :: forall a b. Eval (a -> b) -> Eval a -> Eval b
$cliftA2 :: forall a b c. (a -> b -> c) -> Eval a -> Eval b -> Eval c
liftA2 :: forall a b c. (a -> b -> c) -> Eval a -> Eval b -> Eval c
$c*> :: forall a b. Eval a -> Eval b -> Eval b
*> :: forall a b. Eval a -> Eval b -> Eval b
$c<* :: forall a b. Eval a -> Eval b -> Eval a
<* :: forall a b. Eval a -> Eval b -> Eval a
Applicative
    , Applicative Eval
Applicative Eval =>
(forall a. Eval a)
-> (forall a. Eval a -> Eval a -> Eval a)
-> (forall a. Eval a -> Eval [a])
-> (forall a. Eval a -> Eval [a])
-> Alternative Eval
forall a. Eval a
forall a. Eval a -> Eval [a]
forall a. Eval a -> Eval a -> Eval a
forall (f :: Type -> Type).
Applicative f =>
(forall a. f a)
-> (forall a. f a -> f a -> f a)
-> (forall a. f a -> f [a])
-> (forall a. f a -> f [a])
-> Alternative f
$cempty :: forall a. Eval a
empty :: forall a. Eval a
$c<|> :: forall a. Eval a -> Eval a -> Eval a
<|> :: forall a. Eval a -> Eval a -> Eval a
$csome :: forall a. Eval a -> Eval [a]
some :: forall a. Eval a -> Eval [a]
$cmany :: forall a. Eval a -> Eval [a]
many :: forall a. Eval a -> Eval [a]
Alternative
    , Applicative Eval
Applicative Eval =>
(forall a b. Eval a -> (a -> Eval b) -> Eval b)
-> (forall a b. Eval a -> Eval b -> Eval b)
-> (forall a. a -> Eval a)
-> Monad Eval
forall a. a -> Eval a
forall a b. Eval a -> Eval b -> Eval b
forall a b. Eval a -> (a -> Eval b) -> Eval b
forall (m :: Type -> Type).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. Eval a -> (a -> Eval b) -> Eval b
>>= :: forall a b. Eval a -> (a -> Eval b) -> Eval b
$c>> :: forall a b. Eval a -> Eval b -> Eval b
>> :: forall a b. Eval a -> Eval b -> Eval b
$creturn :: forall a. a -> Eval a
return :: forall a. a -> Eval a
Monad
    , Monad Eval
Monad Eval => (forall a. [Char] -> Eval a) -> MonadFail Eval
forall a. [Char] -> Eval a
forall (m :: Type -> Type).
Monad m =>
(forall a. [Char] -> m a) -> MonadFail m
$cfail :: forall a. [Char] -> Eval a
fail :: forall a. [Char] -> Eval a
MonadFail
    , Monad Eval
Monad Eval => (forall a. IO a -> Eval a) -> MonadIO Eval
forall a. IO a -> Eval a
forall (m :: Type -> Type).
Monad m =>
(forall a. IO a -> m a) -> MonadIO m
$cliftIO :: forall a. IO a -> Eval a
liftIO :: forall a. IO a -> Eval a
MonadIO
    , MonadReader LocalEnv
    , MonadState GlobalEnv
    , Monad Eval
Monad Eval =>
(forall e a. (HasCallStack, Exception e) => e -> Eval a)
-> MonadThrow Eval
forall e a. (HasCallStack, Exception e) => e -> Eval a
forall (m :: Type -> Type).
Monad m =>
(forall e a. (HasCallStack, Exception e) => e -> m a)
-> MonadThrow m
$cthrowM :: forall e a. (HasCallStack, Exception e) => e -> Eval a
throwM :: forall e a. (HasCallStack, Exception e) => e -> Eval a
MonadThrow
    , MonadThrow Eval
MonadThrow Eval =>
(forall e a.
 (HasCallStack, Exception e) =>
 Eval a -> (e -> Eval a) -> Eval a)
-> MonadCatch Eval
forall e a.
(HasCallStack, Exception e) =>
Eval a -> (e -> Eval a) -> Eval a
forall (m :: Type -> Type).
MonadThrow m =>
(forall e a.
 (HasCallStack, Exception e) =>
 m a -> (e -> m a) -> m a)
-> MonadCatch m
$ccatch :: forall e a.
(HasCallStack, Exception e) =>
Eval a -> (e -> Eval a) -> Eval a
catch :: forall e a.
(HasCallStack, Exception e) =>
Eval a -> (e -> Eval a) -> Eval a
MonadCatch
    , MonadCatch Eval
MonadCatch Eval =>
(forall b.
 HasCallStack =>
 ((forall a. Eval a -> Eval a) -> Eval b) -> Eval b)
-> (forall b.
    HasCallStack =>
    ((forall a. Eval a -> Eval a) -> Eval b) -> Eval b)
-> (forall a b c.
    HasCallStack =>
    Eval a
    -> (a -> ExitCase b -> Eval c) -> (a -> Eval b) -> Eval (b, c))
-> MonadMask Eval
forall b.
HasCallStack =>
((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
forall a b c.
HasCallStack =>
Eval a
-> (a -> ExitCase b -> Eval c) -> (a -> Eval b) -> Eval (b, c)
forall (m :: Type -> Type).
MonadCatch m =>
(forall b. HasCallStack => ((forall a. m a -> m a) -> m b) -> m b)
-> (forall b.
    HasCallStack =>
    ((forall a. m a -> m a) -> m b) -> m b)
-> (forall a b c.
    HasCallStack =>
    m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c))
-> MonadMask m
$cmask :: forall b.
HasCallStack =>
((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
mask :: forall b.
HasCallStack =>
((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
$cuninterruptibleMask :: forall b.
HasCallStack =>
((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
uninterruptibleMask :: forall b.
HasCallStack =>
((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
$cgeneralBracket :: forall a b c.
HasCallStack =>
Eval a
-> (a -> ExitCase b -> Eval c) -> (a -> Eval b) -> Eval (b, c)
generalBracket :: forall a b c.
HasCallStack =>
Eval a
-> (a -> ExitCase b -> Eval c) -> (a -> Eval b) -> Eval (b, c)
MonadMask
    )

-- | Evaluate an action in the partial evaluator, returning the result,
-- and the final state of the global environment.
--
runEval :: GlobalEnv -> LocalEnv -> Eval a -> IO (a, GlobalEnv)
runEval :: forall a. GlobalEnv -> LocalEnv -> Eval a -> IO (a, GlobalEnv)
runEval GlobalEnv
g LocalEnv
l Eval a
x =
  let extract :: (a, b, c) -> (a, b)
extract (a
a, b
g', c
_) = (a
a, b
g')
   in (a, GlobalEnv, ()) -> (a, GlobalEnv)
forall {a} {b} {c}. (a, b, c) -> (a, b)
extract ((a, GlobalEnv, ()) -> (a, GlobalEnv))
-> IO (a, GlobalEnv, ()) -> IO (a, GlobalEnv)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> RWST LocalEnv () GlobalEnv IO a
-> LocalEnv -> GlobalEnv -> IO (a, GlobalEnv, ())
forall r w s (m :: Type -> Type) a.
RWST r w s m a -> r -> s -> m (a, s, w)
RWS.runRWST (Eval a -> RWST LocalEnv () GlobalEnv IO a
forall a. Eval a -> RWST LocalEnv () GlobalEnv IO a
unEval Eval a
x) LocalEnv
l GlobalEnv
g
{-# INLINE runEval #-}

getLocalEnv :: Eval LocalEnv
getLocalEnv :: Eval LocalEnv
getLocalEnv = Eval LocalEnv
forall r (m :: Type -> Type). MonadReader r m => m r
RWS.ask
{-# INLINE getLocalEnv #-}

setLocalEnv :: LocalEnv -> Eval a -> Eval a
setLocalEnv :: forall a. LocalEnv -> Eval a -> Eval a
setLocalEnv = (LocalEnv -> LocalEnv) -> Eval a -> Eval a
forall a. (LocalEnv -> LocalEnv) -> Eval a -> Eval a
forall r (m :: Type -> Type) a.
MonadReader r m =>
(r -> r) -> m a -> m a
RWS.local ((LocalEnv -> LocalEnv) -> Eval a -> Eval a)
-> (LocalEnv -> LocalEnv -> LocalEnv)
-> LocalEnv
-> Eval a
-> Eval a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LocalEnv -> LocalEnv -> LocalEnv
forall a b. a -> b -> a
const
{-# INLINE setLocalEnv #-}

modifyLocalEnv :: (LocalEnv -> LocalEnv) -> Eval a -> Eval a
modifyLocalEnv :: forall a. (LocalEnv -> LocalEnv) -> Eval a -> Eval a
modifyLocalEnv = (LocalEnv -> LocalEnv) -> Eval a -> Eval a
forall a. (LocalEnv -> LocalEnv) -> Eval a -> Eval a
forall r (m :: Type -> Type) a.
MonadReader r m =>
(r -> r) -> m a -> m a
RWS.local
{-# INLINE modifyLocalEnv #-}

getGlobalEnv :: Eval GlobalEnv
getGlobalEnv :: Eval GlobalEnv
getGlobalEnv = Eval GlobalEnv
forall s (m :: Type -> Type). MonadState s m => m s
RWS.get
{-# INLINE getGlobalEnv #-}

modifyGlobalEnv :: (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv :: (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv = (GlobalEnv -> GlobalEnv) -> Eval ()
forall s (m :: Type -> Type). MonadState s m => (s -> s) -> m ()
RWS.modify'
{-# INLINE modifyGlobalEnv #-}

getContext :: Eval Id
getContext :: Eval Id
getContext = LocalEnv -> Id
lenvContext (LocalEnv -> Id) -> Eval LocalEnv -> Eval Id
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval LocalEnv
getLocalEnv

withContext :: Id -> Eval a -> Eval a
withContext :: forall a. Id -> Eval a -> Eval a
withContext Id
i = (LocalEnv -> LocalEnv) -> Eval a -> Eval a
forall a. (LocalEnv -> LocalEnv) -> Eval a -> Eval a
modifyLocalEnv LocalEnv -> LocalEnv
go
 where
  go :: LocalEnv -> LocalEnv
go LocalEnv
env = LocalEnv
env { lenvContext = i }

findTyVar :: TyVar -> Eval (Maybe Type)
findTyVar :: TyVar -> Eval (Maybe Type)
findTyVar TyVar
i = TyVar -> Map TyVar Type -> Maybe Type
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TyVar
i (Map TyVar Type -> Maybe Type)
-> (LocalEnv -> Map TyVar Type) -> LocalEnv -> Maybe Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LocalEnv -> Map TyVar Type
lenvTypes (LocalEnv -> Maybe Type) -> Eval LocalEnv -> Eval (Maybe Type)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval LocalEnv
getLocalEnv

withTyVar :: TyVar -> Type -> Eval a -> Eval a
withTyVar :: forall a. TyVar -> Type -> Eval a -> Eval a
withTyVar TyVar
i Type
a Eval a
x = do
  (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv GlobalEnv -> GlobalEnv
goGlobal
  (LocalEnv -> LocalEnv) -> Eval a -> Eval a
forall a. (LocalEnv -> LocalEnv) -> Eval a -> Eval a
modifyLocalEnv LocalEnv -> LocalEnv
goLocal Eval a
x
 where
  goGlobal :: GlobalEnv -> GlobalEnv
goGlobal env :: GlobalEnv
env@GlobalEnv{genvInScope :: GlobalEnv -> InScopeSet
genvInScope=InScopeSet
inScope} =
    let fvs :: VarSet
fvs = TyVar -> VarSet
forall a. Var a -> VarSet
unitVarSet TyVar
i VarSet -> VarSet -> VarSet
`unionVarSet` Type -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf Type
a
        iss :: InScopeSet
iss = VarSet -> InScopeSet
mkInScopeSet VarSet
fvs InScopeSet -> InScopeSet -> InScopeSet
`unionInScope` InScopeSet
inScope
     in GlobalEnv
env { genvInScope = iss }

  goLocal :: LocalEnv -> LocalEnv
goLocal env :: LocalEnv
env@LocalEnv{lenvTypes :: LocalEnv -> Map TyVar Type
lenvTypes=Map TyVar Type
types} =
    LocalEnv
env { lenvTypes = Map.insert i a types }

withTyVars :: [(TyVar, Type)] -> Eval a -> Eval a
withTyVars :: forall a. [(TyVar, Type)] -> Eval a -> Eval a
withTyVars = (Eval a -> [(TyVar, Type)] -> Eval a)
-> [(TyVar, Type)] -> Eval a -> Eval a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Eval a -> [(TyVar, Type)] -> Eval a)
 -> [(TyVar, Type)] -> Eval a -> Eval a)
-> (Eval a -> [(TyVar, Type)] -> Eval a)
-> [(TyVar, Type)]
-> Eval a
-> Eval a
forall a b. (a -> b) -> a -> b
$ ((TyVar, Type) -> Eval a -> Eval a)
-> Eval a -> [(TyVar, Type)] -> Eval a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((TyVar -> Type -> Eval a -> Eval a)
-> (TyVar, Type) -> Eval a -> Eval a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TyVar -> Type -> Eval a -> Eval a
forall a. TyVar -> Type -> Eval a -> Eval a
withTyVar)

getTvSubst :: Eval Subst
getTvSubst :: Eval Subst
getTvSubst = do
  InScopeSet
inScope <- Eval InScopeSet
getInScope
  Map TyVar Type
tys <- LocalEnv -> Map TyVar Type
lenvTypes (LocalEnv -> Map TyVar Type)
-> Eval LocalEnv -> Eval (Map TyVar Type)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval LocalEnv
getLocalEnv
  let vars :: VarEnv Type
vars = [(TyVar, Type)] -> VarEnv Type
forall a b. [(Var a, b)] -> VarEnv b
mkVarEnv (Map TyVar Type -> [(TyVar, Type)]
forall k a. Map k a -> [(k, a)]
Map.toList Map TyVar Type
tys)

  Subst -> Eval Subst
forall a. a -> Eval a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (InScopeSet -> VarEnv Type -> Subst
mkTvSubst InScopeSet
inScope VarEnv Type
vars)

findId :: Id -> Eval (Maybe Value)
findId :: Id -> Eval (Maybe Value)
findId Id
i = Id -> Map Id Value -> Maybe Value
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Id
i (Map Id Value -> Maybe Value)
-> (LocalEnv -> Map Id Value) -> LocalEnv -> Maybe Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LocalEnv -> Map Id Value
lenvValues (LocalEnv -> Maybe Value) -> Eval LocalEnv -> Eval (Maybe Value)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval LocalEnv
getLocalEnv

withId :: Id -> Value -> Eval a -> Eval a
withId :: forall a. Id -> Value -> Eval a -> Eval a
withId Id
i Value
v Eval a
x = do
  (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv GlobalEnv -> GlobalEnv
goGlobal
  (LocalEnv -> LocalEnv) -> Eval a -> Eval a
forall a. (LocalEnv -> LocalEnv) -> Eval a -> Eval a
modifyLocalEnv LocalEnv -> LocalEnv
goLocal Eval a
x
 where
  goGlobal :: GlobalEnv -> GlobalEnv
goGlobal env :: GlobalEnv
env@GlobalEnv{genvInScope :: GlobalEnv -> InScopeSet
genvInScope=InScopeSet
inScope} =
    -- TODO Change this to use an instance HasFreeVars Value
    let fvs :: VarSet
fvs = Id -> VarSet
forall a. Var a -> VarSet
unitVarSet Id
i VarSet -> VarSet -> VarSet
`unionVarSet` Term -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf (Value -> Term
forall a. AsTerm a => a -> Term
asTerm Value
v)
        iss :: InScopeSet
iss = VarSet -> InScopeSet
mkInScopeSet VarSet
fvs InScopeSet -> InScopeSet -> InScopeSet
`unionInScope` InScopeSet
inScope
     in GlobalEnv
env { genvInScope = iss }

  goLocal :: LocalEnv -> LocalEnv
goLocal env :: LocalEnv
env@LocalEnv{lenvValues :: LocalEnv -> Map Id Value
lenvValues=Map Id Value
values} =
    LocalEnv
env { lenvValues = Map.insert i v values }

withIds :: [(Id, Value)] -> Eval a -> Eval a
withIds :: forall a. [(Id, Value)] -> Eval a -> Eval a
withIds = (Eval a -> [(Id, Value)] -> Eval a)
-> [(Id, Value)] -> Eval a -> Eval a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Eval a -> [(Id, Value)] -> Eval a)
 -> [(Id, Value)] -> Eval a -> Eval a)
-> (Eval a -> [(Id, Value)] -> Eval a)
-> [(Id, Value)]
-> Eval a
-> Eval a
forall a b. (a -> b) -> a -> b
$ ((Id, Value) -> Eval a -> Eval a)
-> Eval a -> [(Id, Value)] -> Eval a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Id -> Value -> Eval a -> Eval a)
-> (Id, Value) -> Eval a -> Eval a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Id -> Value -> Eval a -> Eval a
forall a. Id -> Value -> Eval a -> Eval a
withId)

withoutId :: Id -> Eval a -> Eval a
withoutId :: forall a. Id -> Eval a -> Eval a
withoutId Id
i = (LocalEnv -> LocalEnv) -> Eval a -> Eval a
forall a. (LocalEnv -> LocalEnv) -> Eval a -> Eval a
modifyLocalEnv LocalEnv -> LocalEnv
go
 where
  go :: LocalEnv -> LocalEnv
go env :: LocalEnv
env@LocalEnv{lenvValues :: LocalEnv -> Map Id Value
lenvValues=Map Id Value
values} =
    LocalEnv
env { lenvValues = Map.delete i values }

findBinding :: Id -> Eval (Maybe (Binding Value))
findBinding :: Id -> Eval (Maybe (Binding Value))
findBinding Id
i = Id -> VarEnv (Binding Value) -> Maybe (Binding Value)
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Id
i (VarEnv (Binding Value) -> Maybe (Binding Value))
-> (GlobalEnv -> VarEnv (Binding Value))
-> GlobalEnv
-> Maybe (Binding Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GlobalEnv -> VarEnv (Binding Value)
genvBindings (GlobalEnv -> Maybe (Binding Value))
-> Eval GlobalEnv -> Eval (Maybe (Binding Value))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval GlobalEnv
getGlobalEnv

replaceBinding :: Binding Value -> Eval ()
replaceBinding :: Binding Value -> Eval ()
replaceBinding Binding Value
b = (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv GlobalEnv -> GlobalEnv
go
 where
  go :: GlobalEnv -> GlobalEnv
go env :: GlobalEnv
env@GlobalEnv{genvBindings :: GlobalEnv -> VarEnv (Binding Value)
genvBindings=VarEnv (Binding Value)
bindings} =
    GlobalEnv
env { genvBindings = extendVarEnv (bindingId b) b bindings }

getRef :: Int -> Eval Value
getRef :: Int -> Eval Value
getRef Int
addr = do
  IntMap Value
heap <- GlobalEnv -> IntMap Value
genvHeap (GlobalEnv -> IntMap Value)
-> Eval GlobalEnv -> Eval (IntMap Value)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval GlobalEnv
getGlobalEnv

  case Int -> IntMap Value -> Maybe Value
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
addr IntMap Value
heap of
    Just Value
val -> Value -> Eval Value
forall a. a -> Eval a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Value
val
    Maybe Value
Nothing  -> [Char] -> Eval Value
forall a. HasCallStack => [Char] -> a
error ([Char]
"getHeap: Address " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
addr [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" out of bounds")

setRef :: Int -> Value -> Eval ()
setRef :: Int -> Value -> Eval ()
setRef Int
addr Value
val = (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv GlobalEnv -> GlobalEnv
go
 where
  go :: GlobalEnv -> GlobalEnv
go env :: GlobalEnv
env@GlobalEnv{genvHeap :: GlobalEnv -> IntMap Value
genvHeap=IntMap Value
heap,genvAddr :: GlobalEnv -> Int
genvAddr=Int
next}
    | Int
addr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
next =
        GlobalEnv
env { genvHeap = IntMap.insert addr val heap, genvAddr = addr + 1 }

    | Bool
otherwise =
        GlobalEnv
env { genvHeap = IntMap.insert addr val heap }

isKeepingLifted :: Eval Bool
isKeepingLifted :: Eval Bool
isKeepingLifted = LocalEnv -> Bool
lenvKeepLifted (LocalEnv -> Bool) -> Eval LocalEnv -> Eval Bool
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval LocalEnv
getLocalEnv

keepLifted :: Eval a -> Eval a
keepLifted :: forall a. Eval a -> Eval a
keepLifted = (LocalEnv -> LocalEnv) -> Eval a -> Eval a
forall a. (LocalEnv -> LocalEnv) -> Eval a -> Eval a
modifyLocalEnv LocalEnv -> LocalEnv
forceLifted
 where
  forceLifted :: LocalEnv -> LocalEnv
forceLifted LocalEnv
env = LocalEnv
env { lenvKeepLifted = True }

getFuel :: Eval Word
getFuel :: Eval Word
getFuel = do
  LocalEnv
lenv <- Eval LocalEnv
getLocalEnv
  GlobalEnv
genv <- Eval GlobalEnv
getGlobalEnv

  Word -> Eval Word
forall a. a -> Eval a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Word -> Word -> Word
forall a. Ord a => a -> a -> a
min (LocalEnv -> Word
lenvFuel LocalEnv
lenv) (GlobalEnv -> Word
genvFuel GlobalEnv
genv))

withFuel :: Eval a -> Eval a
withFuel :: forall a. Eval a -> Eval a
withFuel Eval a
x = (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv GlobalEnv -> GlobalEnv
go Eval () -> Eval a -> Eval a
forall a b. Eval a -> Eval b -> Eval b
forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
>> Eval a
x
 where
  go :: GlobalEnv -> GlobalEnv
go env :: GlobalEnv
env@GlobalEnv{genvFuel :: GlobalEnv -> Word
genvFuel=Word
fuel} =
    GlobalEnv
env { genvFuel = fuel - 1 }

preserveFuel :: Eval a -> Eval a
preserveFuel :: forall a. Eval a -> Eval a
preserveFuel Eval a
x = do
  Word
fuel <- Eval Word
getFuel
  a
res  <- Eval a
x

  (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv (Word -> GlobalEnv -> GlobalEnv
go Word
fuel)
  a -> Eval a
forall a. a -> Eval a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure a
res
 where
  go :: Word -> GlobalEnv -> GlobalEnv
go Word
fuel GlobalEnv
env = GlobalEnv
env { genvFuel = fuel }

getTyConMap :: Eval TyConMap
getTyConMap :: Eval TyConMap
getTyConMap = GlobalEnv -> TyConMap
genvTyConMap (GlobalEnv -> TyConMap) -> Eval GlobalEnv -> Eval TyConMap
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval GlobalEnv
getGlobalEnv

getInScope :: Eval InScopeSet
getInScope :: Eval InScopeSet
getInScope = GlobalEnv -> InScopeSet
genvInScope (GlobalEnv -> InScopeSet) -> Eval GlobalEnv -> Eval InScopeSet
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval GlobalEnv
getGlobalEnv

getUniqueId :: OccName -> Type -> Eval Id
getUniqueId :: OccName -> Type -> Eval Id
getUniqueId = ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Id))
-> OccName -> Type -> Eval Id
forall a.
((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Var a))
-> OccName -> Type -> Eval (Var a)
getUniqueVar (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Id)
mkUniqSystemId

getUniqueTyVar :: OccName -> Kind -> Eval TyVar
getUniqueTyVar :: OccName -> Type -> Eval TyVar
getUniqueTyVar = ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), TyVar))
-> OccName -> Type -> Eval TyVar
forall a.
((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Var a))
-> OccName -> Type -> Eval (Var a)
getUniqueVar (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), TyVar)
mkUniqSystemTyVar

getUniqueVar
  :: ((Supply, InScopeSet)
         -> (OccName, KindOrType)
         -> ((Supply, InScopeSet), Var a))
  -> OccName
  -> KindOrType
  -> Eval (Var a)
getUniqueVar :: forall a.
((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Var a))
-> OccName -> Type -> Eval (Var a)
getUniqueVar (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Var a)
f OccName
name Type
ty = do
  GlobalEnv
env <- Eval GlobalEnv
getGlobalEnv
  let iss :: InScopeSet
iss = GlobalEnv -> InScopeSet
genvInScope GlobalEnv
env
      ids :: Supply
ids = GlobalEnv -> Supply
genvSupply GlobalEnv
env
      ((Supply
ids', InScopeSet
iss'), Var a
i) = (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Var a)
f (Supply
ids, InScopeSet
iss) (OccName
name, Type
ty)

  (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv (Supply -> InScopeSet -> GlobalEnv -> GlobalEnv
go Supply
ids' InScopeSet
iss')
  Var a -> Eval (Var a)
forall a. a -> Eval a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Var a
i
 where
  go :: Supply -> InScopeSet -> GlobalEnv -> GlobalEnv
go Supply
ids InScopeSet
iss GlobalEnv
env =
    GlobalEnv
env { genvInScope = iss, genvSupply = ids }

workFreeValue :: Value -> Eval Bool
workFreeValue :: Value -> Eval Bool
workFreeValue = \case
  VNeutral Neutral Value
_ -> Bool -> Eval Bool
forall a. a -> Eval a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Bool
False
  VThunk Term
x LocalEnv
_ -> do
    UniqMap (Binding Term)
bindings <- (Binding Value -> Binding Term)
-> VarEnv (Binding Value) -> UniqMap (Binding Term)
forall a b. (a -> b) -> UniqMap a -> UniqMap b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Value -> Term) -> Binding Value -> Binding Term
forall a b. (a -> b) -> Binding a -> Binding b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Term
forall a. AsTerm a => a -> Term
asTerm) (VarEnv (Binding Value) -> UniqMap (Binding Term))
-> (GlobalEnv -> VarEnv (Binding Value))
-> GlobalEnv
-> UniqMap (Binding Term)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GlobalEnv -> VarEnv (Binding Value)
genvBindings (GlobalEnv -> UniqMap (Binding Term))
-> Eval GlobalEnv -> Eval (UniqMap (Binding Term))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eval GlobalEnv
getGlobalEnv
    Lens' GlobalEnv (VarEnv Bool)
-> UniqMap (Binding Term) -> Term -> Eval Bool
forall s (m :: Type -> Type).
(HasCallStack, MonadState s m) =>
Lens' s (VarEnv Bool) -> UniqMap (Binding Term) -> Term -> m Bool
isWorkFree (VarEnv Bool -> f (VarEnv Bool)) -> GlobalEnv -> f GlobalEnv
Lens' GlobalEnv (VarEnv Bool)
workFreeCache UniqMap (Binding Term)
bindings Term
x

  Value
_ -> Bool -> Eval Bool
forall a. a -> Eval a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Bool
True