{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_HADDOCK not-home #-}
-- | Implementation of sequential and concurrent unlifts.
--
-- This module is intended for internal use only, and may change without warning
-- in subsequent releases.
module Effectful.Internal.Unlift
  ( -- * Unlifting strategies
    UnliftStrategy(..)
  , Persistence(..)
  , Limit(..)

    -- * Unlifting functions
  , seqUnlift
  , concUnlift
  ) where

import Control.Concurrent
import Control.Monad
import GHC.Conc.Sync (ThreadId(..))
import GHC.Exts (mkWeak#, mkWeakNoFinalizer#)
import GHC.Generics (Generic)
import GHC.IO (IO(..))
import GHC.Stack (HasCallStack)
import GHC.Weak (Weak(..))
import System.Mem.Weak (deRefWeak)
import qualified Data.IntMap.Strict as IM

import Effectful.Internal.Env
import Effectful.Internal.Utils

----------------------------------------
-- Unlift strategies

-- | The strategy to use when unlifting 'Effectful.Eff' computations via
-- 'Effectful.withEffToIO' or the 'Effectful.Dispatch.Dynamic.localUnlift'
-- family.
data UnliftStrategy
  = SeqUnlift
  -- ^ The sequential strategy is the fastest and a default setting for
  -- t'Effectful.IOE'. Any attempt of calling the unlifting function in threads
  -- distinct from its creator will result in a runtime error.
  | ConcUnlift !Persistence !Limit
  -- ^ The concurrent strategy makes it possible for the unlifting function to
  -- be called in threads distinct from its creator. See 'Persistence' and
  -- 'Limit' settings for more information.
  deriving (UnliftStrategy -> UnliftStrategy -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: UnliftStrategy -> UnliftStrategy -> Bool
$c/= :: UnliftStrategy -> UnliftStrategy -> Bool
== :: UnliftStrategy -> UnliftStrategy -> Bool
$c== :: UnliftStrategy -> UnliftStrategy -> Bool
Eq, forall x. Rep UnliftStrategy x -> UnliftStrategy
forall x. UnliftStrategy -> Rep UnliftStrategy x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep UnliftStrategy x -> UnliftStrategy
$cfrom :: forall x. UnliftStrategy -> Rep UnliftStrategy x
Generic, Eq UnliftStrategy
UnliftStrategy -> UnliftStrategy -> Bool
UnliftStrategy -> UnliftStrategy -> Ordering
UnliftStrategy -> UnliftStrategy -> UnliftStrategy
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: UnliftStrategy -> UnliftStrategy -> UnliftStrategy
$cmin :: UnliftStrategy -> UnliftStrategy -> UnliftStrategy
max :: UnliftStrategy -> UnliftStrategy -> UnliftStrategy
$cmax :: UnliftStrategy -> UnliftStrategy -> UnliftStrategy
>= :: UnliftStrategy -> UnliftStrategy -> Bool
$c>= :: UnliftStrategy -> UnliftStrategy -> Bool
> :: UnliftStrategy -> UnliftStrategy -> Bool
$c> :: UnliftStrategy -> UnliftStrategy -> Bool
<= :: UnliftStrategy -> UnliftStrategy -> Bool
$c<= :: UnliftStrategy -> UnliftStrategy -> Bool
< :: UnliftStrategy -> UnliftStrategy -> Bool
$c< :: UnliftStrategy -> UnliftStrategy -> Bool
compare :: UnliftStrategy -> UnliftStrategy -> Ordering
$ccompare :: UnliftStrategy -> UnliftStrategy -> Ordering
Ord, Int -> UnliftStrategy -> ShowS
[UnliftStrategy] -> ShowS
UnliftStrategy -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [UnliftStrategy] -> ShowS
$cshowList :: [UnliftStrategy] -> ShowS
show :: UnliftStrategy -> String
$cshow :: UnliftStrategy -> String
showsPrec :: Int -> UnliftStrategy -> ShowS
$cshowsPrec :: Int -> UnliftStrategy -> ShowS
Show)

-- | Persistence setting for the 'ConcUnlift' strategy.
--
-- Different functions require different persistence strategies. Examples:
--
-- - Lifting 'pooledMapConcurrentlyN' from the @unliftio@ library requires the
--   'Ephemeral' strategy as we don't want jobs to share environment changes
--   made by previous jobs run in the same worker thread.
--
-- - Lifting 'Control.Concurrent.forkIOWithUnmask' requires the 'Persistent'
--   strategy, otherwise the unmasking function would start with a fresh
--   environment each time it's called.
data Persistence
  = Ephemeral
  -- ^ Don't persist the environment between calls to the unlifting function in
  -- threads distinct from its creator.
  | Persistent
  -- ^ Persist the environment between calls to the unlifting function within a
  -- particular thread.
  deriving (Persistence -> Persistence -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Persistence -> Persistence -> Bool
$c/= :: Persistence -> Persistence -> Bool
== :: Persistence -> Persistence -> Bool
$c== :: Persistence -> Persistence -> Bool
Eq, forall x. Rep Persistence x -> Persistence
forall x. Persistence -> Rep Persistence x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Persistence x -> Persistence
$cfrom :: forall x. Persistence -> Rep Persistence x
Generic, Eq Persistence
Persistence -> Persistence -> Bool
Persistence -> Persistence -> Ordering
Persistence -> Persistence -> Persistence
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Persistence -> Persistence -> Persistence
$cmin :: Persistence -> Persistence -> Persistence
max :: Persistence -> Persistence -> Persistence
$cmax :: Persistence -> Persistence -> Persistence
>= :: Persistence -> Persistence -> Bool
$c>= :: Persistence -> Persistence -> Bool
> :: Persistence -> Persistence -> Bool
$c> :: Persistence -> Persistence -> Bool
<= :: Persistence -> Persistence -> Bool
$c<= :: Persistence -> Persistence -> Bool
< :: Persistence -> Persistence -> Bool
$c< :: Persistence -> Persistence -> Bool
compare :: Persistence -> Persistence -> Ordering
$ccompare :: Persistence -> Persistence -> Ordering
Ord, Int -> Persistence -> ShowS
[Persistence] -> ShowS
Persistence -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Persistence] -> ShowS
$cshowList :: [Persistence] -> ShowS
show :: Persistence -> String
$cshow :: Persistence -> String
showsPrec :: Int -> Persistence -> ShowS
$cshowsPrec :: Int -> Persistence -> ShowS
Show)

-- | Limit setting for the 'ConcUnlift' strategy.
data Limit
  = Limited !Int
  -- ^ Behavior dependent on the 'Persistence' setting.
  --
  -- For 'Ephemeral', it limits the amount of uses of the unlifting function in
  -- threads distinct from its creator to @N@. The unlifting function will
  -- create @N@ copies of the environment when called @N@ times and @K+1@ copies
  -- when called @K < N@ times.
  --
  -- For 'Persistent', it limits the amount of threads, distinct from the
  -- creator of the unlifting function, it can be called in to @N@. The amount
  -- of calls to the unlifting function within a particular threads is
  -- unlimited. The unlifting function will create @N@ copies of the environment
  -- when called in @N@ threads and @K+1@ copies when called in @K < N@ threads.
  | Unlimited
  -- ^ Unlimited use of the unlifting function.
  deriving (Limit -> Limit -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Limit -> Limit -> Bool
$c/= :: Limit -> Limit -> Bool
== :: Limit -> Limit -> Bool
$c== :: Limit -> Limit -> Bool
Eq, forall x. Rep Limit x -> Limit
forall x. Limit -> Rep Limit x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Limit x -> Limit
$cfrom :: forall x. Limit -> Rep Limit x
Generic, Eq Limit
Limit -> Limit -> Bool
Limit -> Limit -> Ordering
Limit -> Limit -> Limit
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Limit -> Limit -> Limit
$cmin :: Limit -> Limit -> Limit
max :: Limit -> Limit -> Limit
$cmax :: Limit -> Limit -> Limit
>= :: Limit -> Limit -> Bool
$c>= :: Limit -> Limit -> Bool
> :: Limit -> Limit -> Bool
$c> :: Limit -> Limit -> Bool
<= :: Limit -> Limit -> Bool
$c<= :: Limit -> Limit -> Bool
< :: Limit -> Limit -> Bool
$c< :: Limit -> Limit -> Bool
compare :: Limit -> Limit -> Ordering
$ccompare :: Limit -> Limit -> Ordering
Ord, Int -> Limit -> ShowS
[Limit] -> ShowS
Limit -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Limit] -> ShowS
$cshowList :: [Limit] -> ShowS
show :: Limit -> String
$cshow :: Limit -> String
showsPrec :: Int -> Limit -> ShowS
$cshowsPrec :: Int -> Limit -> ShowS
Show)

