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

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

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

    -- * Operations
  , emptyEnv
  , cloneEnv
  , forkEnv
  , sizeEnv
  , checkSizeEnv
  , tailEnv

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

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

import Control.Monad
import Control.Monad.Primitive
import Data.IORef
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'.
--
-- Supports forking, i.e. introduction of local branches for encapsulation of
-- effects specific to effect handlers.
--
-- __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(1)@/ (amortized).
--
-- - Shrinking: /@O(1)@/.
--
-- - Indexing via '(:>)': /@O(1)@/
--
-- - Modification of a specific element: /@O(1)@/.
--
-- - Forking: /@O(n)@/, where @n@ is the size of the effect stack.
--
-- - Cloning: /@O(N + Σ(n_i))@/, where @N@ is the size of the 'Storage', while
--   @i@ ranges over handlers of dynamically dispatched effects in the 'Storage'
--   and @n_i@ is the size of the effect stack of @i@-th handler.
--
data Env (es :: [Effect]) = Env
  { Env es -> Int
envSize    :: !Int
  , Env es -> IORef References
envRefs    :: !(IORef References)
  , Env es -> IORef Storage
envStorage :: !(IORef Storage)
  }

-- | An array of references to effects in the 'Storage'.
data References = References
  { References -> Int
refSize    :: !Int
  , References -> MutablePrimArray RealWorld Int
refIndices :: !(MutablePrimArray RealWorld Int)
  }

