{-|
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
  { Eval a -> RWST LocalEnv () GlobalEnv IO a
unEval :: RWST LocalEnv () GlobalEnv IO a }
  deriving
    ( a -> Eval b -> Eval a
(a -> b) -> Eval a -> Eval b
(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
<$ :: a -> Eval b -> Eval a
$c<$ :: forall a b. a -> Eval b -> Eval a
fmap :: (a -> b) -> Eval a -> Eval b
$cfmap :: forall a b. (a -> b) -> Eval a -> Eval b
Functor
    , Functor Eval
a -> Eval a
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
Eval a -> Eval b -> Eval b
Eval a -> Eval b -> Eval a
Eval (a -> b) -> Eval a -> Eval b
(a -> b -> c) -> Eval a -> Eval b -> Eval c
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
<* :: Eval a -> Eval b -> Eval a
$c<* :: forall a b. Eval a -> Eval b -> Eval a
*> :: Eval a -> Eval b -> Eval b
$c*> :: forall a b. Eval a -> Eval b -> Eval b
liftA2 :: (a -> b -> c) -> Eval a -> Eval b -> Eval c
$cliftA2 :: forall a b c. (a -> b -> c) -> Eval a -> Eval b -> Eval c
<*> :: Eval (a -> b) -> Eval a -> Eval b
$c<*> :: forall a b. Eval (a -> b) -> Eval a -> Eval b
pure :: a -> Eval a
$cpure :: forall a. a -> Eval a
$cp1Applicative :: Functor Eval
Applicative
    , Applicative Eval
Eval a
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
Eval a -> Eval a -> Eval a
Eval a -> Eval [a]
Eval a -> Eval [a]
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
many :: Eval a -> Eval [a]
$cmany :: forall a. Eval a -> Eval [a]
some :: Eval a -> Eval [a]
$csome :: forall a. Eval a -> Eval [a]
<|> :: Eval a -> Eval a -> Eval a
$c<|> :: forall a. Eval a -> Eval a -> Eval a
empty :: Eval a
$cempty :: forall a. Eval a
$cp1Alternative :: Applicative Eval
Alternative
    , Applicative Eval
a -> Eval a
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
Eval a -> (a -> Eval b) -> Eval b
Eval a -> Eval b -> Eval b
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
return :: a -> Eval a
$creturn :: forall a. a -> Eval a
>> :: Eval a -> Eval b -> Eval b
$c>> :: forall a b. Eval a -> Eval b -> Eval b
>>= :: Eval a -> (a -> Eval b) -> Eval b
$c>>= :: forall a b. Eval a -> (a -> Eval b) -> Eval b
$cp1Monad :: Applicative Eval
Monad
    , Monad Eval
Monad Eval -> (forall a. String -> Eval a) -> MonadFail Eval
String -> Eval a
forall a. String -> Eval a
forall (m :: Type -> Type).
Monad m -> (forall a. String -> m a) -> MonadFail m
fail :: String -> Eval a
$cfail :: forall a. String -> Eval a
$cp1MonadFail :: Monad Eval
MonadFail
    , Monad Eval
Monad Eval -> (forall a. IO a -> Eval a) -> MonadIO Eval
IO a -> Eval a
forall a. IO a -> Eval a
forall (m :: Type -> Type).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
liftIO :: IO a -> Eval a
$cliftIO :: forall a. IO a -> Eval a
$cp1MonadIO :: Monad Eval
MonadIO
    , MonadReader LocalEnv
    , MonadState GlobalEnv
    , Monad Eval
e -> Eval a
Monad Eval
-> (forall e a. Exception e => e -> Eval a) -> MonadThrow Eval
forall e a. Exception e => e -> Eval a
forall (m :: Type -> Type).
Monad m -> (forall e a. Exception e => e -> m a) -> MonadThrow m
throwM :: e -> Eval a
$cthrowM :: forall e a. Exception e => e -> Eval a
$cp1MonadThrow :: Monad Eval
MonadThrow
    , MonadThrow Eval
MonadThrow Eval
-> (forall e a. Exception e => Eval a -> (e -> Eval a) -> Eval a)
-> MonadCatch Eval
Eval a -> (e -> Eval a) -> Eval a
forall e a. Exception e => Eval a -> (e -> Eval a) -> Eval a
forall (m :: Type -> Type).
MonadThrow m
-> (forall e a. Exception e => m a -> (e -> m a) -> m a)
-> MonadCatch m
catch :: Eval a -> (e -> Eval a) -> Eval a
$ccatch :: forall e a. Exception e => Eval a -> (e -> Eval a) -> Eval a
$cp1MonadCatch :: MonadThrow Eval
MonadCatch
    , MonadCatch Eval
MonadCatch Eval
-> (forall b. ((forall a. Eval a -> Eval a) -> Eval b) -> Eval b)
-> (forall b. ((forall a. Eval a -> Eval a) -> Eval b) -> Eval b)
-> (forall a b c.
    Eval a
    -> (a -> ExitCase b -> Eval c) -> (a -> Eval b) -> Eval (b, c))
-> MonadMask Eval
Eval a
-> (a -> ExitCase b -> Eval c) -> (a -> Eval b) -> Eval (b, c)
((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
forall b. ((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
forall a b c.
Eval a
-> (a -> ExitCase b -> Eval c) -> (a -> Eval b) -> Eval (b, c)
forall (m :: Type -> Type).
MonadCatch m
-> (forall b. ((forall a. m a -> m a) -> m b) -> m b)
-> (forall b. ((forall a. m a -> m a) -> m b) -> m b)
-> (forall a b c.
    m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c))
-> MonadMask m
generalBracket :: Eval a
-> (a -> ExitCase b -> Eval c) -> (a -> Eval b) -> Eval (b, c)
$cgeneralBracket :: forall a b c.
Eval a
-> (a -> ExitCase b -> Eval c) -> (a -> Eval b) -> Eval (b, c)
uninterruptibleMask :: ((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
$cuninterruptibleMask :: forall b. ((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
mask :: ((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
$cmask :: forall b. ((forall a. Eval a -> Eval a) -> Eval b) -> Eval b
$cp1MonadMask :: MonadCatch Eval
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 :: 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 :: LocalEnv -> Eval a -> Eval a
setLocalEnv = (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 :: (LocalEnv -> LocalEnv) -> Eval a -> Eval a
modifyLocalEnv = (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 :: 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 :: Id
lenvContext = Id
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 :: 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 :: InScopeSet
genvInScope = InScopeSet
iss }

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

withTyVars :: [(TyVar, Type)] -> Eval a -> Eval a
withTyVars :: [(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 (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 (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 :: 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 :: InScopeSet
genvInScope = InScopeSet
iss }

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

withIds :: [(Id, Value)] -> Eval a -> Eval a
withIds :: [(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 (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 :: 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 Id Value
lenvValues = Id -> Map Id Value -> Map Id Value
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Id
i Map Id Value
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 :: VarEnv (Binding Value)
genvBindings = Id
-> Binding Value
-> VarEnv (Binding Value)
-> VarEnv (Binding Value)
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv (Binding Value -> Id
forall a. Binding a -> Id
bindingId Binding Value
b) Binding Value
b VarEnv (Binding Value)
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 (f :: Type -> Type) a. Applicative f => a -> f a
pure Value
val
    Maybe Value
Nothing  -> String -> Eval Value
forall a. HasCallStack => String -> a
error (String
"getHeap: Address " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
addr String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" 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 Value
genvHeap = Int -> Value -> IntMap Value -> IntMap Value
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
addr Value
val IntMap Value
heap, genvAddr :: Int
genvAddr = Int
addr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 }

    | Bool
otherwise =
        GlobalEnv
env { genvHeap :: IntMap Value
genvHeap = Int -> Value -> IntMap Value -> IntMap Value
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
addr Value
val IntMap Value
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 :: 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 :: Bool
lenvKeepLifted = Bool
True }

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

  Word -> Eval Word
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 :: Eval a -> Eval a
withFuel Eval a
x = (GlobalEnv -> GlobalEnv) -> Eval ()
modifyGlobalEnv GlobalEnv -> GlobalEnv
go Eval () -> Eval a -> Eval a
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 :: Word
genvFuel = Word
fuel Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1 }

preserveFuel :: Eval a -> Eval a
preserveFuel :: 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 (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 :: Word
genvFuel = Word
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 :: ((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 (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 :: InScopeSet
genvInScope = InScopeSet
iss, genvSupply :: Supply
genvSupply = Supply
ids }

workFreeValue :: Value -> Eval Bool
workFreeValue :: Value -> Eval Bool
workFreeValue = \case
  VNeutral Neutral Value
_ -> Bool -> Eval Bool
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 (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Value -> Term) -> Binding Value -> Binding Term
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 Lens' GlobalEnv (VarEnv Bool)
workFreeCache UniqMap (Binding Term)
bindings Term
x

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