----------------------------------------
-- Unlift functions

-- | Sequential unlift.
seqUnlift
  :: HasCallStack
  => ((forall r. m r -> IO r) -> IO a)
  -> Env es
  -> (forall r. m r -> Env es -> IO r)
  -> IO a
seqUnlift :: forall (m :: Type -> Type) a (es :: [Effect]).
HasCallStack =>
((forall r. m r -> IO r) -> IO a)
-> Env es -> (forall r. m r -> Env es -> IO r) -> IO a
seqUnlift (forall r. m r -> IO r) -> IO a
k Env es
es forall r. m r -> Env es -> IO r
unEff = do
  ThreadId
tid0 <- IO ThreadId
myThreadId
  (forall r. m r -> IO r) -> IO a
k forall a b. (a -> b) -> a -> b
$ \m r
m -> do
    ThreadId
tid <- IO ThreadId
myThreadId
    if ThreadId
tid ThreadId -> ThreadId -> Bool
`eqThreadId` ThreadId
tid0
      then forall r. m r -> Env es -> IO r
unEff m r
m Env es
es
      else forall a. HasCallStack => String -> a
error
         forall a b. (a -> b) -> a -> b
$ String
"If you want to use the unlifting function to run Eff computations "
        forall a. [a] -> [a] -> [a]