-- | A storage of effects.
--
-- Shared between all forks of the environment within the same thread.
data Storage = Storage
  { Storage -> Int
stSize      :: !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 :: Relinker rep e
dummyRelinker = ((forall (es :: [Effect]). Env es -> IO (Env es))
 -> rep e -> IO (rep e))
-> Relinker rep e
forall (rep :: Effect -> Type) (e :: Effect).
((forall (es :: [Effect]). Env es -> IO (Env es))
 -> rep e -> IO (rep e))
-> Relinker rep e
Relinker (((forall (es :: [Effect]). Env es -> IO (Env es))
  -> rep e -> IO (rep e))
 -> Relinker rep e)
-> ((forall (es :: [Effect]). Env es -> IO (Env es))
    -> rep e -> IO (rep e))
-> Relinker rep e
forall a b. (a -> b) -> a -> b
$ \forall (es :: [Effect]). Env es -> IO (Env es)
_ -> rep e -> IO (rep e)
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 = Int -> IORef References -> IORef Storage -> Env '[]
forall (es :: [Effect]).
Int -> IORef References -> IORef Storage -> Env es
Env (Int -> IORef References -> IORef Storage -> Env '[])
-> IO Int -> IO (IORef References -> IORef Storage -> Env '[])
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO Int
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Int
0
  IO (IORef References -> IORef Storage -> Env '[])
-> IO (IORef References) -> IO (IORef Storage -> Env '[])
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (References -> IO (IORef References)
forall a. a -> IO (IORef a)
newIORef (References -> IO (IORef References))
-> (MutablePrimArray RealWorld Int -> References)
-> MutablePrimArray RealWorld Int
-> IO (IORef References)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> MutablePrimArray RealWorld Int -> References
References Int
0 (MutablePrimArray RealWorld Int -> IO (IORef References))
-> IO (MutablePrimArray RealWorld Int) -> IO (IORef References)
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
0)
  IO (IORef Storage -> Env '[]) -> IO (IORef Storage) -> IO (Env '[])
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (Storage -> IO (IORef Storage)
forall a. a -> IO (IORef a)
newIORef (Storage -> IO (IORef Storage)) -> IO Storage -> IO (IORef Storage)
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 :: Env es -> IO (Env es)
cloneEnv (Env Int
size IORef References
mrefs0 IORef Storage
storage0) = do
  References Int
n MutablePrimArray RealWorld Int
refs0 <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs0
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  IORef References
mrefs <- References -> IO (IORef References)
forall a. a -> IO (IORef a)
newIORef (References -> IO (IORef References))
-> (MutablePrimArray RealWorld Int -> References)
-> MutablePrimArray RealWorld Int
-> IO (IORef References)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> MutablePrimArray RealWorld Int -> References
References Int
n
    (MutablePrimArray RealWorld Int -> IO (IORef References))
-> IO (MutablePrimArray RealWorld Int) -> IO (IORef References)
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< MutablePrimArray (PrimState IO) Int
-> Int -> Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0 Int
0 (MutablePrimArray RealWorld Int -> Int
forall s a. Prim a => MutablePrimArray s a -> Int
sizeofMutablePrimArray MutablePrimArray RealWorld Int
refs0)
  Storage Int
storageSize SmallMutableArray RealWorld Any
es0 SmallMutableArray RealWorld Any
fs0 <- IORef Storage -> IO Storage
forall a. IORef a -> IO a
readIORef IORef Storage
storage0
  let esSize :: Int
esSize = SmallMutableArray RealWorld Any -> Int
forall s a. SmallMutableArray s a -> Int
sizeofSmallMutableArray SmallMutableArray RealWorld Any
es0
      fsSize :: Int
fsSize = SmallMutableArray RealWorld Any -> Int
forall s a. SmallMutableArray s a -> Int
sizeofSmallMutableArray SmallMutableArray RealWorld Any
fs0
  Bool -> IO () -> IO ()
forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
when (Int
esSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
fsSize) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char]
"esSize (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
esSize [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") /= fsSize (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
fsSize [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
  SmallMutableArray RealWorld Any
es <- SmallMutableArray (PrimState IO) Any
-> Int -> Int -> IO (SmallMutableArray (PrimState IO) Any)
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> Int -> m (SmallMutableArray (PrimState m) a)
cloneSmallMutableArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es0 Int
0 Int
esSize
  SmallMutableArray RealWorld Any
fs <- SmallMutableArray (PrimState IO) Any
-> Int -> Int -> IO (SmallMutableArray (PrimState IO) Any)
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> Int -> m (SmallMutableArray (PrimState m) a)
cloneSmallMutableArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
fs0 Int
0 Int
esSize
  IORef Storage
storage <- Storage -> IO (IORef Storage)
forall a. a -> IO (IORef a)
newIORef (Storage -> IO (IORef Storage)) -> Storage -> IO (IORef Storage)
forall a b. (a -> b) -> a -> b
$ Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage Int
storageSize SmallMutableArray RealWorld Any
es SmallMutableArray RealWorld Any
fs
  let relinkEffects :: Int -> IO ()
relinkEffects = \case
        Int
0 -> () -> IO ()
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()
        Int
k -> do
          let i :: Int
i = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
          Relinker (forall (es :: [Effect]). Env es -> IO (Env es))
-> Any Any -> IO (Any Any)
f <- Any -> Relinker Any Any
forall a. Any -> a
fromAny (Any -> Relinker Any Any) -> IO Any -> IO (Relinker Any Any)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> SmallMutableArray (PrimState IO) Any -> Int -> IO Any
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
fs Int
i
          SmallMutableArray (PrimState IO) Any -> Int -> IO Any
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
i
            IO Any -> (Any -> IO (Any Any)) -> IO (Any Any)
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 (IORef Storage -> Env es -> IO (Env es)
forall (es :: [Effect]). IORef Storage -> Env es -> IO (Env es)
relinkEnv IORef Storage
storage) (Any Any -> IO (Any Any))
-> (Any -> Any Any) -> Any -> IO (Any Any)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Any -> Any Any
forall a. Any -> a
fromAny
            IO (Any Any) -> (Any Any -> IO ()) -> IO ()
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
i (Any -> IO ()) -> (Any Any -> Any) -> Any Any -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Any Any -> Any
forall a. a -> Any
toAny
          Int -> IO ()
relinkEffects Int
i
  Int -> IO ()
relinkEffects Int
storageSize
  Env es -> IO (Env es)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Env es -> IO (Env es)) -> Env es -> IO (Env es)
forall a b. (a -> b) -> a -> b
$ Int -> IORef References -> IORef Storage -> Env es
forall (es :: [Effect]).
Int -> IORef References -> IORef Storage -> Env es
Env Int
size IORef References
mrefs IORef Storage
storage
{-# NOINLINE cloneEnv #-}

-- | Create a fork of the environment.
--
-- Forked environment can be updated independently of the original one within
-- the same thread.
forkEnv :: Env es -> IO (Env es)
forkEnv :: Env es -> IO (Env es)
forkEnv (Env Int
size IORef References
mrefs0 IORef Storage
storage) = do
  References Int
n MutablePrimArray RealWorld Int
refs0 <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs0
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  IORef References
mrefs <- References -> IO (IORef References)
forall a. a -> IO (IORef a)
newIORef (References -> IO (IORef References))
-> (MutablePrimArray RealWorld Int -> References)
-> MutablePrimArray RealWorld Int
-> IO (IORef References)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> MutablePrimArray RealWorld Int -> References
References Int
size
    (MutablePrimArray RealWorld Int -> IO (IORef References))
-> IO (MutablePrimArray RealWorld Int) -> IO (IORef References)
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< MutablePrimArray (PrimState IO) Int
-> Int -> Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0 Int
0 (MutablePrimArray RealWorld Int -> Int
forall s a. Prim a => MutablePrimArray s a -> Int
sizeofMutablePrimArray MutablePrimArray RealWorld Int
refs0)
  Env es -> IO (Env es)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Env es -> IO (Env es)) -> Env es -> IO (Env es)
forall a b. (a -> b) -> a -> b
$ Int -> IORef References -> IORef Storage -> Env es
forall (es :: [Effect]).
Int -> IORef References -> IORef Storage -> Env es
Env Int
size IORef References
mrefs IORef Storage
storage
{-# NOINLINE forkEnv #-}

-- | Check that the size of the environment is internally consistent.
checkSizeEnv :: Env es -> IO ()
checkSizeEnv :: Env es -> IO ()
checkSizeEnv (Env Int
size IORef References
mrefs IORef Storage
_) = do
  References Int
n MutablePrimArray RealWorld Int
_ <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
{-# NOINLINE checkSizeEnv #-}

-- | Get the current size of the environment.
sizeEnv :: Env es -> IO Int
sizeEnv :: Env es -> IO Int
sizeEnv Env es
env = Int -> IO Int
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Env es -> Int
forall (es :: [Effect]). Env es -> Int
envSize Env es
env

-- | Access the tail of the environment.
tailEnv :: Env (e : es) -> IO (Env es)
tailEnv :: Env (e : es) -> IO (Env es)
tailEnv (Env Int
size IORef References
mrefs0 IORef Storage
storage) = do
  References Int
n MutablePrimArray RealWorld Int
refs0 <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs0
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  IORef References
mrefs <- References -> IO (IORef References)
forall a. a -> IO (IORef a)
newIORef (References -> IO (IORef References))
-> (MutablePrimArray RealWorld Int -> References)
-> MutablePrimArray RealWorld Int
-> IO (IORef References)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> MutablePrimArray RealWorld Int -> References
References (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    (MutablePrimArray RealWorld Int -> IO (IORef References))
-> IO (MutablePrimArray RealWorld Int) -> IO (IORef References)
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< MutablePrimArray (PrimState IO) Int
-> Int -> Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0 Int
0 (MutablePrimArray RealWorld Int -> Int
forall s a. Prim a => MutablePrimArray s a -> Int
sizeofMutablePrimArray MutablePrimArray RealWorld Int
refs0)
  Env es -> IO (Env es)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Env es -> IO (Env es)) -> Env es -> IO (Env es)
forall a b. (a -> b) -> a -> b
$ Int -> IORef References -> IORef Storage -> Env es
forall (es :: [Effect]).
Int -> IORef References -> IORef Storage -> Env es
Env (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) IORef References
mrefs IORef Storage
storage
{-# NOINLINE tailEnv #-}

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

-- | Extend the environment with a new data type (in place).
consEnv
  :: EffectRep (DispatchOf e) e
  -- ^ The representation of the effect.
  -> Relinker (EffectRep (DispatchOf e)) e
  -> Env es
  -> IO (Env (e : es))
consEnv :: 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
size IORef References
mrefs IORef Storage
storage) = do
  References Int
n MutablePrimArray RealWorld Int
refs0 <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  Int
len0 <- MutablePrimArray (PrimState IO) Int -> IO Int
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0
  MutablePrimArray RealWorld Int
refs <- case Int
size Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
len0 of
    Ordering
GT -> [Char] -> IO (MutablePrimArray RealWorld Int)
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO (MutablePrimArray RealWorld Int))
-> [Char] -> IO (MutablePrimArray RealWorld Int)
forall a b. (a -> b) -> a -> b
$ [Char]
"size (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
size [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") > len0 (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
len0 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
    Ordering
LT -> MutablePrimArray RealWorld Int
-> IO (MutablePrimArray RealWorld Int)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MutablePrimArray RealWorld Int
refs0
    Ordering
EQ -> MutablePrimArray (PrimState IO) Int
-> Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
resizeMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0 (Int -> Int
doubleCapacity Int
len0)
  Int
ref <- IORef Storage
-> EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e
-> IO Int
forall (e :: Effect).
IORef Storage
-> EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e
-> IO Int
insertEffect IORef Storage
storage EffectRep (DispatchOf e) e
e Relinker (EffectRep (DispatchOf e)) e
f
  MutablePrimArray (PrimState IO) Int -> Int -> Int -> IO ()
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs Int
size Int
ref
  IORef References -> References -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef References
mrefs (References -> IO ()) -> References -> IO ()
forall a b. (a -> b) -> a -> b
$! Int -> MutablePrimArray RealWorld Int -> References
References (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) MutablePrimArray RealWorld Int
refs
  Env (e : es) -> IO (Env (e : es))
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Env (e : es) -> IO (Env (e : es)))
-> Env (e : es) -> IO (Env (e : es))
forall a b. (a -> b) -> a -> b
$ Int -> IORef References -> IORef Storage -> Env (e : es)
forall (es :: [Effect]).
Int -> IORef References -> IORef Storage -> Env es
Env (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) IORef References
mrefs IORef Storage
storage
{-# NOINLINE consEnv #-}

-- | Shrink the environment by one data type (in place).
--
-- /Note:/ after calling this function the input environment is no longer
-- usable.
unconsEnv :: Env (e : es) -> IO ()
unconsEnv :: Env (e : es) -> IO ()
unconsEnv (Env Int
size IORef References
mrefs IORef Storage
storage) = do
  References Int
n MutablePrimArray RealWorld Int
refs <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  Int
ref <- MutablePrimArray (PrimState IO) Int -> Int -> IO Int
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  IORef Storage -> Int -> IO ()
deleteEffect IORef Storage
storage Int
ref
  IORef References -> References -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef References
mrefs (References -> IO ()) -> References -> IO ()
forall a b. (a -> b) -> a -> b
$! Int -> MutablePrimArray RealWorld Int -> References
References (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) MutablePrimArray RealWorld Int
refs
{-# NOINLINE unconsEnv #-}

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

-- | Replace a specific effect in the stack with a new value.
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 :: 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
size IORef References
mrefs0 IORef Storage
storage) = do
  References Int
n MutablePrimArray RealWorld Int
refs0 <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs0
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  Int
len0 <- MutablePrimArray (PrimState IO) Int -> IO Int
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0
  Bool -> IO () -> IO ()
forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
when (Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
len0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char]
"size (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
size [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") > len0 (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
len0 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
  MutablePrimArray RealWorld Int
refs <- MutablePrimArray (PrimState IO) Int
-> Int -> Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0 Int
0 Int
len0
  Int
ref <- IORef Storage
-> EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e
-> IO Int
forall (e :: Effect).
IORef Storage
-> EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e
-> IO Int
insertEffect IORef Storage
storage EffectRep (DispatchOf e) e
e Relinker (EffectRep (DispatchOf e)) e
f
  MutablePrimArray (PrimState IO) Int -> Int -> Int -> IO ()
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs (Int -> Int -> Int
mkIndex ((e :> es) => Int
forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es) Int
size) Int
ref
  IORef References
mrefs <- References -> IO (IORef References)
forall a. a -> IO (IORef a)
newIORef (References -> IO (IORef References))
-> References -> IO (IORef References)
forall a b. (a -> b) -> a -> b
$ Int -> MutablePrimArray RealWorld Int -> References
References Int
n MutablePrimArray RealWorld Int
refs
  Env es -> IO (Env es)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Env es -> IO (Env es)) -> Env es -> IO (Env es)
forall a b. (a -> b) -> a -> b
$ Int -> IORef References -> IORef Storage -> Env es
forall (es :: [Effect]).
Int -> IORef References -> IORef Storage -> Env es
Env Int
size IORef References
mrefs 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 :: Env es -> IO ()
unreplaceEnv (Env Int
size IORef References
mrefs IORef Storage
storage) = do
  References Int
n MutablePrimArray RealWorld Int
refs <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  Int
ref <- MutablePrimArray (PrimState IO) Int -> Int -> IO Int
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
mkIndex ((e :> es) => Int
forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es) Int
size
  IORef Storage -> Int -> IO ()
deleteEffect IORef Storage
storage Int
ref
{-# NOINLINE unreplaceEnv #-}

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

-- | Reference an existing effect from the top of the stack (in place).
subsumeEnv :: forall e es. e :> es => Env es -> IO (Env (e : es))
subsumeEnv :: Env es -> IO (Env (e : es))
subsumeEnv (Env Int
size IORef References
mrefs IORef Storage
storage) = do
  References Int
n MutablePrimArray RealWorld Int
refs0 <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  Int
len0 <- MutablePrimArray (PrimState IO) Int -> IO Int
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0
  MutablePrimArray RealWorld Int
refs <- case Int
size Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
len0 of
    Ordering
GT -> [Char] -> IO (MutablePrimArray RealWorld Int)
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO (MutablePrimArray RealWorld Int))
-> [Char] -> IO (MutablePrimArray RealWorld Int)
forall a b. (a -> b) -> a -> b
$ [Char]
"size (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
size [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") > len0 (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
len0 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
    Ordering
LT -> MutablePrimArray RealWorld Int
-> IO (MutablePrimArray RealWorld Int)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MutablePrimArray RealWorld Int
refs0
    Ordering
EQ -> MutablePrimArray (PrimState IO) Int
-> Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
resizeMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0 (Int -> Int
doubleCapacity Int
len0)
  Int
ref <- MutablePrimArray (PrimState IO) Int -> Int -> IO Int
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
mkIndex ((e :> es) => Int
forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es) Int
size
  MutablePrimArray (PrimState IO) Int -> Int -> Int -> IO ()
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs Int
size Int
ref
  IORef References -> References -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef References
mrefs (References -> IO ()) -> References -> IO ()
forall a b. (a -> b) -> a -> b
$! Int -> MutablePrimArray RealWorld Int -> References
References (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) MutablePrimArray RealWorld Int
refs
  Env (e : es) -> IO (Env (e : es))
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Env (e : es) -> IO (Env (e : es)))
-> Env (e : es) -> IO (Env (e : es))
forall a b. (a -> b) -> a -> b
$ Int -> IORef References -> IORef Storage -> Env (e : es)
forall (es :: [Effect]).
Int -> IORef References -> IORef Storage -> Env es
Env (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) IORef References
mrefs IORef Storage
storage
{-# NOINLINE subsumeEnv #-}

-- | Remove a reference to an existing effect from the top of the stack.
--
-- /Note:/ after calling this function the input environment is no longer
-- usable.
unsubsumeEnv :: e :> es => Env (e : es) -> IO ()
unsubsumeEnv :: Env (e : es) -> IO ()
unsubsumeEnv (Env Int
size IORef References
mrefs IORef Storage
_) = do
  References Int
n MutablePrimArray RealWorld Int
refs <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  IORef References -> References -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef References
mrefs (References -> IO ()) -> References -> IO ()
forall a b. (a -> b) -> a -> b
$! Int -> MutablePrimArray RealWorld Int -> References
References (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) MutablePrimArray RealWorld Int
refs
{-# NOINLINE unsubsumeEnv #-}

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

-- | 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 :: Env es -> IO (Env xs)
injectEnv (Env Int
size0 IORef References
mrefs0 IORef Storage
storage) = do
  References Int
n0 MutablePrimArray RealWorld Int
refs0 <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs0
  HasCallStack => Int -> Int -> IO ()
Int -> Int -> IO ()
errorWhenDifferent Int
size0 Int
n0
  let makeRefs :: Int -> [Int] -> [Int] -> IO (MutablePrimArray RealWorld Int)
makeRefs Int
k [Int]
acc = \case
        []       -> PrimArray Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: Type -> Type) a.
PrimMonad m =>
PrimArray a -> m (MutablePrimArray (PrimState m) a)
unsafeThawPrimArray (PrimArray Int -> IO (MutablePrimArray (PrimState IO) Int))
-> PrimArray Int -> IO (MutablePrimArray (PrimState IO) Int)
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> PrimArray Int
forall a. Prim a => Int -> [a] -> PrimArray a
primArrayFromListN Int
k [Int]
acc
        (Int
e : [Int]
es) -> do
          Int
i <- MutablePrimArray (PrimState IO) Int -> Int -> IO Int
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0 (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
mkIndex Int
e Int
size0
          Int -> [Int] -> [Int] -> IO (MutablePrimArray RealWorld Int)
makeRefs (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
acc) [Int]
es
  MutablePrimArray RealWorld Int
refs <- Int -> [Int] -> [Int] -> IO (MutablePrimArray RealWorld Int)
makeRefs Int
0 [] (Subset xs es => [Int]
forall (xs :: [Effect]) (es :: [Effect]). Subset xs es => [Int]
reifyIndices @xs @es)
  Int
size <- MutablePrimArray (PrimState IO) Int -> IO Int
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs
  IORef References
mrefs <- References -> IO (IORef References)
forall a. a -> IO (IORef a)
newIORef (References -> IO (IORef References))
-> References -> IO (IORef References)
forall a b. (a -> b) -> a -> b
$ Int -> MutablePrimArray RealWorld Int -> References
References Int
size MutablePrimArray RealWorld Int
refs
  Env xs -> IO (Env xs)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Env xs -> IO (Env xs)) -> Env xs -> IO (Env xs)
forall a b. (a -> b) -> a -> b
$ Int -> IORef References -> IORef Storage -> Env xs
forall (es :: [Effect]).
Int -> IORef References -> IORef Storage -> Env es
Env Int
size IORef References
mrefs 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 :: Env es -> IO (EffectRep (DispatchOf e) e)
getEnv Env es
env = do
  (Int
i, SmallMutableArray RealWorld Any
es) <- Env es -> IO (Int, SmallMutableArray RealWorld Any)
forall (e :: Effect) (es :: [Effect]).
(e :> es) =>
Env es -> IO (Int, SmallMutableArray RealWorld Any)
getLocation @e Env es
env
  Any -> EffectRep (DispatchOf e) e
forall a. Any -> a
fromAny (Any -> EffectRep (DispatchOf e) e)
-> IO Any -> IO (EffectRep (DispatchOf e) e)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> SmallMutableArray (PrimState IO) Any -> Int -> IO Any
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) 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 :: Env es -> EffectRep (DispatchOf e) e -> IO ()
putEnv Env es
env EffectRep (DispatchOf e) e
e = do
  (Int
i, SmallMutableArray RealWorld Any
es) <- Env es -> IO (Int, SmallMutableArray RealWorld Any)
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 () -> IO ()
`seq` SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
i (EffectRep (DispatchOf e) e -> Any
forall a. a -> Any
toAny EffectRep (DispatchOf e) e
e)

-- | Modify the data type in the environment (in place) and return a value.
stateEnv
  :: forall e es a. e :> es
  => Env es -- ^ The environment.
  -> (EffectRep (DispatchOf e) e -> (a, EffectRep (DispatchOf e) e))
  -> IO a
stateEnv :: Env es
-> (EffectRep (DispatchOf e) e -> (a, EffectRep (DispatchOf e) e))
-> IO a
stateEnv Env es
env EffectRep (DispatchOf e) e -> (a, EffectRep (DispatchOf e) e)
f = do
  (Int
i, SmallMutableArray RealWorld Any
es) <- Env es -> IO (Int, SmallMutableArray RealWorld Any)
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 -> (a, EffectRep (DispatchOf e) e)
f (EffectRep (DispatchOf e) e -> (a, EffectRep (DispatchOf e) e))
-> (Any -> EffectRep (DispatchOf e) e)
-> Any
-> (a, EffectRep (DispatchOf e) e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Any -> EffectRep (DispatchOf e) e
forall a. Any -> a
fromAny (Any -> (a, EffectRep (DispatchOf e) e))
-> IO Any -> IO (a, EffectRep (DispatchOf e) e)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> SmallMutableArray (PrimState IO) Any -> Int -> IO Any
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
i
  EffectRep (DispatchOf e) e
e EffectRep (DispatchOf e) e -> IO () -> IO ()
`seq` SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
i (EffectRep (DispatchOf e) e -> Any
forall a. a -> Any
toAny EffectRep (DispatchOf e) e
e)
  a -> IO a
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 -> EffectRep (DispatchOf e) e)
  -> IO ()
modifyEnv :: Env es
-> (EffectRep (DispatchOf e) e -> EffectRep (DispatchOf e) e)
-> IO ()
modifyEnv Env es
env EffectRep (DispatchOf e) e -> EffectRep (DispatchOf e) e
f = do
  (Int
i, SmallMutableArray RealWorld Any
es) <- Env es -> IO (Int, SmallMutableArray RealWorld Any)
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 -> EffectRep (DispatchOf e) e
f (EffectRep (DispatchOf e) e -> EffectRep (DispatchOf e) e)
-> (Any -> EffectRep (DispatchOf e) e)
-> Any
-> EffectRep (DispatchOf e) e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Any -> EffectRep (DispatchOf e) e
forall a. Any -> a
fromAny (Any -> EffectRep (DispatchOf e) e)
-> IO Any -> IO (EffectRep (DispatchOf e) e)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> SmallMutableArray (PrimState IO) Any -> Int -> IO Any
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
i
  EffectRep (DispatchOf e) e
e EffectRep (DispatchOf e) e -> IO () -> IO ()
`seq` SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
i (EffectRep (DispatchOf e) e -> Any
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 :: Env es -> IO (Int, SmallMutableArray RealWorld Any)
getLocation (Env Int
size IORef References
mrefs IORef Storage
storage) = do
  MutablePrimArray RealWorld Int
refs <- References -> MutablePrimArray RealWorld Int
refIndices (References -> MutablePrimArray RealWorld Int)
-> IO References -> IO (MutablePrimArray RealWorld Int)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs
  Int
i <- MutablePrimArray (PrimState IO) Int -> Int -> IO Int
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
mkIndex ((e :> es) => Int
forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es) Int
size
  SmallMutableArray RealWorld Any
es <- Storage -> SmallMutableArray RealWorld Any
stEffects (Storage -> SmallMutableArray RealWorld Any)
-> IO Storage -> IO (SmallMutableArray RealWorld Any)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef Storage -> IO Storage
forall a. IORef a -> IO a
readIORef IORef Storage
storage
  (Int, SmallMutableArray RealWorld Any)
-> IO (Int, SmallMutableArray RealWorld Any)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int
i, SmallMutableArray RealWorld Any
es)

-- | Get the index of a reference to an effect.
mkIndex :: Int -> Int -> Int
mkIndex :: Int -> Int -> Int
mkIndex Int
ix Int
size = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

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

-- | Create an empty storage.
emptyStorage :: IO Storage
emptyStorage :: IO Storage
emptyStorage = Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage
  (Int
 -> SmallMutableArray RealWorld Any
 -> SmallMutableArray RealWorld Any
 -> Storage)
-> IO Int
-> IO
     (SmallMutableArray RealWorld Any
      -> SmallMutableArray RealWorld Any -> Storage)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO Int
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Int
0
  IO
  (SmallMutableArray RealWorld Any
   -> SmallMutableArray RealWorld Any -> Storage)
-> IO (SmallMutableArray RealWorld Any)
-> IO (SmallMutableArray RealWorld Any -> Storage)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Int -> Any -> IO (SmallMutableArray (PrimState IO) Any)
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
0 Any
forall a. HasCallStack => a
undefinedData
  IO (SmallMutableArray RealWorld Any -> Storage)
-> IO (SmallMutableArray RealWorld Any) -> IO Storage
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Int -> Any -> IO (SmallMutableArray (PrimState IO) Any)
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
0 Any
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
insertEffect :: IORef Storage
-> EffectRep (DispatchOf e) e
-> Relinker (EffectRep (DispatchOf e)) e
-> IO Int
insertEffect IORef Storage
storage EffectRep (DispatchOf e) e
e Relinker (EffectRep (DispatchOf e)) e
f = do
  Storage Int
size SmallMutableArray RealWorld Any
es0 SmallMutableArray RealWorld Any
fs0 <- IORef Storage -> IO Storage
forall a. IORef a -> IO a
readIORef IORef Storage
storage
  let len0 :: Int
len0 = SmallMutableArray RealWorld Any -> Int
forall s a. SmallMutableArray s a -> Int
sizeofSmallMutableArray SmallMutableArray RealWorld Any
es0
  case Int
size Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
len0 of
    Ordering
GT -> [Char] -> IO Int
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO Int) -> [Char] -> IO Int
forall a b. (a -> b) -> a -> b
$ [Char]
"size (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
size [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") > len0 (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
len0 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
    Ordering
LT -> do
      EffectRep (DispatchOf e) e
e EffectRep (DispatchOf e) e -> IO () -> IO ()
`seq` SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es0 Int
size (EffectRep (DispatchOf e) e -> Any
forall a. a -> Any
toAny EffectRep (DispatchOf e) e
e)
      Relinker (EffectRep (DispatchOf e)) e
f Relinker (EffectRep (DispatchOf e)) e -> IO () -> IO ()
`seq` SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
fs0 Int
size (Relinker (EffectRep (DispatchOf e)) e -> Any
forall a. a -> Any
toAny Relinker (EffectRep (DispatchOf e)) e
f)
      IORef Storage -> Storage -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Storage
storage (Storage -> IO ()) -> Storage -> IO ()
forall a b. (a -> b) -> a -> b
$! Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) SmallMutableArray RealWorld Any
es0 SmallMutableArray RealWorld Any
fs0
      Int -> IO Int
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Int
size
    Ordering
EQ -> do
      let len :: Int
len = Int -> Int
doubleCapacity Int
len0
      SmallMutableArray RealWorld Any
es <- Int -> Any -> IO (SmallMutableArray (PrimState IO) Any)
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len Any
forall a. HasCallStack => a
undefinedData
      SmallMutableArray RealWorld Any
fs <- Int -> Any -> IO (SmallMutableArray (PrimState IO) Any)
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len Any
forall a. HasCallStack => a
undefinedData
      SmallMutableArray (PrimState IO) Any
-> Int
-> SmallMutableArray (PrimState IO) Any
-> Int
-> Int
-> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
copySmallMutableArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
0 SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es0 Int
0 Int
size
      SmallMutableArray (PrimState IO) Any
-> Int
-> SmallMutableArray (PrimState IO) Any
-> Int
-> Int
-> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
copySmallMutableArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
fs Int
0 SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
fs0 Int
0 Int
size
      EffectRep (DispatchOf e) e
e EffectRep (DispatchOf e) e -> IO () -> IO ()
`seq` SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
size (EffectRep (DispatchOf e) e -> Any
forall a. a -> Any
toAny EffectRep (DispatchOf e) e
e)
      Relinker (EffectRep (DispatchOf e)) e
f Relinker (EffectRep (DispatchOf e)) e -> IO () -> IO ()
`seq` SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
fs Int
size (Relinker (EffectRep (DispatchOf e)) e -> Any
forall a. a -> Any
toAny Relinker (EffectRep (DispatchOf e)) e
f)
      IORef Storage -> Storage -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Storage
storage (Storage -> IO ()) -> Storage -> IO ()
forall a b. (a -> b) -> a -> b
$! Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) SmallMutableArray RealWorld Any
es SmallMutableArray RealWorld Any
fs
      Int -> IO Int
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Int
size

-- | Given a reference to an effect, delete it from the storage.
--
-- /Note:/ the reference needs to point to the end of the storage. Normally it's
-- not a problem as it turns out effects are put and taken from the storage in
-- the same order across all forks, unless someone tries to do something
-- unexpected.
deleteEffect :: IORef Storage -> Int -> IO ()
deleteEffect :: IORef Storage -> Int -> IO ()
deleteEffect IORef Storage
storage Int
ref = do
  Storage Int
size SmallMutableArray RealWorld Any
es SmallMutableArray RealWorld Any
fs <- IORef Storage -> IO Storage
forall a. IORef a -> IO a
readIORef IORef Storage
storage
  Bool -> IO () -> IO ()
forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
when (Int
ref Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char]
"ref (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
ref [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") /= size - 1 (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
  SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
es Int
ref Any
forall a. HasCallStack => a
undefinedData
  SmallMutableArray (PrimState IO) Any -> Int -> Any -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Any
SmallMutableArray (PrimState IO) Any
fs Int
ref Any
forall a. HasCallStack => a
undefinedData
  IORef Storage -> Storage -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Storage
storage (Storage -> IO ()) -> Storage -> IO ()
forall a b. (a -> b) -> a -> b
$! Int
-> SmallMutableArray RealWorld Any
-> SmallMutableArray RealWorld Any
-> Storage
Storage (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) 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 :: IORef Storage -> Env es -> IO (Env es)
relinkEnv IORef Storage
storage (Env Int
size IORef References
mrefs0 IORef Storage
_) = do
  References Int
n MutablePrimArray RealWorld Int
refs0 <- IORef References -> IO References
forall a. IORef a -> IO a
readIORef IORef References
mrefs0
  IORef References
mrefs <- References -> IO (IORef References)
forall a. a -> IO (IORef a)
newIORef (References -> IO (IORef References))
-> (MutablePrimArray RealWorld Int -> References)
-> MutablePrimArray RealWorld Int
-> IO (IORef References)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> MutablePrimArray RealWorld Int -> References
References Int
n
    (MutablePrimArray RealWorld Int -> IO (IORef References))
-> IO (MutablePrimArray RealWorld Int) -> IO (IORef References)
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< MutablePrimArray (PrimState IO) Int
-> Int -> Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
refs0 Int
0 (MutablePrimArray RealWorld Int -> Int
forall s a. Prim a => MutablePrimArray s a -> Int
sizeofMutablePrimArray MutablePrimArray RealWorld Int
refs0)
  Env es -> IO (Env es)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Env es -> IO (Env es)) -> Env es -> IO (Env es)
forall a b. (a -> b) -> a -> b
$ Int -> IORef References -> IORef Storage -> Env es
forall (es :: [Effect]).
Int -> IORef References -> IORef Storage -> Env es
Env Int
size IORef References
mrefs IORef Storage
storage

-- | Throw an error if array sizes do not agree.
errorWhenDifferent :: HasCallStack => Int -> Int -> IO ()
errorWhenDifferent :: Int -> Int -> IO ()
errorWhenDifferent Int
size Int
n
  | Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n = [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char]
"size (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
size [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") /= n (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
  | Bool
otherwise = () -> IO ()
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()

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

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