Env.hs

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

    -- ** StorageData
  , StorageData(..)
  , cloneStorage
  , replaceStorage
  , backupStorageData
  , restoreStorageData

    -- *** Utils
  , AnyRelinker
  , toAnyRelinker
  , fromAnyRelinker
  , AnyEffect
  , toAnyEffect
  , fromAnyEffect

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

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

    -- * Operations
  , emptyEnv
  , cloneEnv
  , 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.IORef.Strict qualified as S
import Data.Primitive.PrimArray
import Data.Primitive.SmallArray
import Data.Primitive.Types
import Data.Proxy
import GHC.Exts ((*#), (+#))
import GHC.Stack

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' function.
--
-- 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
  { offset  :: !Int
  , refs    :: !(PrimArray Ref)
  , storage :: !(S.IORef Storage)
  }

-- | Reference to the effect in 'Storage'.
data Ref = Ref !Int !Version

instance Prim Ref where
  sizeOfType# _ = 2# *# sizeOfType# (Proxy @Int)
  alignmentOfType# _ = alignmentOfType# (Proxy @Int)
  indexByteArray# arr i =
    let n = 2# *# i
        ref = indexByteArray# arr n
        version = indexByteArray# arr (n +# 1#)
    in Ref ref version
  readByteArray# arr i s0 =
    let n = 2# *# i
        !(# s1, ref #) = readByteArray# arr n s0
        !(# s2, version #) = readByteArray# arr (n +# 1#) s1
    in (# s2, Ref ref version #)
  writeByteArray# arr i (Ref ref version) s0 =
    let n = 2# *# i
        s1 = writeByteArray# arr n ref s0
        s2 = writeByteArray# arr (n +# 1#) version s1
    in s2
  indexOffAddr# addr i =
    let n = 2# *# i
        ref = indexOffAddr# addr n
        version = indexOffAddr# addr (n +# 1#)
    in Ref ref version
  readOffAddr# addr i s0 =
    let n = 2# *# i
        !(# s1, ref #) = readOffAddr# addr n s0
        !(# s2, version #) = readOffAddr# addr (n +# 1#) s1
    in (# s2, Ref ref version #)
  writeOffAddr# addr i (Ref ref version) s0 =
    let n = 2# *# i
        s1 = writeOffAddr# addr n ref s0
        s2 = writeOffAddr# addr (n +# 1#) version s1
    in s2

-- | Version of the effect.
newtype Version = Version Int
  deriving newtype (Eq, Ord, Prim, Show)

-- | A storage of effects.
data Storage = Storage
  { version :: !Version
  , data_   :: {-# UNPACK #-} !StorageData
  }

----------------------------------------
-- StorageData

-- | Effect in 'Storage'.
newtype AnyEffect = AnyEffect Any

toAnyEffect :: EffectRep (DispatchOf e) e -> AnyEffect
toAnyEffect = AnyEffect . toAny

fromAnyEffect :: AnyEffect -> EffectRep (DispatchOf e) e
fromAnyEffect (AnyEffect e) = fromAny e

-- | Relinker in 'Storage'.
newtype AnyRelinker = AnyRelinker Any

toAnyRelinker :: Relinker (EffectRep (DispatchOf e)) e -> AnyRelinker
toAnyRelinker = AnyRelinker . toAny

fromAnyRelinker :: AnyRelinker -> Relinker (EffectRep (DispatchOf e)) e
fromAnyRelinker (AnyRelinker f) = fromAny f

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

data StorageData = StorageData
  { size      :: !Int
  , versions  :: !(MutablePrimArray RealWorld Version)
  , effects   :: !(SmallMutableArray RealWorld AnyEffect)
  , relinkers :: !(SmallMutableArray RealWorld AnyRelinker)
  }

-- | Clone the storage to use it in a different thread.
--
-- @since 2.7.0.0
cloneStorage :: HasCallStack => S.IORef Storage -> IO (S.IORef Storage)
cloneStorage storage0 = do
  Storage version storageData0 <- S.readIORef storage0
  storageData <- copyStorageData storageData0
  storage <- S.newIORef $ Storage version storageData
  relinkStorageData storageData storage
  pure storage

-- | Replace the storage of the environment.
--
-- @since 2.7.0.0
replaceStorage :: Env es -> S.IORef Storage -> IO (Env es)
replaceStorage (Env offset refs _) storage = pure $ Env offset refs storage

-- | Make a shallow copy of the 'StorageData'.
--
-- @since 2.5.0.0
copyStorageData :: HasCallStack => StorageData -> IO StorageData
copyStorageData (StorageData storageSize vs0 es0 fs0) = do
  vsSize <- getSizeofMutablePrimArray  vs0
  esSize <- getSizeofSmallMutableArray es0
  fsSize <- getSizeofSmallMutableArray fs0
  when (vsSize /= esSize) $ do
    error $ "vsSize (" ++ show vsSize ++ ") /= esSize (" ++ show esSize ++ ")"
  when (esSize /= fsSize) $ do
    error $ "esSize (" ++ show esSize ++ ") /= fsSize (" ++ show fsSize ++ ")"
  vs <- cloneMutablePrimArray  vs0 0 vsSize
  es <- cloneSmallMutableArray es0 0 esSize
  fs <- cloneSmallMutableArray fs0 0 fsSize
  pure $ StorageData storageSize vs es fs

-- | Relink effects in the storage data to the given storage.
--
-- @since 2.7.0.0
relinkStorageData :: HasCallStack => StorageData -> S.IORef Storage -> IO ()
relinkStorageData (StorageData storageSize _ es fs) storage = go storageSize
  where
    go = \case
      0 -> pure ()
      k -> do
        let i = k - 1
        Relinker relinker <- fromAnyRelinker <$> readSmallArray fs i
        readSmallArray es i
          >>= relinker (relinkEnv storage) . fromAnyEffect
          >>= writeSmallArray' es i . toAnyEffect
        go i

-- | Backup storage data of the environment.
--
-- It can be restored later with 'restoreStorageData'.
--
-- @since 2.7.0.0
backupStorageData :: HasCallStack => Env es -> IO StorageData
backupStorageData env = do
  storageData <- copyStorageData . (.data_) =<< S.readIORef env.storage
  -- Relinking to the same storage might seem weird, but relinkers need to run
  -- and make a copy of mutable data associated with statically dispatched
  -- effects if appropriate.
  relinkStorageData storageData env.storage
  pure storageData

-- | Restore a copy of the 'StorageData'.
--
-- The copy needs to be from the same 'Env' as the target. It's consumed by this
-- operation and must not be used afterwards.
--
-- @since 2.5.0.0
restoreStorageData :: HasCallStack => StorageData -> Env es -> IO ()
restoreStorageData (StorageData newSize vs1 es1 fs1) env = do
  Storage version (StorageData oldSize vs0 es0 fs0) <- S.readIORef env.storage
  when (newSize /= oldSize) $ do
    error $ "newSize (" ++ show newSize ++ ") /= oldSize (" ++ show oldSize ++ ")"
  -- Since the time the backup was made the storage might've been grown by
  -- 'insertEffect', so if necessary create new arrays matching the current
  -- capacity, as shrinking it would violate the invariant that out of date
  -- references in 'getLocation' never read out of bounds.
  vs0size <- getSizeofMutablePrimArray vs0
  vs1size <- getSizeofMutablePrimArray vs1
  vs <- if vs0size > vs1size
    then do
      vs <- newPrimArray vs0size
      copyMutablePrimArray vs 0 vs1 0 newSize
      -- Fill the unused part of the versions array with
      -- 'undefinedVersion' to maintain the invariant that slots beyond
      -- the size of the storage never contain garbage (see the note on
      -- 'undefinedVersion').
      setPrimArray vs newSize (vs0size - newSize) undefinedVersion
      pure vs
    else pure vs1
  es0size <- getSizeofSmallMutableArray es0
  es1size <- getSizeofSmallMutableArray es1
  es <- if es0size > es1size
    then do
      es <- newSmallArray es0size undefinedEffect
      copySmallMutableArray es 0 es1 0 newSize
      pure es
    else pure es1
  fs0size <- getSizeofSmallMutableArray fs0
  fs1size <- getSizeofSmallMutableArray fs1
  fs <- if fs0size > fs1size
    then do
      fs <- newSmallArray fs0size undefinedRelinker
      copySmallMutableArray fs 0 fs1 0 newSize
      pure fs
    else pure fs1
  S.writeIORef env.storage $ Storage version (StorageData newSize vs es fs)

----------------------------------------
-- 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
    :: (HasCallStack => (forall es. Env es -> IO (Env es)) -> rep e -> IO (rep e))
    -> Relinker rep e

-- | A dummy 'Relinker'.
dummyRelinker :: Relinker rep e
dummyRelinker = Relinker $ \_ -> 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 :: HasCallStack => IO (Env '[])
emptyEnv = Env 0
  <$> (unsafeFreezePrimArray =<< newPrimArray 0)
  <*> (S.newIORef =<< emptyStorage)

-- | Clone the environment to use it in a different thread.
cloneEnv :: HasCallStack => Env es -> IO (Env es)
cloneEnv env = replaceStorage env =<< cloneStorage env.storage

-- | Get the current size of the environment.
sizeEnv :: Env es -> IO Int
sizeEnv (Env offset refs _) = do
  pure $ sizeofPrimArray refs - offset

-- | Access the tail of the environment.
tailEnv :: Env (e : es) -> IO (Env es)
tailEnv (Env offset refs storage) = do
  pure $ Env (offset + 1) refs storage

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

-- | Extend the environment with a new data type.
consEnv
  :: HasCallStack
  => EffectRep (DispatchOf e) e
  -- ^ The representation of the effect.
  -> Relinker (EffectRep (DispatchOf e)) e
  -> Env es
  -> IO (Env (e : es))
consEnv e f (Env offset refs0 storage) = do
  let size = sizeofPrimArray refs0 - offset
  mrefs <- newPrimArray (size + 1)
  copyPrimArray mrefs 1 refs0 offset size
  ref <- insertEffect storage e f
  writePrimArray mrefs 0 ref
  refs <- unsafeFreezePrimArray mrefs
  pure $ Env 0 refs storage

-- | Shrink the environment by one data type.
--
-- The environment needs to come from 'consEnv', i.e. the intended usage is
-- @bracket (consEnv e f env) unconsEnv@.
--
-- /Note:/ after calling this function @e@ from the input environment is no
-- longer usable.
unconsEnv :: HasCallStack => Env (e : es) -> IO ()
unconsEnv (Env offset refs storage) = do
  when (offset /= 0) $ do
    error $ "offset (" ++ show offset ++ ") /= 0"
  deleteEffect storage (indexPrimArray refs 0)

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

-- | Replace a specific effect in the stack with a new value.
--
-- /Note:/ unlike in 'putEnv' the value is not changed in place, so only the new
-- environment will see it.
replaceEnv
  :: forall e es. (HasCallStack, e :> es)
  => EffectRep (DispatchOf e) e
  -- ^ The representation of the effect.
  -> Relinker (EffectRep (DispatchOf e)) e
  -> Env es
  -> IO (Env es)
replaceEnv e f (Env offset refs0 storage) = do
  let size = sizeofPrimArray refs0 - offset
  mrefs <- newPrimArray size
  copyPrimArray mrefs 0 refs0 offset size
  ref <- insertEffect storage e f
  writePrimArray mrefs (reifyIndex @e @es) ref
  refs <- unsafeFreezePrimArray mrefs
  pure $ Env 0 refs storage

-- | Remove a reference to the replaced effect.
--
-- The environment needs to come from 'replaceEnv', i.e. the intended usage is
-- @bracket (replaceEnv e f env) unreplaceEnv@.
--
-- /Note:/ after calling this function the input environment is no longer
-- usable.
unreplaceEnv :: forall e es. (HasCallStack, e :> es) => Env es -> IO ()
unreplaceEnv (Env offset refs storage) = do
  when (offset /= 0) $ do
    error $ "offset (" ++ show offset ++ ") /= 0"
  deleteEffect storage $ indexPrimArray refs (reifyIndex @e @es)

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

-- | Reference an existing effect from the top of the stack.
subsumeEnv :: forall e es. e :> es => Env es -> IO (Env (e : es))
subsumeEnv (Env offset refs0 storage) = do
  let size = sizeofPrimArray refs0 - offset
  mrefs <- newPrimArray (size + 1)
  copyPrimArray mrefs 1 refs0 offset size
  writePrimArray mrefs 0 $ indexPrimArray refs0 (offset + reifyIndex @e @es)
  refs <- unsafeFreezePrimArray mrefs
  pure $ Env 0 refs storage

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

-- | Construct an environment containing a permutation (with possible
-- duplicates) of a subset of effects from the input environment.
injectEnv :: forall subEs es. Subset subEs es => Env es -> IO (Env subEs)
injectEnv (Env offset refs0 storage) = do
  let subEs      = reifyIndices @subEs @es
      subEsSize  = length subEs
      prefixSize = prefixLength @es
      suffixSize = if subsetFullyKnown @subEs @es
                   then 0
                   else sizeofPrimArray refs0 - offset - prefixSize
  mrefs <- newPrimArray (subEsSize + suffixSize)
  copyPrimArray mrefs subEsSize refs0 (offset + prefixSize) suffixSize
  let writeRefs i = \case
        []       -> pure ()
        (x : xs) -> do
          writePrimArray mrefs i $ indexPrimArray refs0 (offset + x)
          writeRefs (i + 1) xs
  writeRefs 0 subEs
  refs <- unsafeFreezePrimArray mrefs
  pure $ Env 0 refs storage
{-# NOINLINE injectEnv #-}

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

-- | Extract a specific data type from the environment.
getEnv
  :: forall e es. (HasCallStack, e :> es)
  => Env es -- ^ The environment.
  -> IO (EffectRep (DispatchOf e) e)
getEnv env = do
  (i, es) <- getLocation @e env
  fromAnyEffect <$> readSmallArray es i

-- | Replace the data type in the environment with a new value (in place).
putEnv
  :: forall e es. (HasCallStack, e :> es)
  => Env es -- ^ The environment.
  -> EffectRep (DispatchOf e) e
  -> IO ()
putEnv env e = do
  (i, es) <- getLocation @e env
  writeSmallArray' es i (toAnyEffect e)

-- | Modify the data type in the environment and return a value (in place).
stateEnv
  :: forall e es a. (HasCallStack, e :> es)
  => Env es -- ^ The environment.
  -> (EffectRep (DispatchOf e) e -> (a, EffectRep (DispatchOf e) e))
  -> IO a
stateEnv env f = do
  (i, es) <- getLocation @e env
  (a, e) <- f . fromAnyEffect <$> readSmallArray es i
  writeSmallArray' es i (toAnyEffect e)
  pure a

-- | Modify the data type in the environment (in place).
modifyEnv
  :: forall e es. (HasCallStack, e :> es)
  => Env es -- ^ The environment.
  -> (EffectRep (DispatchOf e) e -> (EffectRep (DispatchOf e) e))
  -> IO ()
modifyEnv env f = do
  (i, es) <- getLocation @e env
  e <- f . fromAnyEffect <$> readSmallArray es i
  writeSmallArray' es i (toAnyEffect e)

-- | Determine location of the effect in the environment.
getLocation
  :: forall e es. (HasCallStack, e :> es)
  => Env es
  -> IO (Int, SmallMutableArray RealWorld AnyEffect)
getLocation (Env offset refs storage) = do
  Storage _ (StorageData _ vs es _) <- S.readIORef storage
  storageVersion <- readPrimArray vs 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.
  when (version /= storageVersion) $ do
    error $ "version (" ++ show version ++ ") /= storageVersion ("
         ++ show storageVersion ++ ")\n"
         ++ "If you're attempting to run an unlifting function outside "
         ++ "of the scope of effects it captures, have a look at "
         ++ "UnliftStrategy (SeqForkUnlift)."
  pure (ref, es)
  where
    Ref ref version = indexPrimArray refs (offset + reifyIndex @e @es)

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

-- | Create an empty storage.
emptyStorage :: HasCallStack => IO Storage
emptyStorage = Storage initialVersion <$> storageData
  where
    storageData = StorageData 0
      <$> newPrimArray 0
      <*> newSmallArray 0 undefinedEffect
      <*> newSmallArray 0 undefinedRelinker

-- | Insert an effect into the storage and return its reference.
insertEffect
  :: HasCallStack
  => S.IORef Storage
  -> EffectRep (DispatchOf e) e
  -- ^ The representation of the effect.
  -> Relinker (EffectRep (DispatchOf e)) e
  -> IO Ref
insertEffect storage e f = do
  Storage version (StorageData size vs0 es0 fs0) <- S.readIORef storage
  len0 <- getSizeofSmallMutableArray es0
  case size `compare` len0 of
    GT -> error $ "size (" ++ show size ++ ") > len0 (" ++ show len0 ++ ")"
    LT -> do
      writePrimArray   vs0 size version
      writeSmallArray' es0 size (toAnyEffect e)
      writeSmallArray' fs0 size (toAnyRelinker f)
      S.writeIORef storage $
        Storage (bumpVersion version) (StorageData (size + 1) vs0 es0 fs0)
      pure $ Ref size version
    EQ -> do
      let len = growCapacity len0
      vs <- newPrimArray len
      es <- newSmallArray len undefinedEffect
      fs <- newSmallArray len undefinedRelinker
      -- Fill the unused part of the versions array with 'undefinedVersion' to
      -- maintain the invariant that slots beyond the size of the storage never
      -- contain garbage (see the note on 'undefinedVersion').
      setPrimArray vs size (len - size) undefinedVersion
      copyMutablePrimArray  vs 0 vs0 0 size
      copySmallMutableArray es 0 es0 0 size
      copySmallMutableArray fs 0 fs0 0 size
      writePrimArray   vs size version
      writeSmallArray' es size (toAnyEffect e)
      writeSmallArray' fs size (toAnyRelinker f)
      S.writeIORef storage $
        Storage (bumpVersion version) (StorageData (size + 1) vs es fs)
      pure $ Ref size version
{-# NOINLINE insertEffect #-}

-- | Given a reference to an effect from the top of the stack, delete it from
-- the storage.
deleteEffect :: HasCallStack => S.IORef Storage -> Ref -> IO ()
deleteEffect storage (Ref ref version) = do
  Storage currentVersion (StorageData size vs es fs) <- S.readIORef storage
  when (ref /= size - 1) $ do
    error $ "ref (" ++ show ref ++ ") /= size - 1 (" ++ show (size - 1) ++ ")"
  storageVersion <- readPrimArray vs ref
  when (version /= storageVersion) $ do
    error $ "version (" ++ show version ++ ") /= storageVersion ("
         ++ show storageVersion ++ ")\n"
  writePrimArray  vs ref undefinedVersion
  writeSmallArray es ref undefinedEffect
  writeSmallArray fs ref undefinedRelinker
  S.writeIORef storage $ Storage currentVersion (StorageData (size - 1) vs es fs)
{-# NOINLINE deleteEffect #-}

-- | Relink the environment to use the new storage.
relinkEnv :: S.IORef Storage -> Env es -> IO (Env es)
relinkEnv storage (Env offset refs _) = pure $ Env offset refs storage

-- | Version of an unused slot.
--
-- /Note:/ slots of the versions array beyond the current size of the storage
-- always contain 'undefinedVersion', so that out of date references to them
-- reliably fail the version check in 'getLocation'. This invariant is
-- maintained by 'insertEffect', 'deleteEffect' and 'restoreStorageData'.
undefinedVersion :: Version
undefinedVersion = Version 0

initialVersion :: Version
initialVersion = Version 1

bumpVersion :: Version -> Version
bumpVersion (Version n) = Version (n + 1)

undefinedEffect :: HasCallStack => AnyEffect
undefinedEffect = toAnyEffect . errorWithoutStackTrace $ unlines
  [ "Undefined effect"
  , "Created at: " ++ prettyCallStack callStack
  ]

undefinedRelinker :: HasCallStack => AnyRelinker
undefinedRelinker = toAnyRelinker $ Relinker $ \_ _ -> do
  errorWithoutStackTrace $ unlines
    [ "Undefined relinker"
    , "Created at: " ++ prettyCallStack creationCallStack
    , "Called at: " ++ prettyCallStack callStack
    ]
  where
    creationCallStack = callStack

-- | A strict version of 'writeSmallArray'.
writeSmallArray' :: SmallMutableArray RealWorld a -> Int -> a -> IO ()
writeSmallArray' arr i a = a `seq` writeSmallArray arr i a