++ String
"in multiple threads, have a look at UnliftStrategy (ConcUnlift)."

-- | Concurrent unlift.
concUnlift
  :: HasCallStack
  => Persistence
  -> Limit
  -> ((forall r. m r -> IO r) -> IO a)
  -> Env es
  -> (forall r. m r -> Env es -> IO r)
  -> IO a
concUnlift :: forall (m :: Type -> Type) a (es :: [Effect]).
HasCallStack =>
Persistence
-> Limit
-> ((forall r. m r -> IO r) -> IO a)
-> Env es
-> (forall r. m r -> Env es -> IO r)
-> IO a
concUnlift Persistence
Ephemeral (Limited Int
uses) (forall r. m r -> IO r) -> IO a
k =
  forall (m :: Type -> Type) a (es :: [Effect]).
HasCallStack =>
Int
-> ((forall r. m r -> IO r) -> IO a)
-> Env es
-> (forall r. m r -> Env es -> IO r)
-> IO a
ephemeralConcUnlift Int
uses (forall r. m r -> IO r) -> IO a
k
concUnlift Persistence
Ephemeral Limit
Unlimited (forall r. m r -> IO r) -> IO a
k =
  forall (m :: Type -> Type) a (es :: [Effect]).
HasCallStack =>
Int
-> ((forall r. m r -> IO r) -> IO a)
-> Env es
-> (forall r. m r -> Env es -> IO r)
-> IO a
ephemeralConcUnlift forall a. Bounded a => a
maxBound (forall r. m r -> IO r) -> IO a
k
concUnlift Persistence
Persistent (Limited Int
threads) (forall r. m r -> IO r) -> IO a
k =
  forall (m :: Type -> Type) a (es :: [Effect]).
