{-# LANGUAGE MagicHash #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# 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
  , ephemeralConcUnlift
  , persistentConcUnlift
  ) where

import Control.Concurrent
import Control.Concurrent.MVar.Strict
import Control.Monad
import Data.Coerce
import Data.IntMap.Strict qualified as IM
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 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.
  | SeqForkUnlift
  -- ^ Like 'SeqUnlift', but all unlifted actions will be executed in a cloned
  -- environment.
  --
  -- The main consequence is that thread local state is forked at the point of
  -- creation of the unlifting function and its modifications in unlifted
  -- actions will not affect the main thread of execution (and vice versa):
  --
  -- >>> import Effectful
  -- >>> import Effectful.State.Dynamic
  -- >>> :{
  --  action :: (IOE :> es, State Int :> es) => Eff es ()
  --  action = do
  --    modify @Int (+1)
  --    withEffToIO SeqForkUnlift $ \unlift -> unlift $ modify @Int (+2)
  --    modify @Int (+4)
  -- :}
  --
  -- >>> runEff . execStateLocal @Int 0 $ action
  -- 5
  --
  -- >>> runEff . execStateShared @Int 0 $ action
  -- 7
  --
  -- Because of this it's possible to safely use the unlifting function outside
  -- of the scope of effects it captures, e.g. by creating an @IO@ action that
  -- executes effectful operations and running it later:
  --
  -- >>> :{
  --   delayed :: UnliftStrategy -> IO (IO String)
  --   delayed strategy = runEff . evalStateLocal "Hey" $ do
  --     r <- withEffToIO strategy $ \unlift -> pure $ unlift get
  --     modify (++ "!!!")
  --     pure r
  -- :}
  --
  -- This doesn't work with the 'SeqUnlift' strategy because when the returned
  -- action runs, @State@ is no longer in scope:
  --
  -- >>> join $ delayed SeqUnlift
  -- *** Exception: version (...) /= storageVersion (0)
  -- ...
  --
  -- However, it does with the 'SeqForkUnlift' strategy:
  --
  -- >>> join $ delayed SeqForkUnlift
  -- "Hey"
  --
  | 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 stock (UnliftStrategy -> UnliftStrategy -> Bool
(UnliftStrategy -> UnliftStrategy -> Bool)
-> (UnliftStrategy -> UnliftStrategy -> Bool) -> Eq UnliftStrategy
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: UnliftStrategy -> UnliftStrategy -> Bool
== :: UnliftStrategy -> UnliftStrategy -> Bool
$c/= :: UnliftStrategy -> UnliftStrategy -> Bool
/= :: UnliftStrategy -> UnliftStrategy -> Bool
Eq, (forall x. UnliftStrategy -> Rep UnliftStrategy x)
-> (forall x. Rep UnliftStrategy x -> UnliftStrategy)
-> Generic UnliftStrategy
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
$cfrom :: forall x. UnliftStrategy -> Rep UnliftStrategy x
from :: forall x. UnliftStrategy -> Rep UnliftStrategy x
$cto :: forall x. Rep UnliftStrategy x -> UnliftStrategy
to :: forall x. Rep UnliftStrategy x -> UnliftStrategy
Generic, Eq UnliftStrategy
Eq UnliftStrategy =>
(UnliftStrategy -> UnliftStrategy -> Ordering)
-> (UnliftStrategy -> UnliftStrategy -> Bool)
-> (UnliftStrategy -> UnliftStrategy -> Bool)
-> (UnliftStrategy -> UnliftStrategy -> Bool)
-> (UnliftStrategy -> UnliftStrategy -> Bool)
-> (UnliftStrategy -> UnliftStrategy -> UnliftStrategy)
-> (UnliftStrategy -> UnliftStrategy -> UnliftStrategy)
-> Ord 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
$ccompare :: UnliftStrategy -> UnliftStrategy -> Ordering
compare :: UnliftStrategy -> UnliftStrategy -> Ordering
$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
>= :: UnliftStrategy -> UnliftStrategy -> Bool
$cmax :: UnliftStrategy -> UnliftStrategy -> UnliftStrategy
max :: UnliftStrategy -> UnliftStrategy -> UnliftStrategy
$cmin :: UnliftStrategy -> UnliftStrategy -> UnliftStrategy
min :: UnliftStrategy -> UnliftStrategy -> UnliftStrategy
Ord, Int -> UnliftStrategy -> ShowS
[UnliftStrategy] -> ShowS
UnliftStrategy -> String
(Int -> UnliftStrategy -> ShowS)
-> (UnliftStrategy -> String)
-> ([UnliftStrategy] -> ShowS)
-> Show UnliftStrategy
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> UnliftStrategy -> ShowS
showsPrec :: Int -> UnliftStrategy -> ShowS
$cshow :: UnliftStrategy -> String
show :: UnliftStrategy -> String
$cshowList :: [UnliftStrategy] -> ShowS
showList :: [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 stock (Persistence -> Persistence -> Bool
(Persistence -> Persistence -> Bool)
-> (Persistence -> Persistence -> Bool) -> Eq Persistence
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Persistence -> Persistence -> Bool
== :: Persistence -> Persistence -> Bool
$c/= :: Persistence -> Persistence -> Bool
/= :: Persistence -> Persistence -> Bool
Eq, (forall x. Persistence -> Rep Persistence x)
-> (forall x. Rep Persistence x -> Persistence)
-> Generic Persistence
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
$cfrom :: forall x. Persistence -> Rep Persistence x
from :: forall x. Persistence -> Rep Persistence x
$cto :: forall x. Rep Persistence x -> Persistence
to :: forall x. Rep Persistence x -> Persistence
Generic, Eq Persistence
Eq Persistence =>
(Persistence -> Persistence -> Ordering)
-> (Persistence -> Persistence -> Bool)
-> (Persistence -> Persistence -> Bool)
-> (Persistence -> Persistence -> Bool)
-> (Persistence -> Persistence -> Bool)
-> (Persistence -> Persistence -> Persistence)
-> (Persistence -> Persistence -> Persistence)
-> Ord 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
$ccompare :: Persistence -> Persistence -> Ordering
compare :: Persistence -> Persistence -> Ordering
$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
>= :: Persistence -> Persistence -> Bool
$cmax :: Persistence -> Persistence -> Persistence
max :: Persistence -> Persistence -> Persistence
$cmin :: Persistence -> Persistence -> Persistence
min :: Persistence -> Persistence -> Persistence
Ord, Int -> Persistence -> ShowS
[Persistence] -> ShowS
Persistence -> String
(Int -> Persistence -> ShowS)
-> (Persistence -> String)
-> ([Persistence] -> ShowS)
-> Show Persistence
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Persistence -> ShowS
showsPrec :: Int -> Persistence -> ShowS
$cshow :: Persistence -> String
show :: Persistence -> String
$cshowList :: [Persistence] -> ShowS
showList :: [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 stock (Limit -> Limit -> Bool
(Limit -> Limit -> Bool) -> (Limit -> Limit -> Bool) -> Eq Limit
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Limit -> Limit -> Bool
== :: Limit -> Limit -> Bool
$c/= :: Limit -> Limit -> Bool
/= :: Limit -> Limit -> Bool
Eq, (forall x. Limit -> Rep Limit x)
-> (forall x. Rep Limit x -> Limit) -> Generic Limit
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
$cfrom :: forall x. Limit -> Rep Limit x
from :: forall x. Limit -> Rep Limit x
$cto :: forall x. Rep Limit x -> Limit
to :: forall x. Rep Limit x -> Limit
Generic, Eq Limit
Eq Limit =>
(Limit -> Limit -> Ordering)
-> (Limit -> Limit -> Bool)
-> (Limit -> Limit -> Bool)
-> (Limit -> Limit -> Bool)
-> (Limit -> Limit -> Bool)
-> (Limit -> Limit -> Limit)
-> (Limit -> Limit -> Limit)
-> Ord 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
$ccompare :: Limit -> Limit -> Ordering
compare :: Limit -> Limit -> Ordering
$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
>= :: Limit -> Limit -> Bool
$cmax :: Limit -> Limit -> Limit
max :: Limit -> Limit -> Limit
$cmin :: Limit -> Limit -> Limit
min :: Limit -> Limit -> Limit
Ord, Int -> Limit -> ShowS
[Limit] -> ShowS
Limit -> String
(Int -> Limit -> ShowS)
-> (Limit -> String) -> ([Limit] -> ShowS) -> Show Limit
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Limit -> ShowS
showsPrec :: Int -> Limit -> ShowS
$cshow :: Limit -> String
show :: Limit -> String
$cshowList :: [Limit] -> ShowS
showList :: [Limit] -> ShowS
Show)

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

-- | Concurrent unlift that doesn't preserve the environment between calls to
-- the unlifting function in threads other than its creator.
ephemeralConcUnlift
  :: (HasCallStack, forall r. Coercible (m r) (Env es -> IO r))
  => Env es
  -> Int
  -- ^ Number of permitted uses of the unlift function.
  -> ((forall r. m r -> IO r) -> IO a)
  -> IO a
ephemeralConcUnlift :: forall (m :: Type -> Type) (es :: [Effect]) a.
(HasCallStack, forall r. Coercible (m r) (Env es -> IO r)) =>
Env es -> Int -> ((forall r. m r -> IO r) -> IO a) -> IO a
ephemeralConcUnlift Env es
es0 Int
uses (forall r. m r -> IO r) -> IO a
k = do
  Bool -> IO () -> IO ()
forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
unless (Int
uses Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    String -> IO ()
forall a. HasCallStack => String -> a
error (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Invalid number of uses: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
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 <- Env es -> IO (Env es)
forall (es :: [Effect]). HasCallStack => Env es -> IO (Env es)
cloneEnv Env es
es0
  MVar' Int
mvUses <- Int -> IO (MVar' Int)
forall a. a -> IO (MVar' a)
newMVar' Int
uses
  (forall r. m r -> IO r) -> IO a
k ((forall r. m r -> IO r) -> IO a)
-> (forall r. m r -> IO r) -> IO a
forall a b. (a -> b) -> a -> b
$ \m r
m -> do
    Env es
es <- IO ThreadId
myThreadId IO ThreadId -> (ThreadId -> IO (Env es)) -> IO (Env es)
forall a b. IO a -> (a -> IO b) -> IO b
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 -> Env es -> IO (Env es)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Env es
es0
      ThreadId
_ -> MVar' Int -> (Int -> IO (Int, Env es)) -> IO (Env es)
forall a b. MVar' a -> (a -> IO (a, b)) -> IO b
modifyMVar' MVar' Int
mvUses ((Int -> IO (Int, Env es)) -> IO (Env es))
-> (Int -> IO (Int, Env es)) -> IO (Env es)
forall a b. (a -> b) -> a -> b
$ \case
        Int
0 -> String -> IO (Int, Env es)
forall a. HasCallStack => String -> a
error
           (String -> IO (Int, Env es)) -> String -> IO (Int, Env es)
forall a b. (a -> b) -> a -> b
$ String
"Number of permitted calls (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
uses String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") to the unlifting "
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"function in other threads was exceeded. Please increase the limit "
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"or use the unlimited variant."
        Int
1 -> (Int, Env es) -> IO (Int, Env es)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int
0, Env es
esTemplate)
        Int
n -> do
          Env es
es <- Env es -> IO (Env es)
forall (es :: [Effect]). HasCallStack => Env es -> IO (Env es)
cloneEnv Env es
esTemplate
          (Int, Env es) -> IO (Int, Env es)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Env es
es)
    m r -> Env es -> IO r
forall a b. Coercible a b => a -> b
coerce m r
m Env es
es
{-# NOINLINE ephemeralConcUnlift #-}

-- | Concurrent unlift that preserves the environment between calls to the
-- unlifting function within a particular thread.
persistentConcUnlift
  :: (HasCallStack, forall r. Coercible (m r) (Env es -> IO r))
  => Env es
  -> Bool
  -> Int
  -- ^ Number of threads that are allowed to use the unlift function.
  -> ((forall r. m r -> IO r) -> IO a)
  -> IO a
persistentConcUnlift :: forall (m :: Type -> Type) (es :: [Effect]) a.
(HasCallStack, forall r. Coercible (m r) (Env es -> IO r)) =>
Env es -> Bool -> Int -> ((forall r. m r -> IO r) -> IO a) -> IO a
persistentConcUnlift Env es
es0 Bool
cleanUp Int
threads (forall r. m r -> IO r) -> IO a
k = do
  Bool -> IO () -> IO ()
forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
unless (Int
threads Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    String -> IO ()
forall a. HasCallStack => String -> a
error (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Invalid number of threads: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
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 <- Env es -> IO (Env es)
forall (es :: [Effect]). HasCallStack => Env es -> IO (Env es)
cloneEnv Env es
es0
  MVar' (ThreadEntries es)
mvEntries <- ThreadEntries es -> IO (MVar' (ThreadEntries es))
forall a. a -> IO (MVar' a)
newMVar' (ThreadEntries es -> IO (MVar' (ThreadEntries es)))
-> ThreadEntries es -> IO (MVar' (ThreadEntries es))
forall a b. (a -> b) -> a -> b
$ Int -> IntMap (ThreadEntry es) -> ThreadEntries es
forall (es :: [Effect]).
Int -> IntMap (ThreadEntry es) -> ThreadEntries es
ThreadEntries Int
threads IntMap (ThreadEntry es)
forall a. IntMap a
IM.empty
  (forall r. m r -> IO r) -> IO a
k ((forall r. m r -> IO r) -> IO a)
-> (forall r. m r -> IO r) -> IO a
forall a b. (a -> b) -> a -> b
$ \m r
m -> do
    Env es
es <- IO ThreadId
myThreadId IO ThreadId -> (ThreadId -> IO (Env es)) -> IO (Env es)
forall a b. IO a -> (a -> IO b) -> IO b
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 -> Env es -> IO (Env es)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Env es
es0
      ThreadId
tid -> MVar' (ThreadEntries es)
-> (ThreadEntries es -> IO (ThreadEntries es, Env es))
-> IO (Env es)
forall a b. MVar' a -> (a -> IO (a, b)) -> IO b
modifyMVar' MVar' (ThreadEntries es)
mvEntries ((ThreadEntries es -> IO (ThreadEntries es, Env es))
 -> IO (Env es))
-> (ThreadEntries es -> IO (ThreadEntries es, Env es))
-> IO (Env es)
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 Int -> IntMap (ThreadEntry es) -> Maybe (ThreadEntry es)
forall a. Int -> IntMap a -> Maybe a
`IM.lookup` ThreadEntries es -> IntMap (ThreadEntry es)
forall (es :: [Effect]).
ThreadEntries es -> IntMap (ThreadEntry es)
teEntries ThreadEntries es
te of
          Just (ThreadEntry EntryId
i ThreadData es
td) -> (, EntryId
i) (Maybe (Env es) -> (Maybe (Env es), EntryId))
-> IO (Maybe (Env es)) -> IO (Maybe (Env es), EntryId)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ThreadId -> ThreadData es -> IO (Maybe (Env es))
forall (es :: [Effect]).
ThreadId -> ThreadData es -> IO (Maybe (Env es))
lookupEnv ThreadId
tid ThreadData es
td
          Maybe (ThreadEntry es)
Nothing                 -> (Maybe (Env es), EntryId) -> IO (Maybe (Env es), EntryId)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Maybe (Env es)
forall a. Maybe a
Nothing, EntryId
newEntryId)
        case Maybe (Env es)
mes of
          Just Env es
es -> (ThreadEntries es, Env es) -> IO (ThreadEntries es, Env es)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ThreadEntries es
te, Env es
es)
          Maybe (Env es)
Nothing -> case ThreadEntries es -> Int
forall (es :: [Effect]). ThreadEntries es -> Int
teCapacity ThreadEntries es
te of
            Int
0 -> String -> IO (ThreadEntries es, Env es)
forall a. HasCallStack => String -> a
error
              (String -> IO (ThreadEntries es, Env es))
-> String -> IO (ThreadEntries es, Env es)
forall a b. (a -> b) -> a -> b
$ String
"Number of other threads (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
threads String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") permitted to "
              String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"use the unlifting function was exceeded. Please increase the "
              String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"limit or use the unlimited variant."
            Int
1 -> do
              Weak (ThreadId, Env es)
wkTidEs <- ThreadId
-> Env es
-> Int
-> EntryId
-> MVar' (ThreadEntries es)
-> Bool
-> IO (Weak (ThreadId, Env es))
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 = ThreadEntries es -> Int
forall (es :: [Effect]). ThreadEntries es -> Int
teCapacity ThreadEntries es
te Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                    , teEntries :: IntMap (ThreadEntry es)
teEntries  = Int
-> EntryId
-> Weak (ThreadId, Env es)
-> IntMap (ThreadEntry es)
-> IntMap (ThreadEntry es)
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 (IntMap (ThreadEntry es) -> IntMap (ThreadEntry es))
-> IntMap (ThreadEntry es) -> IntMap (ThreadEntry es)
forall a b. (a -> b) -> a -> b
$ ThreadEntries es -> IntMap (ThreadEntry es)
forall (es :: [Effect]).
ThreadEntries es -> IntMap (ThreadEntry es)
teEntries ThreadEntries es
te
                    }
              (ThreadEntries es, Env es) -> IO (ThreadEntries es, Env es)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ThreadEntries es
newEntries, Env es
esTemplate)
            Int
_ -> do
              Env es
es      <- Env es -> IO (Env es)
forall (es :: [Effect]). HasCallStack => Env es -> IO (Env es)
cloneEnv Env es
esTemplate
              Weak (ThreadId, Env es)
wkTidEs <- ThreadId
-> Env es
-> Int
-> EntryId
-> MVar' (ThreadEntries es)
-> Bool
-> IO (Weak (ThreadId, Env es))
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 = ThreadEntries es -> Int
forall (es :: [Effect]). ThreadEntries es -> Int
teCapacity ThreadEntries es
te Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                    , teEntries :: IntMap (ThreadEntry es)
teEntries  = Int
-> EntryId
-> Weak (ThreadId, Env es)
-> IntMap (ThreadEntry es)
-> IntMap (ThreadEntry es)
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 (IntMap (ThreadEntry es) -> IntMap (ThreadEntry es))
-> IntMap (ThreadEntry es) -> IntMap (ThreadEntry es)
forall a b. (a -> b) -> a -> b
$ ThreadEntries es -> IntMap (ThreadEntry es)
forall (es :: [Effect]).
ThreadEntries es -> IntMap (ThreadEntry es)
teEntries ThreadEntries es
te
                    }
              (ThreadEntries es, Env es) -> IO (ThreadEntries es, Env es)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ThreadEntries es
newEntries, Env es
es)
    m r -> Env es -> IO r
forall a b. Coercible a b => a -> b
coerce m r
m Env es
es
{-# NOINLINE persistentConcUnlift #-}

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

newtype EntryId = EntryId Int
  deriving newtype EntryId -> EntryId -> Bool
(EntryId -> EntryId -> Bool)
-> (EntryId -> EntryId -> Bool) -> Eq EntryId
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: EntryId -> EntryId -> Bool
== :: EntryId -> EntryId -> Bool
$c/= :: EntryId -> EntryId -> Bool
/= :: 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 Int -> Int -> Int
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 -> (State# RealWorld
 -> (# State# RealWorld, Weak (ThreadId, Env es) #))
-> IO (Weak (ThreadId, Env es))
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld
  -> (# State# RealWorld, Weak (ThreadId, Env es) #))
 -> IO (Weak (ThreadId, Env es)))
-> (State# RealWorld
    -> (# State# RealWorld, Weak (ThreadId, Env es) #))
-> IO (Weak (ThreadId, Env es))
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s0 ->
    case ThreadId#
-> (ThreadId, Env es)
-> (State# RealWorld -> (# State# RealWorld, () #))
-> State# RealWorld
-> (# State# RealWorld, Weak# (ThreadId, Env es) #)
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, Weak# (ThreadId, Env es) -> Weak (ThreadId, Env es)
forall v. Weak# v -> Weak v
Weak Weak# (ThreadId, Env es)
w #)
  Bool
False -> (State# RealWorld
 -> (# State# RealWorld, Weak (ThreadId, Env es) #))
-> IO (Weak (ThreadId, Env es))
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld
  -> (# State# RealWorld, Weak (ThreadId, Env es) #))
 -> IO (Weak (ThreadId, Env es)))
-> (State# RealWorld
    -> (# State# RealWorld, Weak (ThreadId, Env es) #))
-> IO (Weak (ThreadId, Env es))
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s0 ->
    case ThreadId#
-> (ThreadId, Env es)
-> State# RealWorld
-> (# State# RealWorld, Weak# (ThreadId, Env es) #)
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, Weak# (ThreadId, Env es) -> Weak (ThreadId, Env es)
forall v. Weak# v -> Weak v
Weak Weak# (ThreadId, Env es)
w #)
  where
    IO State# RealWorld -> (# State# RealWorld, () #)
finalizer = Int -> EntryId -> MVar' (ThreadEntries es) -> IO ()
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 -> Maybe (Env es) -> IO (Maybe (Env es))
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Maybe (Env es)
forall a. Maybe a
Nothing
  ThreadData EntryId
_ Weak (ThreadId, Env es)
wkTidEs ThreadData es
td -> Weak (ThreadId, Env es) -> IO (Maybe (ThreadId, Env es))
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak (ThreadId, Env es)
wkTidEs IO (Maybe (ThreadId, Env es))
-> (Maybe (ThreadId, Env es) -> IO (Maybe (Env es)))
-> IO (Maybe (Env es))
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe (ThreadId, Env es)
Nothing -> ThreadId -> ThreadData es -> IO (Maybe (Env es))
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 -> Maybe (Env es) -> IO (Maybe (Env es))
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Maybe (Env es) -> IO (Maybe (Env es)))
-> Maybe (Env es) -> IO (Maybe (Env es))
forall a b. (a -> b) -> a -> b
$ Env es -> Maybe (Env es)
forall a. a -> Maybe a
Just Env es
es
      | Bool
otherwise             -> ThreadId -> ThreadData es -> IO (Maybe (Env es))
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 EntryId -> EntryId -> Bool
forall a. Eq a => a -> a -> Bool
== EntryId
newEntryId = Int
-> ThreadEntry es
-> IntMap (ThreadEntry es)
-> IntMap (ThreadEntry es)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
wkTid (EntryId -> Weak (ThreadId, Env es) -> ThreadEntry es
forall (es :: [Effect]).
EntryId -> Weak (ThreadId, Env es) -> ThreadEntry es
newThreadEntry EntryId
i Weak (ThreadId, Env es)
w) IntMap (ThreadEntry es)
teMap
  | Bool
otherwise       = (ThreadEntry es -> ThreadEntry es)
-> Int -> IntMap (ThreadEntry es) -> IntMap (ThreadEntry es)
forall a. (a -> a) -> Int -> IntMap a -> IntMap a
IM.adjust (Weak (ThreadId, Env es) -> ThreadEntry es -> ThreadEntry es
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 = EntryId -> ThreadData es -> ThreadEntry es
forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadEntry es
ThreadEntry (EntryId -> EntryId
nextEntryId EntryId
i) (ThreadData es -> ThreadEntry es)
-> ThreadData es -> ThreadEntry es
forall a b. (a -> b) -> a -> b
$ EntryId
-> Weak (ThreadId, Env es) -> ThreadData es -> ThreadData es
forall (es :: [Effect]).
EntryId
-> Weak (ThreadId, Env es) -> ThreadData es -> ThreadData es
ThreadData EntryId
i Weak (ThreadId, Env es)
w ThreadData es
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) =
  EntryId -> ThreadData es -> ThreadEntry es
forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadEntry es
ThreadEntry (EntryId -> EntryId
nextEntryId EntryId
i) (ThreadData es -> ThreadEntry es)
-> ThreadData es -> ThreadEntry es
forall a b. (a -> b) -> a -> b
$ EntryId
-> Weak (ThreadId, Env es) -> ThreadData es -> ThreadData es
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 = MVar' (ThreadEntries es)
-> (ThreadEntries es -> IO (ThreadEntries es)) -> IO ()
forall a. MVar' a -> (a -> IO a) -> IO ()
modifyMVar'_ MVar' (ThreadEntries es)
v ((ThreadEntries es -> IO (ThreadEntries es)) -> IO ())
-> (ThreadEntries es -> IO (ThreadEntries es)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ThreadEntries es
te -> do
  ThreadEntries es -> IO (ThreadEntries es)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ThreadEntries
    { teCapacity :: Int
teCapacity = case ThreadEntries es -> Int
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    , teEntries :: IntMap (ThreadEntry es)
teEntries = (ThreadEntry es -> Maybe (ThreadEntry es))
-> Int -> IntMap (ThreadEntry es) -> IntMap (ThreadEntry es)
forall a. (a -> Maybe a) -> Int -> IntMap a -> IntMap a
IM.update (EntryId -> ThreadEntry es -> Maybe (ThreadEntry es)
forall (es :: [Effect]).
EntryId -> ThreadEntry es -> Maybe (ThreadEntry es)
cleanThreadEntry EntryId
i) Int
wkTid (IntMap (ThreadEntry es) -> IntMap (ThreadEntry es))
-> IntMap (ThreadEntry es) -> IntMap (ThreadEntry es)
forall a b. (a -> b) -> a -> b
$ ThreadEntries es -> IntMap (ThreadEntry es)
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 EntryId -> ThreadData es -> ThreadData es
forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadData es
cleanThreadData EntryId
i0 ThreadData es
td0 of
  ThreadData es
NoThreadData -> Maybe (ThreadEntry es)
forall a. Maybe a
Nothing
  ThreadData es
td           -> ThreadEntry es -> Maybe (ThreadEntry es)
forall a. a -> Maybe a
Just (EntryId -> ThreadData es -> ThreadEntry es
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 -> ThreadData es
forall (es :: [Effect]). ThreadData es
NoThreadData
  ThreadData EntryId
i Weak (ThreadId, Env es)
w ThreadData es
td
    | EntryId
i0 EntryId -> EntryId -> Bool
forall a. Eq a => a -> a -> Bool
== EntryId
i   -> ThreadData es
td
    | Bool
otherwise -> EntryId
-> Weak (ThreadId, Env es) -> ThreadData es -> ThreadData es
forall (es :: [Effect]).
EntryId
-> Weak (ThreadId, Env es) -> ThreadData es -> ThreadData es
ThreadData EntryId
i Weak (ThreadId, Env es)
w (EntryId -> ThreadData es -> ThreadData es
forall (es :: [Effect]). EntryId -> ThreadData es -> ThreadData es
cleanThreadData EntryId
i0 ThreadData es
td)