{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE CPP #-}
{-# OPTIONS_HADDOCK not-home #-}
module Effectful.Internal.Env
  ( -- * The environment
    Env(..)
  , Storage(..)

    -- ** Relinker
  , Relinker(..)
  , dummyRelinker

    -- * Dispatch
  , Dispatch(..)
  , SideEffects(..)
  , DispatchOf
  , EffectRep

    -- * Operations
  , emptyEnv
  , cloneEnv
  , restoreEnv
  , sizeEnv
  , tailEnv

    -- ** Modification of the effect stack
  , consEnv
  , unconsEnv
  , replaceEnv
  , unreplaceEnv
  , subsumeEnv
  , injectEnv

    -- ** Data retrieval and update
  , getEnv
  , putEnv
  , stateEnv
  , modifyEnv
  ) where

import Control.Monad
import Control.Monad.Primitive
import Data.Primitive.PrimArray
import Data.Primitive.SmallArray
import GHC.Stack (HasCallStack)

import Effectful.Internal.Effect
import Effectful.Internal.Utils

type role Env nominal

-- | A strict (WHNF), __thread local__, mutable, extensible record indexed by types
-- of kind 'Effect'.
--
-- __Warning: the environment is a mutable data structure and cannot be simultaneously used from multiple threads under any circumstances.__
--
-- In order to pass it to a different thread, you need to perform a deep copy
-- with the 'cloneEnv' funtion.
--
-- Offers very good performance characteristics for most often performed
-- operations:
--
-- - Extending: /@O(n)@/, where @n@ is the size of the effect stack.
--
-- - Shrinking: /@O(1)@/.
--
-- - Indexing via '(:>)': /@O(1)@/
--
-- - Modification of a specific element: /@O(1)@/.
--
-- - Getting a tail: /@O(1)@/.
--
-- - Cloning: /@O(N)@/, where @N@ is the size of the 'Storage'.
--
data Env (es :: [Effect]) = Env
  { forall (es :: [Effect]). Env es -> Int
envOffset  :: !Int
  , forall (es :: [Effect]). Env es -> PrimArray Int
envRefs    :: !(PrimArray Int)
  , forall (es :: [Effect]). Env es -> IORef' Storage
envStorage :: !(IORef' Storage)
  }

-- | A storage of effects.
data Storage = Storage
  { Storage -> Int
stSize      :: !Int
  , Storage -> Int
stVersion   :: !Int
  , Storage -> MutablePrimArray RealWorld Int
stVersions  :: !(MutablePrimArray RealWorld Int)
  , Storage -> SmallMutableArray RealWorld Any
stEffects   :: !(SmallMutableArray RealWorld Any)
  , Storage -> SmallMutableArray RealWorld Any
stRelinkers :: !(SmallMutableArray RealWorld Any)
  }

----------------------------------------
-- Relinker

-- | A function for relinking 'Env' objects stored in the handlers and/or making
-- a deep copy of the representation of the effect when cloning the environment.
newtype Relinker :: (Effect -> Type) -> Effect -> Type where
  Relinker
    :: ((forall es. Env es -> IO (Env es)) -> rep e -> IO (rep e))
    -> Relinker rep e

-- | A dummy 'Relinker'.
dummyRelinker :: Relinker rep e
dummyRelinker :: forall (rep :: Effect -> Type) (e :: Effect). Relinker rep e
dummyRelinker = forall (rep :: Effect -> Type) (e :: Effect).
((forall (es :: [Effect]). Env es -> IO (Env es))
 -> rep e -> IO (rep e))
-> Relinker rep e
Relinker forall a b. (a -> b) -> a -> b
$ \forall (es :: [Effect]). Env es -> IO (Env es)
_ -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure

----------------------------------------
-- Dispatch

-- | A type of dispatch. For more information consult the documentation in
-- "Effectful.Dispatch.Dynamic" and "Effectful.Dispatch.Static".
data Dispatch = Dynamic | Static SideEffects

-- | Signifies whether core operations of a statically dispatched effect perform
-- side effects. If an effect is marked as such, the
-- 'Effectful.Dispatch.Static.runStaticRep' family of functions will require the
-- 'Effectful.IOE' effect to be in context via the
-- 'Effectful.Dispatch.Static.MaybeIOE' type family.
data SideEffects = NoSideEffects | WithSideEffects

-- | Dispatch types of effects.
type family DispatchOf (e :: Effect) :: Dispatch

-- | Internal representations of effects.
type family EffectRep (d :: Dispatch) :: Effect -> Type

----------------------------------------
-- Operations

-- | Create an empty environment.
emptyEnv :: IO (Env '[])
emptyEnv :: IO (Env '[])
emptyEnv = forall (es :: [Effect]).
Int -> PrimArray Int -> IORef' Storage -> Env es
Env Int
0
  forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (m :: Type -> Type) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
0)
  forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall a. a -> IO (IORef' a)
newIORef' forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO Storage
emptyStorage)

-- | Clone the environment to use it in a different thread.
cloneEnv :: Env es -> IO (Env es)
cloneEnv :: forall (es :: [Effect]). Env es -> IO (Env es)
cloneEnv (Env Int
offset PrimArray Int
refs IORef' Storage
storage0) = do
  Storage Int
storageSize Int
version MutablePrimArray RealWorld Int
vs0 SmallMutableArray RealWorld Any
es0 SmallMutableArray RealWorld Any
fs0 <- forall a. IORef' a -> IO a
readIORef' IORef' Storage
storage0
  Int
vsSize <- forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray  MutablePrimArray RealWorld Int
vs0
  Int
esSize <- forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m Int
getSizeofSmallMutableArray SmallMutableArray RealWorld Any
es0
  Int
fsSize <- forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m Int
getSizeofSmallMutableArray SmallMutableArray RealWorld Any
fs0
  forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
when (Int
vsSize forall a. Eq a => a -> a -> Bool
/= Int
esSize) forall a b. (a -> b) -> a -> b
$ do
    forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"vsSize (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
vsSize forall a. [a] -> [a] -> [a]
++ [Char]
") /= esSize (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
esSize forall a. [a] -> [a] -> [a]
++ [Char]
")"
  forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
when (Int
esSize forall a. Eq a => a -> a -> Bool
/= Int
fsSize) forall a b. (a -> b) -> a -> b
$ do
    forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"esSize (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
esSize forall a. [a] -> [a] -> [a]
++ [Char]
") /= fsSize (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
fsSize forall a. [a] -> [a] -> [a]
++ [Char]
")"
  MutablePrimArray RealWorld Int
vs <- forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArray  MutablePrimArray RealWorld Int
vs0 Int
0 Int
vsSize
  SmallMutableArray RealWorld Any
es <- forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> Int -> m (SmallMutableArray (PrimState m) a)
cloneSmallMutableArray SmallMutableArray RealWorld Any
es0 Int
0 Int
esSize
  SmallMutableArray RealWorld Any
fs <- forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> Int -> m (SmallMutableArray (PrimState m) a)
cloneSmallMutableArray SmallMutableArray RealWorld Any
fs0 Int
0 Int
fsSize
  IORef' Storage
storage <- forall a. a -> IO (IORef' a)
newIORef' forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> MutablePrimArray RealWorld Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage Int
storageSize Int
version MutablePrimArray RealWorld Int
vs SmallMutableArray RealWorld Any
es SmallMutableArray RealWorld Any
fs
  let relinkEffects :: Int -> IO ()
relinkEffects = \case
        Int
0 -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()
        Int
k -> do
          let i :: Int
i = Int
k forall a. Num a => a -> a -> a
- Int
1
          Relinker (forall (es :: [Effect]). Env es -> IO (Env es))
-> Any Any -> IO (Any Any)
f <- forall a. Any -> a
fromAny forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
fs Int
i
          forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
es Int
i
            forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= (forall (es :: [Effect]). Env es -> IO (Env es))
-> Any Any -> IO (Any Any)
f (forall (es :: [Effect]). IORef' Storage -> Env es -> IO (Env es)
relinkEnv IORef' Storage
storage) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Any -> a
fromAny
            forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a. SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' SmallMutableArray RealWorld Any
es Int
i forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Any
toAny
          Int -> IO ()
relinkEffects Int
i
  Int -> IO ()
relinkEffects Int
storageSize
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
Int -> PrimArray Int -> IORef' Storage -> Env es
Env Int
offset PrimArray Int
refs IORef' Storage
storage
{-# NOINLINE cloneEnv #-}

-- | Restore the environment from its clone.
--
-- @since 2.2.0.0
restoreEnv
  :: Env es -- ^ Destination.
  -> Env es -- ^ Source.
  -> IO ()
restoreEnv :: forall (es :: [Effect]). Env es -> Env es -> IO ()
restoreEnv Env es
dest Env es
src = do
  Storage
destStorage <- forall a. IORef' a -> IO a
readIORef' (forall (es :: [Effect]). Env es -> IORef' Storage
envStorage Env es
dest)
  Storage
srcStorage  <- forall a. IORef' a -> IO a
readIORef' (forall (es :: [Effect]). Env es -> IORef' Storage
envStorage Env es
src)
  let destStorageSize :: Int
destStorageSize = Storage -> Int
stSize Storage
destStorage
      srcStorageSize :: Int
srcStorageSize  = Storage -> Int
stSize Storage
srcStorage
  forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
when (Int
destStorageSize forall a. Eq a => a -> a -> Bool
/= Int
srcStorageSize) forall a b. (a -> b) -> a -> b
$ do
    forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"destStorageSize (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
destStorageSize
         forall a. [a] -> [a] -> [a]
++ [Char]
") /= srcStorageSize (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
srcStorageSize forall a. [a] -> [a] -> [a]
++ [Char]
")"
  forall a. IORef' a -> a -> IO ()
writeIORef' (forall (es :: [Effect]). Env es -> IORef' Storage
envStorage Env es
dest) forall a b. (a -> b) -> a -> b
$ Storage
srcStorage
    -- Decreasing the counter allows leakage of unsafeCoerce (see unsafeCoerce2
    -- in the EnvTests module).
    { stVersion :: Int
stVersion = forall a. Ord a => a -> a -> a
max (Storage -> Int
stVersion Storage
destStorage) (Storage -> Int
stVersion Storage
srcStorage)
    }
{-# NOINLINE restoreEnv #-}

-- | Get the current size of the environment.
sizeEnv :: Env es -> IO Int
sizeEnv :: forall (es :: [Effect]). Env es -> IO Int
sizeEnv (Env Int
offset PrimArray Int
refs IORef' Storage
_) = do
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
refs forall a. Num a => a -> a -> a
- Int
offset) forall a. Integral a => a -> a -> a
`div` Int
2

-- | Access the tail of the environment.
tailEnv :: Env (e : es) -> IO (Env es)
tailEnv :: forall (e :: Effect) (es :: [Effect]). Env (e : es) -> IO (Env es)
tailEnv (Env Int
offset PrimArray Int
refs IORef' Storage
storage) = do
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
Int -> PrimArray Int -> IORef' Storage -> Env es
Env (Int
offset forall a. Num a => a -> a -> a
+ Int
2) PrimArray Int
refs IORef' Storage
storage

----------------------------------------
-- Extending and shrinking

-- | Extend the environment with a new data type.
consEnv
  :: EffectRep (DispatchOf e) e
  -- ^ The representation of the effect.
  -> Relinker (EffectRep (DispatchOf e)) e
  -> Env es
  -> IO (Env (e : es))
consEnv :: forall (e :: Effect) (es :: [Effect]).
EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e
-> Env es
-> IO (Env (e : es))
consEnv EffectRep (DispatchOf e) e
e Relinker (EffectRep (DispatchOf e)) e
f (Env Int
offset PrimArray Int
refs0 IORef' Storage
storage) = do
  let size :: Int
size = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
refs0 forall a. Num a => a -> a -> a
- Int
offset
  MutablePrimArray RealWorld Int
mrefs <- forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (Int
size forall a. Num a => a -> a -> a
+ Int
2)
  forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray RealWorld Int
mrefs Int
2 PrimArray Int
refs0 Int
offset Int
size
  (Int
ref, Int
version) <- forall (e :: Effect).
IORef' Storage
-> EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e
-> IO (Int, Int)
insertEffect IORef' Storage
storage EffectRep (DispatchOf e) e
e Relinker (EffectRep (DispatchOf e)) e
f
  forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
mrefs Int
0 Int
ref
  forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
mrefs Int
1 Int
version
  PrimArray Int
refs <- forall (m :: Type -> Type) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld Int
mrefs
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
Int -> PrimArray Int -> IORef' Storage -> Env es
Env Int
0 PrimArray Int
refs IORef' Storage
storage
{-# NOINLINE consEnv #-}

-- | Shrink the environment by one data type.
--
-- /Note:/ after calling this function @e@ from the input environment is no
-- longer usable.
unconsEnv :: Env (e : es) -> IO ()
unconsEnv :: forall (e :: Effect) (es :: [Effect]). Env (e : es) -> IO ()
unconsEnv (Env Int
_ PrimArray Int
refs IORef' Storage
storage) = do
  IORef' Storage -> Int -> IO ()
deleteEffect IORef' Storage
storage (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
refs Int
0)
{-# NOINLINE unconsEnv #-}

----------------------------------------

-- | Replace a specific effect in the stack with a new value.
--
-- /Note:/ unlike in 'putEnv' the value in not changed in place, so only the new
-- environment will see it.
replaceEnv
  :: forall e es. e :> es
  => EffectRep (DispatchOf e) e
  -- ^ The representation of the effect.
  -> Relinker (EffectRep (DispatchOf e)) e
  -> Env es
  -> IO (Env es)
replaceEnv :: forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e -> Env es -> IO (Env es)
replaceEnv EffectRep (DispatchOf e) e
e Relinker (EffectRep (DispatchOf e)) e
f (Env Int
offset PrimArray Int
refs0 IORef' Storage
storage) = do
  let size :: Int
size = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
refs0 forall a. Num a => a -> a -> a
- Int
offset
  MutablePrimArray RealWorld Int
mrefs <- forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
size
  forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray RealWorld Int
mrefs Int
0 PrimArray Int
refs0 Int
offset Int
size
  (Int
ref, Int
version) <- forall (e :: Effect).
IORef' Storage
-> EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e
-> IO (Int, Int)
insertEffect IORef' Storage
storage EffectRep (DispatchOf e) e
e Relinker (EffectRep (DispatchOf e)) e
f
  let i :: Int
i = Int
2 forall a. Num a => a -> a -> a
* forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es
  forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
mrefs  Int
i      Int
ref
  forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
mrefs (Int
i forall a. Num a => a -> a -> a
+ Int
1) Int
version
  PrimArray Int
refs <- forall (m :: Type -> Type) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld Int
mrefs
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
Int -> PrimArray Int -> IORef' Storage -> Env es
Env Int
0 PrimArray Int
refs IORef' Storage
storage
{-# NOINLINE replaceEnv #-}

-- | Remove a reference to the replaced effect.
--
-- /Note:/ after calling this function the input environment is no longer
-- usable.
unreplaceEnv :: forall e es. e :> es => Env es -> IO ()
unreplaceEnv :: forall (e :: Effect) (es :: [Effect]). (e :> es) => Env es -> IO ()
unreplaceEnv (Env Int
offset PrimArray Int
refs IORef' Storage
storage) = do
  IORef' Storage -> Int -> IO ()
deleteEffect IORef' Storage
storage forall a b. (a -> b) -> a -> b
$ forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
refs (Int
offset forall a. Num a => a -> a -> a
+ Int
2 forall a. Num a => a -> a -> a
* forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es)
{-# NOINLINE unreplaceEnv #-}

----------------------------------------

-- | Reference an existing effect from the top of the stack.
subsumeEnv :: forall e es. e :> es => Env es -> IO (Env (e : es))
subsumeEnv :: forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es -> IO (Env (e : es))
subsumeEnv (Env Int
offset PrimArray Int
refs0 IORef' Storage
storage) = do
  let size :: Int
size = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
refs0 forall a. Num a => a -> a -> a
- Int
offset
  MutablePrimArray RealWorld Int
mrefs <- forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (Int
size forall a. Num a => a -> a -> a
+ Int
2)
  forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray RealWorld Int
mrefs Int
2 PrimArray Int
refs0 Int
offset Int
size
  let ix :: Int
ix = Int
offset forall a. Num a => a -> a -> a
+ Int
2 forall a. Num a => a -> a -> a
* forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es
  forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
mrefs Int
0 forall a b. (a -> b) -> a -> b
$ forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
refs0  Int
ix
  forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
mrefs Int
1 forall a b. (a -> b) -> a -> b
$ forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
refs0 (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
  PrimArray Int
refs <- forall (m :: Type -> Type) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld Int
mrefs
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
Int -> PrimArray Int -> IORef' Storage -> Env es
Env Int
0 PrimArray Int
refs IORef' Storage
storage
{-# NOINLINE subsumeEnv #-}

----------------------------------------

-- | Construct an environment containing a permutation (with possible
-- duplicates) of a subset of effects from the input environment.
injectEnv :: forall xs es. Subset xs es => Env es -> IO (Env xs)
injectEnv :: forall (xs :: [Effect]) (es :: [Effect]).
Subset xs es =>
Env es -> IO (Env xs)
injectEnv (Env Int
offset PrimArray Int
refs0 IORef' Storage
storage) = do
  let xs :: [Int]
xs         = forall (xs :: [Effect]) (es :: [Effect]). Subset xs es => [Int]
reifyIndices @xs @es
      permSize :: Int
permSize   = Int
2 forall a. Num a => a -> a -> a
* forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Int]
xs
      prefixSize :: Int
prefixSize = Int
2 forall a. Num a => a -> a -> a
* forall (es :: [Effect]). KnownPrefix es => Int
prefixLength @es
      suffixSize :: Int
suffixSize = if forall (xs :: [Effect]) (es :: [Effect]). Subset xs es => Bool
subsetFullyKnown @xs @es
                   then Int
0
                   else forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
refs0 forall a. Num a => a -> a -> a
- Int
offset forall a. Num a => a -> a -> a
- Int
prefixSize
  MutablePrimArray RealWorld Int
mrefs <- forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (Int
permSize forall a. Num a => a -> a -> a
+ Int
suffixSize)
  forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray RealWorld Int
mrefs Int
permSize PrimArray Int
refs0 (Int
offset forall a. Num a => a -> a -> a
+ Int
prefixSize) Int
suffixSize
  let writePermRefs :: Int -> [Int] -> IO ()
writePermRefs Int
i = \case
        []       -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()
        (Int
e : [Int]
es) -> do
          let ix :: Int
ix = Int
offset forall a. Num a => a -> a -> a
+ Int
2 forall a. Num a => a -> a -> a
* Int
e
          forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
mrefs  Int
i      forall a b. (a -> b) -> a -> b
$ forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
refs0  Int
ix
          forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
mrefs (Int
i forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
refs0 (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
          Int -> [Int] -> IO ()
writePermRefs (Int
i forall a. Num a => a -> a -> a
+ Int
2) [Int]
es
  Int -> [Int] -> IO ()
writePermRefs Int
0 [Int]
xs
  PrimArray Int
refs <- forall (m :: Type -> Type) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld Int
mrefs
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
Int -> PrimArray Int -> IORef' Storage -> Env es
Env Int
0 PrimArray Int
refs IORef' Storage
storage
{-# NOINLINE injectEnv #-}

----------------------------------------
-- Data retrieval and update

-- | Extract a specific data type from the environment.
getEnv
  :: forall e es. e :> es
  => Env es -- ^ The environment.
  -> IO (EffectRep (DispatchOf e) e)
getEnv :: forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es -> IO (EffectRep (DispatchOf e) e)
getEnv Env es
env = do
  (Int
i, SmallMutableArray RealWorld Any
es) <- forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es -> IO (Int, SmallMutableArray RealWorld Any)
getLocation @e Env es
env
  forall a. Any -> a
fromAny forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
es Int
i

-- | Replace the data type in the environment with a new value (in place).
putEnv
  :: forall e es. e :> es
  => Env es -- ^ The environment.
  -> EffectRep (DispatchOf e) e
  -> IO ()
putEnv :: forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es -> EffectRep (DispatchOf e) e -> IO ()
putEnv Env es
env EffectRep (DispatchOf e) e
e = do
  (Int
i, SmallMutableArray RealWorld Any
es) <- forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es -> IO (Int, SmallMutableArray RealWorld Any)
getLocation @e Env es
env
  forall a. SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' SmallMutableArray RealWorld Any
es Int
i (forall a. a -> Any
toAny EffectRep (DispatchOf e) e
e)

-- | Modify the data type in the environment and return a value (in place).
stateEnv
  :: forall e es a. e :> es
  => Env es -- ^ The environment.
  -> (EffectRep (DispatchOf e) e -> IO (a, EffectRep (DispatchOf e) e))
  -> IO a
stateEnv :: forall (e :: Effect) (es :: [Effect]) a.
(e :> es) =>
Env es
-> (EffectRep (DispatchOf e) e
    -> IO (a, EffectRep (DispatchOf e) e))
-> IO a
stateEnv Env es
env EffectRep (DispatchOf e) e -> IO (a, EffectRep (DispatchOf e) e)
f = do
  (Int
i, SmallMutableArray RealWorld Any
es) <- forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es -> IO (Int, SmallMutableArray RealWorld Any)
getLocation @e Env es
env
  (a
a, EffectRep (DispatchOf e) e
e) <- EffectRep (DispatchOf e) e -> IO (a, EffectRep (DispatchOf e) e)
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Any -> a
fromAny forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
es Int
i
  forall a. SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' SmallMutableArray RealWorld Any
es Int
i (forall a. a -> Any
toAny EffectRep (DispatchOf e) e
e)
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure a
a

-- | Modify the data type in the environment (in place).
modifyEnv
  :: forall e es. e :> es
  => Env es -- ^ The environment.
  -> (EffectRep (DispatchOf e) e -> IO (EffectRep (DispatchOf e) e))
  -> IO ()
modifyEnv :: forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es
-> (EffectRep (DispatchOf e) e -> IO (EffectRep (DispatchOf e) e))
-> IO ()
modifyEnv Env es
env EffectRep (DispatchOf e) e -> IO (EffectRep (DispatchOf e) e)
f = do
  (Int
i, SmallMutableArray RealWorld Any
es) <- forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es -> IO (Int, SmallMutableArray RealWorld Any)
getLocation @e Env es
env
  EffectRep (DispatchOf e) e
e <- EffectRep (DispatchOf e) e -> IO (EffectRep (DispatchOf e) e)
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Any -> a
fromAny forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
es Int
i
  forall a. SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' SmallMutableArray RealWorld Any
es Int
i (forall a. a -> Any
toAny EffectRep (DispatchOf e) e
e)

-- | Determine location of the effect in the environment.
getLocation
  :: forall e es. e :> es
  => Env es
  -> IO (Int, SmallMutableArray RealWorld Any)
getLocation :: forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es -> IO (Int, SmallMutableArray RealWorld Any)
getLocation (Env Int
offset PrimArray Int
refs IORef' Storage
storage) = do
  let i :: Int
i       = Int
offset forall a. Num a => a -> a -> a
+ Int
2 forall a. Num a => a -> a -> a
* forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es
      ref :: Int
ref     = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
refs  Int
i
      version :: Int
version = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
refs (Int
i forall a. Num a => a -> a -> a
+ Int
1)
  Storage Int
_ Int
_ MutablePrimArray RealWorld Int
vs SmallMutableArray RealWorld Any
es SmallMutableArray RealWorld Any
_ <- forall a. IORef' a -> IO a
readIORef' IORef' Storage
storage
  Int
storageVersion <- forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray RealWorld Int
vs Int
ref
  -- If version of the reference is different than version in the storage, it
  -- means that the effect in the storage is not the one that was initially
  -- referenced.
  forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
when (Int
version forall a. Eq a => a -> a -> Bool
/= Int
storageVersion) forall a b. (a -> b) -> a -> b
$ do
    forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"version (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
version forall a. [a] -> [a] -> [a]
++ [Char]
") /= storageVersion ("
         forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
storageVersion forall a. [a] -> [a] -> [a]
++ [Char]
")"
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int
ref, SmallMutableArray RealWorld Any
es)

----------------------------------------
-- Internal helpers

-- | Create an empty storage.
emptyStorage :: IO Storage
emptyStorage :: IO Storage
emptyStorage = Int
-> Int
-> MutablePrimArray RealWorld Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage Int
0 (Int
noVersion forall a. Num a => a -> a -> a
+ Int
1)
  forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
0
  forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
0 forall a. HasCallStack => a
undefinedData
  forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
0 forall a. HasCallStack => a
undefinedData

-- | Insert an effect into the storage and return its reference.
insertEffect
  :: IORef' Storage
  -> EffectRep (DispatchOf e) e
  -- ^ The representation of the effect.
  -> Relinker (EffectRep (DispatchOf e)) e
  -> IO (Int, Int)
insertEffect :: forall (e :: Effect).
IORef' Storage
-> EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e
-> IO (Int, Int)
insertEffect IORef' Storage
storage EffectRep (DispatchOf e) e
e Relinker (EffectRep (DispatchOf e)) e
f = do
  Storage Int
size Int
version MutablePrimArray RealWorld Int
vs0 SmallMutableArray RealWorld Any
es0 SmallMutableArray RealWorld Any
fs0 <- forall a. IORef' a -> IO a
readIORef' IORef' Storage
storage
  Int
len0 <- forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m Int
getSizeofSmallMutableArray SmallMutableArray RealWorld Any
es0
  case Int
size forall a. Ord a => a -> a -> Ordering
`compare` Int
len0 of
    Ordering
GT -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"size (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
size forall a. [a] -> [a] -> [a]
++ [Char]
") > len0 (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
len0 forall a. [a] -> [a] -> [a]
++ [Char]
")"
    Ordering
LT -> do
      forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray   MutablePrimArray RealWorld Int
vs0 Int
size Int
version
      forall a. SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' SmallMutableArray RealWorld Any
es0 Int
size (forall a. a -> Any
toAny EffectRep (DispatchOf e) e
e)
      forall a. SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' SmallMutableArray RealWorld Any
fs0 Int
size (forall a. a -> Any
toAny Relinker (EffectRep (DispatchOf e)) e
f)
      forall a. IORef' a -> a -> IO ()
writeIORef' IORef' Storage
storage forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> MutablePrimArray RealWorld Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage (Int
size forall a. Num a => a -> a -> a
+ Int
1) (Int
version forall a. Num a => a -> a -> a
+ Int
1) MutablePrimArray RealWorld Int
vs0 SmallMutableArray RealWorld Any
es0 SmallMutableArray RealWorld Any
fs0
      forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int
size, Int
version)
    Ordering
EQ -> do
      let len :: Int
len = Int -> Int
doubleCapacity Int
len0
      MutablePrimArray RealWorld Int
vs <- forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
      SmallMutableArray RealWorld Any
es <- forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len forall a. HasCallStack => a
undefinedData
      SmallMutableArray RealWorld Any
fs <- forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len forall a. HasCallStack => a
undefinedData
      forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
copyMutablePrimArray  MutablePrimArray RealWorld Int
vs Int
0 MutablePrimArray RealWorld Int
vs0 Int
0 Int
size
      forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
copySmallMutableArray SmallMutableArray RealWorld Any
es Int
0 SmallMutableArray RealWorld Any
es0 Int
0 Int
size
      forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
copySmallMutableArray SmallMutableArray RealWorld Any
fs Int
0 SmallMutableArray RealWorld Any
fs0 Int
0 Int
size
      forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray   MutablePrimArray RealWorld Int
vs Int
size Int
version
      forall a. SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' SmallMutableArray RealWorld Any
es Int
size (forall a. a -> Any
toAny EffectRep (DispatchOf e) e
e)
      forall a. SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' SmallMutableArray RealWorld Any
fs Int
size (forall a. a -> Any
toAny Relinker (EffectRep (DispatchOf e)) e
f)
      forall a. IORef' a -> a -> IO ()
writeIORef' IORef' Storage
storage forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> MutablePrimArray RealWorld Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage (Int
size forall a. Num a => a -> a -> a
+ Int
1) (Int
version forall a. Num a => a -> a -> a
+ Int
1) MutablePrimArray RealWorld Int
vs SmallMutableArray RealWorld Any
es SmallMutableArray RealWorld Any
fs
      forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int
size, Int
version)

-- | Given a reference to an effect from the top of the stack, delete it from
-- the storage.
deleteEffect :: IORef' Storage -> Int -> IO ()
deleteEffect :: IORef' Storage -> Int -> IO ()
deleteEffect IORef' Storage
storage Int
ref = do
  Storage Int
size Int
version MutablePrimArray RealWorld Int
vs SmallMutableArray RealWorld Any
es SmallMutableArray RealWorld Any
fs <- forall a. IORef' a -> IO a
readIORef' IORef' Storage
storage
  forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
when (Int
ref forall a. Eq a => a -> a -> Bool
/= Int
size forall a. Num a => a -> a -> a
- Int
1) forall a b. (a -> b) -> a -> b
$ do
    forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"ref (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
ref forall a. [a] -> [a] -> [a]
++ [Char]
") /= size - 1 (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show (Int
size forall a. Num a => a -> a -> a
- Int
1) forall a. [a] -> [a] -> [a]
++ [Char]
")"
  forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray  MutablePrimArray RealWorld Int
vs Int
ref Int
noVersion
  forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
es Int
ref forall a. HasCallStack => a
undefinedData
  forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
fs Int
ref forall a. HasCallStack => a
undefinedData
  forall a. IORef' a -> a -> IO ()
writeIORef' IORef' Storage
storage forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> MutablePrimArray RealWorld Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage (Int
size forall a. Num a => a -> a -> a
- Int
1) Int
version MutablePrimArray RealWorld Int
vs SmallMutableArray RealWorld Any
es SmallMutableArray RealWorld Any
fs

-- | Relink the environment to use the new storage.
relinkEnv :: IORef' Storage -> Env es -> IO (Env es)
relinkEnv :: forall (es :: [Effect]). IORef' Storage -> Env es -> IO (Env es)
relinkEnv IORef' Storage
storage (Env Int
offset PrimArray Int
refs IORef' Storage
_) = forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
Int -> PrimArray Int -> IORef' Storage -> Env es
Env Int
offset PrimArray Int
refs IORef' Storage
storage

-- | Double the capacity of an array.
doubleCapacity :: Int -> Int
doubleCapacity :: Int -> Int
doubleCapacity Int
n = forall a. Ord a => a -> a -> a
max Int
1 Int
n forall a. Num a => a -> a -> a
* Int
2

noVersion :: Int
noVersion :: Int
noVersion = Int
0

undefinedData :: HasCallStack => a
undefinedData :: forall a. HasCallStack => a
undefinedData = forall a. HasCallStack => [Char] -> a
error [Char]
"undefined data"

-- | A strict version of 'writeSmallArray'.
writeSmallArray' :: SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' :: forall a. SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' SmallMutableArray RealWorld a
arr Int
i a
a = a
a seq :: forall a b. a -> b -> b
`seq` forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld a
arr Int
i a
a

#if !MIN_VERSION_primitive(0,9,0)
getSizeofSmallMutableArray :: SmallMutableArray RealWorld a -> IO Int
getSizeofSmallMutableArray arr = pure $! sizeofSmallMutableArray arr
#endif