HasCallStack =>
Bool
-> Int
-> ((forall r. m r -> IO r) -> IO a)
-> Env es
-> (forall r. m r -> Env es -> IO r)
-> IO a
persistentConcUnlift Bool
False Int
threads (forall r. m r -> IO r) -> IO a
k
concUnlift Persistence
Persistent Limit
Unlimited (forall r. m r -> IO r) -> IO a
k =
  forall (m :: Type -> Type) a (es :: [Effect]).
HasCallStack =>
Bool
-> Int
-> ((forall r. m r -> IO r) -> IO a)
-> Env es
-> (forall r. m r -> Env es -> IO r)
-> IO a
persistentConcUnlift Bool
True forall a. Bounded a => a
maxBound (forall r. m r -> IO r) -> IO a
k

----------------------------------------
-- Internal

-- | Concurrent unlift that doesn't preserve the environment between calls to
-- the unlifting function in threads other than its creator.
ephemeralConcUnlift
  :: HasCallStack
  => Int
  -- ^ Number of permitted uses of the unlift function.
  -> ((forall r. m r -> IO r) -> IO a)
  -> Env es
  -> (forall r. m r -> Env es -> IO r)
  -> IO a
ephemeralConcUnlift :: forall (m :: Type -> Type) a (es :: [Effect]).
HasCallStack =>
Int
-> ((forall r. m r -> IO r) -> IO a)
-> Env es
-> (forall r. m r -> Env es -> IO r)
-> IO a
ephemeralConcUnlift Int
uses (forall r. m r -> IO r) -> IO a
k Env es
es0 forall r. m r -> Env es -> IO r
unEff = do
  forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
unless (Int
uses forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$ do
    forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Invalid number of uses: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
uses
  ThreadId
tid0 <- IO ThreadId
myThreadId
  -- Create a copy of the environment as a template for the other threads to
  -- use. This can't be done from inside the callback as the environment might
  -- have already changed by then.
  Env es
esTemplate <- forall (es :: [Effect]). Env es -> IO (Env es)
cloneEnv Env es
es0
  MVar' Int
mvUses <- forall a. a -> IO (MVar' a)
newMVar' Int
uses
  (forall r. m r -> IO r) -> IO a
k forall a b. (a -> b) -> a -> b
$ \m r
m -> do
    Env es
es <- IO ThreadId
myThreadId forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      ThreadId
tid | ThreadId
tid0 ThreadId -> ThreadId -> Bool
`eqThreadId` ThreadId
tid -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Env es
es0
      ThreadId
_ -> forall a r. MVar' a -> (a -> IO (a, r)) -> IO r
modifyMVar' MVar' Int
mvUses forall a b. (a -> b) -> a -> b
$ \case
        Int
0 -> forall a. HasCallStack => String -> a
error
           forall a b. (a -> b) -> a -> b
$ String
"Number of permitted calls (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
uses forall a. [a] -> [a] -> [a]
++ String
") to the unlifting "
          forall a. [a] -> [a] -> [a]
++ String
"function in other threads was exceeded. Please increase the limit "
          forall a. [a] -> [a] -> [a]
++ String
"or use the unlimited variant."
        Int
1 -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int
0, Env es
esTemplate)
        Int
n -> do
          Env es
es <- forall (es :: [Effect]). Env es -> IO (Env es)
cloneEnv Env es
esTemplate
          forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int
n forall a. Num a => a -> a -> a
- Int
1, Env es
es)
    forall r. m r -> Env es -> IO r
unEff m r
m Env es
es

-- | Concurrent unlift that preserves the environment between calls to the
-- unlifting function within a particular thread.
persistentConcUnlift
  :: HasCallStack
  => Bool
  -> Int
  -- ^ Number of threads that are allowed to use the unlift function.
  -> ((forall r. m r -> IO r) -> IO a)
  -> Env es
  -> (forall r. m r -> Env es -> IO r)
  -> IO a
persistentConcUnlift :: forall (m :: Type -> Type) a (es :: [Effect]).
HasCallStack =>
Bool
-> Int
-> ((forall r. m r -> IO r) -> IO a)
-> Env es
-> (forall r. m r -> Env es -> IO r)
-> IO a
persistentConcUnlift Bool
cleanUp Int
threads (forall r. m r -> IO r) -> IO a
k Env es
es0 forall r. m r -> Env es -> IO r
unEff = do
  forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
unless (Int
threads forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$ do
    forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Invalid number of threads: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
threads
  ThreadId
tid0 <- IO ThreadId
myThreadId
  -- Create a copy of the environment as a template for the other threads to
  -- use. This can't be done from inside the callback as the environment might
  -- have already changed by then.
  Env es
esTemplate <- forall (es :: [Effect]). Env es -> IO (Env es)
cloneEnv Env es
es0
  MVar' (ThreadEntries es)
mvEntries <- forall a. a -> IO (MVar' a)
newMVar' forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
Int -> IntMap (ThreadEntry es) -> ThreadEntries es
ThreadEntries Int
threads forall a. IntMap a
IM.empty
  (forall r. m r -> IO r) -> IO a
k forall a b. (a -> b) -> a -> b
$ \m r
m -> do
    Env es
es <- IO ThreadId
myThreadId forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      ThreadId
tid | ThreadId
tid0 ThreadId -> ThreadId -> Bool
`eqThreadId` ThreadId
tid -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Env es
es0
      ThreadId
tid -> forall a r. MVar' a -> (a -> IO (a, r)) -> IO r
modifyMVar' MVar' (ThreadEntries es)
mvEntries forall a b. (a -> b) -> a -> b
$ \ThreadEntries es
te -> do
        let wkTid :: Int
wkTid = ThreadId -> Int
weakThreadId ThreadId
tid
        (Maybe (Env es)
mes, EntryId
i) <- case Int
wkTid forall a. Int -> IntMap a -> Maybe a
`IM.lookup` forall (es :: [Effect]).
ThreadEntries es -> IntMap (ThreadEntry es)
teEntries ThreadEntries es
te of
          Just (ThreadEntry EntryId
i ThreadData es
td) -> (, EntryId
i) forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (es :: [Effect]).
ThreadId -> ThreadData es -> IO (Maybe (Env es))
lookupEnv ThreadId
tid ThreadData es
td
          Maybe (ThreadEntry es)
Nothing                 -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (forall a. Maybe a
Nothing, EntryId
newEntryId)
        case Maybe (Env es)
mes of
          Just Env es
es -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ThreadEntries es
te, Env es
es)
          Maybe (Env es)
Nothing -> case forall (es :: [Effect]). ThreadEntries es -> Int
teCapacity ThreadEntries es
te of
            Int
0 -> forall a. HasCallStack => String -> a
error
              forall a b. (a -> b) -> a -> b
$ String
"Number of other threads (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
threads forall a. [a] -> [a] -> [a]
++ String
") permitted to "
              forall a. [a] -> [a] -> [a]
++ String
"use the unlifting function was exceeded. Please increase the "
              forall a. [a] -> [a] -> [a]
++ String
"limit or use the unlimited variant."
            Int
1 -> do
              Weak (ThreadId, Env es)
wkTidEs <- forall (es :: [Effect]).
ThreadId
-> Env es
-> Int
-> EntryId
-> MVar' (ThreadEntries es)
-> Bool
-> IO (Weak (ThreadId, Env es))
mkWeakThreadIdEnv ThreadId
tid Env es
esTemplate Int
wkTid EntryId
i MVar' (ThreadEntries es)
mvEntries Bool
cleanUp
              let newEntries :: ThreadEntries es
newEntries = ThreadEntries
                    { teCapacity :: Int
teCapacity = forall (es :: [Effect]). ThreadEntries es -> Int
teCapacity ThreadEntries es
te forall a. Num a => a -> a -> a
- Int
1
                    , teEntries :: IntMap (ThreadEntry es)
teEntries  = forall (es :: [Effect]).
Int
-> EntryId
-> Weak (ThreadId, Env es)
-> IntMap (ThreadEntry es)
-> IntMap (ThreadEntry es)
addThreadData Int
wkTid EntryId
i Weak (ThreadId, Env es)
wkTidEs forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
ThreadEntries es -> IntMap (ThreadEntry es)
teEntries ThreadEntries es
te
                    }
              forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ThreadEntries es
newEntries, Env es
esTemplate)
            Int
_ -> do
              Env es
es      <- forall (es :: [Effect]). Env es -> IO (Env es)
cloneEnv Env es
esTemplate
              Weak (ThreadId, Env es)
wkTidEs <- forall (es :: [Effect]).
ThreadId
-> Env es
-> Int
-> EntryId
-> MVar' (ThreadEntries es)
-> Bool
-> IO (Weak (ThreadId, Env es))
mkWeakThreadIdEnv ThreadId
tid Env es
es Int
wkTid EntryId
i MVar' (ThreadEntries es)
mvEntries Bool
cleanUp
              let newEntries :: ThreadEntries es
newEntries = ThreadEntries
                    { teCapacity :: Int
teCapacity = forall (es :: [Effect]). ThreadEntries es -> Int
teCapacity ThreadEntries es
te forall a. Num a => a -> a -> a
- Int
1
                    , teEntries :: IntMap (ThreadEntry es)
teEntries  = forall (es :: [Effect]).
Int
-> EntryId
-> Weak (ThreadId, Env es)
-> IntMap (ThreadEntry es)
-> IntMap (ThreadEntry es)
addThreadData Int
wkTid EntryId
i Weak (ThreadId, Env es)
wkTidEs forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
ThreadEntries es -> IntMap (ThreadEntry es)
teEntries ThreadEntries es
te
                    }
              forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ThreadEntries es
newEntries, Env es
es)
    forall r. m r -> Env es -> IO r
unEff m r
m Env es
es

----------------------------------------
-- Data types

newtype EntryId = EntryId Int
  deriving EntryId -> EntryId -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: EntryId -> EntryId -> Bool
$c/= :: EntryId -> EntryId -> Bool
== :: EntryId -> EntryId -> Bool
$c== :: EntryId -> EntryId -> Bool
Eq

newEntryId :: EntryId
newEntryId :: EntryId
newEntryId = Int -> EntryId
EntryId Int
0

nextEntryId :: EntryId -> EntryId
nextEntryId :: EntryId -> EntryId
nextEntryId (EntryId Int
i) = Int -> EntryId
EntryId (Int
i forall a. Num a => a -> a -> a
+ Int
1)

data ThreadEntries es = ThreadEntries
  { forall (es :: [Effect]). ThreadEntries es -> Int
teCapacity :: !Int
  , forall (es :: [Effect]).
ThreadEntries es -> IntMap (ThreadEntry es)
teEntries  :: !(IM.IntMap (ThreadEntry es))
  }

-- | In GHC < 9 weak thread ids are 32bit long, while ThreadIdS are 64bit long,
-- so there is potential for collisions. This is solved by keeping, for a
-- particular weak thread id, a list of ThreadIdS with unique EntryIdS.
data ThreadEntry es = ThreadEntry !EntryId !(ThreadData es)

data ThreadData es
  = ThreadData !EntryId !(Weak (ThreadId, Env es)) (ThreadData es)
  | NoThreadData

----------------------------------------
-- Weak references to threads

mkWeakThreadIdEnv
  :: ThreadId
  -> Env es
  -> Int
  -> EntryId
  -> MVar' (ThreadEntries es)
  -> Bool
  -> IO (Weak (ThreadId, Env es))
mkWeakThreadIdEnv :: forall (es :: [Effect]).
ThreadId
-> Env es
-> Int
-> EntryId
-> MVar' (ThreadEntries es)
-> Bool
-> IO (Weak (ThreadId, Env es))
mkWeakThreadIdEnv t :: ThreadId
t@(ThreadId ThreadId#
t#) Env es
es Int
wkTid EntryId
i MVar' (ThreadEntries es)
v = \case
  Bool
True -> forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s0 ->
    case mkWeak# :: forall a b c.
a
-> b
-> (State# RealWorld -> (# State# RealWorld, c #))
-> State# RealWorld
-> (# State# RealWorld, Weak# b #)
mkWeak# ThreadId#
t# (ThreadId
t, Env es
es) State# RealWorld -> (# State# RealWorld, () #)
finalizer State# RealWorld
s0 of
      (# State# RealWorld
s1, Weak# (ThreadId, Env es)
w #) -> (# State# RealWorld
s1, forall v. Weak# v -> Weak v
Weak Weak# (ThreadId, Env es)
w #)
  Bool
False -> forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s0 ->
    case mkWeakNoFinalizer# :: forall a b.
a -> b -> State# RealWorld -> (# State# RealWorld, Weak# b #)
mkWeakNoFinalizer# ThreadId#
t# (ThreadId
t, Env es
es) State# RealWorld
s0 of
      (# State# RealWorld
s1, Weak# (ThreadId, Env es)
w #) -> (# State# RealWorld
s1, forall v. Weak# v -> Weak v
Weak Weak# (ThreadId, Env es)
w #)
  where
    IO State# RealWorld -> (# State# RealWorld, () #)
finalizer = forall (es :: [Effect]).
Int -> EntryId -> MVar' (ThreadEntries es) -> IO ()
deleteThreadData Int
wkTid EntryId
i MVar' (ThreadEntries es)
v

----------------------------------------
-- Manipulation of ThreadEntries

lookupEnv :: ThreadId -> ThreadData es -> IO (Maybe (Env es))
lookupEnv :: forall (es :: [Effect]).
ThreadId -> ThreadData es -> IO (Maybe (Env es))
lookupEnv ThreadId
tid0 = \case
  ThreadData es
NoThreadData -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
  ThreadData EntryId
_ Weak (ThreadId, Env es)
wkTidEs ThreadData es
td -> forall v. Weak v -> IO (Maybe v)
deRefWeak Weak (ThreadId, Env es)
wkTidEs forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe (ThreadId, Env es)
Nothing -> forall (es :: [Effect]).
ThreadId -> ThreadData es -> IO (Maybe (Env es))
lookupEnv ThreadId
tid0 ThreadData es
td
    Just (ThreadId
tid, Env es
es)
      | ThreadId
tid0 ThreadId -> ThreadId -> Bool
`eqThreadId` ThreadId
tid -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just Env es
es
      | Bool
otherwise             -> forall (es :: [Effect]).
ThreadId -> ThreadData es -> IO (Maybe (Env es))
lookupEnv ThreadId
tid0 ThreadData es
td

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

addThreadData
  :: Int
  -> EntryId
  -> Weak (ThreadId, Env es)
  -> IM.IntMap (ThreadEntry es)
  -> IM.IntMap (ThreadEntry es)
addThreadData :: forall (es :: [Effect]).
Int
-> EntryId
-> Weak (ThreadId, Env es)
-> IntMap (ThreadEntry es)
-> IntMap (ThreadEntry es)
addThreadData Int
wkTid EntryId
i Weak (ThreadId, Env es)
w IntMap (ThreadEntry es)
teMap
  | EntryId
i forall a. Eq a => a -> a -> Bool
== EntryId
newEntryId = forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
wkTid (forall (es :: [Effect]).
EntryId -> Weak (ThreadId, Env es) -> ThreadEntry es
newThreadEntry EntryId
i Weak (ThreadId, Env es)
w) IntMap (ThreadEntry es)
teMap
  | Bool
otherwise       = forall a. (a -> a) -> Int -> IntMap a -> IntMap a
IM.adjust (forall (es :: [Effect]).
Weak (ThreadId, Env es) -> ThreadEntry es -> ThreadEntry es
consThreadData Weak (ThreadId, Env es)
w) Int
wkTid IntMap (ThreadEntry es)
teMap

newThreadEntry :: EntryId -> Weak (ThreadId, Env es) -> ThreadEntry es
newThreadEntry :: forall (es :: [Effect]).
EntryId -> Weak (ThreadId, Env es) -> ThreadEntry es
newThreadEntry EntryId
i Weak (ThreadId, Env es)
w = forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadEntry es
ThreadEntry (EntryId -> EntryId
nextEntryId EntryId
i) forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
EntryId
-> Weak (ThreadId, Env es) -> ThreadData es -> ThreadData es
ThreadData EntryId
i Weak (ThreadId, Env es)
w forall (es :: [Effect]). ThreadData es
NoThreadData

consThreadData :: Weak (ThreadId, Env es) -> ThreadEntry es -> ThreadEntry es
consThreadData :: forall (es :: [Effect]).
Weak (ThreadId, Env es) -> ThreadEntry es -> ThreadEntry es
consThreadData Weak (ThreadId, Env es)
w (ThreadEntry EntryId
i ThreadData es
td) =
  forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadEntry es
ThreadEntry (EntryId -> EntryId
nextEntryId EntryId
i) forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
EntryId
-> Weak (ThreadId, Env es) -> ThreadData es -> ThreadData es
ThreadData EntryId
i Weak (ThreadId, Env es)
w ThreadData es
td

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

deleteThreadData :: Int -> EntryId -> MVar' (ThreadEntries es) -> IO ()
deleteThreadData :: forall (es :: [Effect]).
Int -> EntryId -> MVar' (ThreadEntries es) -> IO ()
deleteThreadData Int
wkTid EntryId
i MVar' (ThreadEntries es)
v = forall a. MVar' a -> (a -> IO a) -> IO ()
modifyMVar_' MVar' (ThreadEntries es)
v forall a b. (a -> b) -> a -> b
$ \ThreadEntries es
te -> do
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ThreadEntries
    { teCapacity :: Int
teCapacity = case forall (es :: [Effect]). ThreadEntries es -> Int
teCapacity ThreadEntries es
te of
        -- If the template copy of the environment hasn't been consumed
        -- yet, the capacity can be restored.
        Int
0 -> Int
0
        Int
n -> Int
n forall a. Num a => a -> a -> a
+ Int
1
    , teEntries :: IntMap (ThreadEntry es)
teEntries = forall a. (a -> Maybe a) -> Int -> IntMap a -> IntMap a
IM.update (forall (es :: [Effect]).
EntryId -> ThreadEntry es -> Maybe (ThreadEntry es)
cleanThreadEntry EntryId
i) Int
wkTid forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
ThreadEntries es -> IntMap (ThreadEntry es)
teEntries ThreadEntries es
te
    }

cleanThreadEntry :: EntryId -> ThreadEntry es -> Maybe (ThreadEntry es)
cleanThreadEntry :: forall (es :: [Effect]).
EntryId -> ThreadEntry es -> Maybe (ThreadEntry es)
cleanThreadEntry EntryId
i0 (ThreadEntry EntryId
i ThreadData es
td0) = case forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadData es
cleanThreadData EntryId
i0 ThreadData es
td0 of
  ThreadData es
NoThreadData -> forall a. Maybe a
Nothing
  ThreadData es
td           -> forall a. a -> Maybe a
Just (forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadEntry es
ThreadEntry EntryId
i ThreadData es
td)

cleanThreadData :: EntryId -> ThreadData es -> ThreadData es
cleanThreadData :: forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadData es
cleanThreadData EntryId
i0 = \case
  ThreadData es
NoThreadData -> forall (es :: [Effect]). ThreadData es
NoThreadData
  ThreadData EntryId
i Weak (ThreadId, Env es)
w ThreadData es
td
    | EntryId
i0 forall a. Eq a => a -> a -> Bool
== EntryId
i   -> ThreadData es
td
    | Bool
otherwise -> forall (es :: [Effect]).
EntryId
-> Weak (ThreadId, Env es) -> ThreadData es -> ThreadData es
ThreadData EntryId
i Weak (ThreadId, Env es)
w (forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadData es
cleanThreadData EntryId
i0 ThreadData es
td)