@@ -1,40 +1,52 @@
1−{-# LANGUAGE AllowAmbiguousTypes #-}
1+{-# LANGUAGE MagicHash #-}
22{-# OPTIONS_HADDOCK not-home #-}
33-- | The environment for the 'Effectful.Internal.Monad.Eff' monad.
44--
55-- This module is intended for internal use only, and may change without warning
66-- in subsequent releases.
77module Effectful.Internal.Env
8− ( Env
8+ ( -- * The environment
9+ Env(..)
10+ , Forks(..)
11+ , EnvRef(..)
12+
13+ -- ** ForkId
14+ , ForkId(..)
15+ , ForkIdGen(..)
16+ , newForkIdGen
17+ , cloneForkIdGen
18+ , newForkId
19+
20+ -- ** Relinker
921 , Relinker(..)
1022 , noRelinker
1123
12− -- * Safe operations
24+ -- * Operations
1325 , emptyEnv
1426 , cloneEnv
1527 , forkEnv
1628 , sizeEnv
17− , getEnv
1829 , checkSizeEnv
1930
20− -- * Extending and shrinking
31+ -- ** Extending and shrinking
2132 , unsafeConsEnv
2233 , unsafeTailEnv
2334
24− -- * Data retrieval and update
25− , unsafePutEnv
26− , unsafeModifyEnv
27− , unsafeStateEnv
35+ -- ** Data retrieval and update
36+ , getEnv
37+ , putEnv
38+ , stateEnv
39+ , modifyEnv
2840 ) where
2941
3042import Control.Monad
3143import Control.Monad.Primitive
3244import Data.IORef
3345import Data.Primitive.SmallArray
34−import Data.Word
35−import GHC.Exts (Any)
46+import GHC.Exts (Any, SmallMutableArray#)
47+import GHC.Stack (HasCallStack)
3648import Unsafe.Coerce
37−import qualified Data.Map.Strict as M
49+import qualified Data.IntMap.Strict as IM
3850
3951import Effectful.Internal.Effect
4052
@@ -44,7 +56,7 @@ type role Env nominal
4456-- of kind 'Effect'.
4557--
4658-- Supports forking, i.e. introduction of local branches for encapsulation of
47−-- data specific to effect interpreters.
59+-- data specific to effect handlers.
4860--
4961-- Offers very good performance characteristics:
5062--
@@ -69,17 +81,17 @@ type role Env nominal
6981-- @
7082--
7183-- The point of execution is currently at (*), i.e. the mainline. Moreover,
72−-- currently:
84+-- currently (brackets represent forks):
7385--
7486-- - Mainline sees @[e00, e01, e05, e07]@.
7587--
76−-- - @e01@ is an interpreted effect, its handler sees @[e00, e02, e03]@.
88+-- - @e01@ is an interpreted effect, its handler sees @[e00][e02, e03]@.
7789--
78−-- - @e03@ is an interpreted effect, its handler sees @[e00, e02, e04]@.
90+-- - @e03@ is an interpreted effect, its handler sees @[e00][e02][e04]@.
7991--
80−-- - @e05@ is an interpreted effect, its handler sees @[e00, e01, e06]@.
92+-- - @e05@ is an interpreted effect, its handler sees @[e00, e01][e06]@.
8193--
82−-- - @e07@ is an interpreted effect, its handler sees @[e00, e01, e05, e08,
94+-- - @e07@ is an interpreted effect, its handler sees @[e00, e01, e05][e08,
8395-- e09]@.
8496--
8597-- If an operation from @e01@ is invoked, the environment in the middle of
@@ -102,45 +114,42 @@ type role Env nominal
102114-- introduced @e12@), so its local environment was temporarily extended with
103115-- it. Moreover, currently:
104116--
105−-- - Handler of @e01@ sees @[e00, e02, e03, e10]@.
117+-- - Handler of @e01@ sees @[e00][e02, e03, e10]@.
106118--
107−-- - Handler of @e10@ sees @[e00, e02, e03, e11]@.
119+-- - Handler of @e10@ sees @[e00][e02, e03][e11]@.
108120--
109−-- - Handler of @e11@ sees @[e00, e02, e03, e12]@.
121+-- - Handler of @e11@ sees @[e00][e02, e03][e12]@.
110122--
111123data Env (es :: [Effect]) = Env
112− { _forks :: Forks
113− , _globalRef :: IORef EnvRef
114− , _forkIdGen :: ForkIdGen
124+ { _forks :: !Forks
125+ , _globalRef :: !(IORef EnvRef)
126+ , _forkIdGen :: !ForkIdGen
115127 }
116128
117129-- | Local forks of the environment.
118−data Forks = Forks ForkId Int (IORef EnvRef) Forks | NoFork
130+data Forks = Forks !ForkId !Int !(IORef EnvRef) Forks | NoFork
119131
132+-- | Data held in the environment.
120133data EnvRef = EnvRef
121− { _size :: Int
122− , _data :: SmallMutableArray RealWorld Any
123− , _relinkers :: SmallMutableArray RealWorld Any
134+ { _size :: !Int
135+ , _data :: !(SmallMutableArray RealWorld Any)
136+ , _relinkers :: !(SmallMutableArray RealWorld Any)
124137 }
125138
126139----------------------------------------
127140-- ForkId
128141
129−newtype ForkId = ForkId Word64
130− deriving (Eq, Ord, Show)
131−
132−-- | Implicit 'ForkId' of the 'EnvRef' directly available from 'Env'.
133−globalFid :: ForkId
134−globalFid = ForkId 0
142+-- | Internal id of the fork.
143+newtype ForkId = ForkId { unForkId :: Int }
135144
136145-- | 'ForkId' generation.
137146newtype ForkIdGen = ForkIdGen (IORef ForkId)
138147
139−-- | Create a new thread local unique generator.
148+-- | Create a new thread local 'ForkId' generator.
140149newForkIdGen :: IO ForkIdGen
141−newForkIdGen = ForkIdGen <$> newIORef (ForkId 1)
150+newForkIdGen = ForkIdGen <$> newIORef (ForkId 0)
142151
143−-- | Clone the unique generator for use in a different thread.
152+-- | Clone the 'ForkId' generator for use in a different thread.
144153cloneForkIdGen :: ForkIdGen -> IO ForkIdGen
145154cloneForkIdGen (ForkIdGen ref) = fmap ForkIdGen . newIORef =<< readIORef ref
146155
@@ -154,12 +163,15 @@ newForkId (ForkIdGen ref) = do
154163----------------------------------------
155164-- Relinker
156165
157−newtype Relinker i (e :: Effect) where
166+-- | A function for relinking 'Env' objects stored in the handlers when cloning
167+-- the environment.
168+newtype Relinker :: (Effect -> Type) -> Effect -> Type where
158169 Relinker
159− :: ((forall es. Env es -> IO (Env es)) -> i e -> IO (i e))
160− -> Relinker i e
170+ :: ((forall es. Env es -> IO (Env es)) -> adapter e -> IO (adapter e))
171+ -> Relinker adapter e
161172
162−noRelinker :: Relinker i e
173+-- | A dummy 'Relinker' that does nothing.
174+noRelinker :: Relinker adapter e
163175noRelinker = Relinker $ \_ -> pure
164176
165177----------------------------------------
@@ -172,23 +184,27 @@ emptyEnv = Env NoFork <$> emptyEnvRef <*> newForkIdGen
172184-- | Clone the environment.
173185cloneEnv :: Env es -> IO (Env es)
174186cloneEnv (Env NoFork gref0 gen0) = do
187+ EnvRef n es0 fs0 <- readIORef gref0
188+ es <- cloneSmallMutableArray es0 0 (sizeofSmallMutableArray es0)
189+ fs <- cloneSmallMutableArray fs0 0 (sizeofSmallMutableArray fs0)
175190 gen <- cloneForkIdGen gen0
176− cache <- newIORef M.empty
177− gref <- cloneEnvRef gen cache globalFid gref0
191+ gref <- newIORef $ EnvRef n es fs
192+ store <- newIORef IM.empty
193+ relinkData gref gen store es fs n
178194 pure $ Env NoFork gref gen
179195cloneEnv (Env forks@(Forks _ baseIx lref0 _) gref0 gen0) = do
180196 EnvRef _ es0 fs0 <- readIORef gref0
181197 EnvRef n _ _ <- readIORef lref0
182198 let len = baseIx + n
183− es <- newSmallArray len (error "undefined field")
184− fs <- newSmallArray len (error "undefined field")
199+ es <- newSmallArray len undefinedData
200+ fs <- newSmallArray len undefinedData
185201 baseN <- copyForks es fs len forks
186202 copySmallMutableArray es 0 es0 0 baseN
187203 copySmallMutableArray fs 0 fs0 0 baseN
188204 gen <- cloneForkIdGen gen0
189205 gref <- newIORef $ EnvRef len es fs
190− cache <- newIORef $ M.singleton globalFid gref
191− relinkData gen cache es fs n
206+ store <- newIORef IM.empty
207+ relinkData gref gen store es fs n
192208 -- The forked environment is flattened and becomes the global one.
193209 pure $ Env NoFork gref gen
194210{-# NOINLINE cloneEnv #-}
@@ -238,33 +254,18 @@ copyForks es fs len = \case
238254 copySmallMutableArray fs baseIx fs0 0 n
239255 copyForks es fs baseIx forks
240256
241−type EnvRefCache = IORef (M.Map ForkId (IORef EnvRef))
242−
243−-- | Clone the 'EnvRef' and put it in a cache.
244−cloneEnvRef
245− :: ForkIdGen
246− -> EnvRefCache
247− -> ForkId
248− -> IORef EnvRef
249− -> IO (IORef EnvRef)
250−cloneEnvRef gen cache fid ref0 = do
251− EnvRef n es0 fs0 <- readIORef ref0
252− es <- cloneSmallMutableArray es0 0 (sizeofSmallMutableArray es0)
253− fs <- cloneSmallMutableArray fs0 0 (sizeofSmallMutableArray fs0)
254− ref <- newIORef $ EnvRef n es fs
255− modifyIORef' cache $ M.insert fid ref
256− relinkData gen cache es fs n
257− pure ref
257+type EnvRefStore = IORef (IM.IntMap (IORef EnvRef))
258258
259−-- | Relink local environments hiding in the interpreters.
259+-- | Relink local environments hiding in the handlers.
260260relinkData
261− :: ForkIdGen
262− -> EnvRefCache
261+ :: IORef EnvRef
262+ -> ForkIdGen
263+ -> EnvRefStore
263264 -> SmallMutableArray RealWorld Any
264265 -> SmallMutableArray RealWorld Any
265266 -> Int
266267 -> IO ()
267−relinkData gen cache es fs = \case
268+relinkData gref gen store es fs = \case
268269 0 -> pure ()
269270 n -> do
270271 let i = n - 1
@@ -272,29 +273,46 @@ relinkData gen cache es fs = \case
272273 readSmallArray es i
273274 >>= f relinkEnv . fromAny
274275 >>= writeSmallArray es i . toAny
275− relinkData gen cache es fs i
276+ relinkData gref gen store es fs i
276277 where
277278 relinkEnv :: Env es -> IO (Env es)
278− relinkEnv (Env forks _ _) = do
279− Just gref <- M.lookup globalFid <$> readIORef cache
280− Env <$> relinkForks gen cache forks <*> pure gref <*> pure gen
279+ relinkEnv (Env forks _ _) = Env <$> relinkForks gref gen store forks
280+ <*> pure gref
281+ <*> pure gen
281282
282−relinkForks :: ForkIdGen -> EnvRefCache -> Forks -> IO Forks
283−relinkForks gen cache = \case
283+relinkForks :: IORef EnvRef -> ForkIdGen -> EnvRefStore -> Forks -> IO Forks
284+relinkForks gref gen store = \case
284285 NoFork -> pure NoFork
285286 Forks fid baseIx lref0 forks -> do
286287 -- A specific IORef EnvRef can be held by more than one local environment
287288 -- and we need to replace all its occurences with the same, new value
288289 -- containing its clone.
289− readIORef cache >>= pure . M.lookup fid >>= \case
290+ readIORef store >>= pure . IM.lookup (unForkId fid) >>= \case
290291 Just lref -> Forks fid baseIx <$> pure lref
291− <*> relinkForks gen cache forks
292− Nothing -> Forks fid baseIx <$> cloneEnvRef gen cache fid lref0
293− <*> relinkForks gen cache forks
292+ <*> relinkForks gref gen store forks
293+ Nothing -> Forks fid baseIx <$> cloneEnvRef gref gen store fid lref0
294+ <*> relinkForks gref gen store forks
295+
296+-- | Clone the local 'EnvRef' and put it in a store.
297+cloneEnvRef
298+ :: IORef EnvRef
299+ -> ForkIdGen
300+ -> EnvRefStore
301+ -> ForkId
302+ -> IORef EnvRef
303+ -> IO (IORef EnvRef)
304+cloneEnvRef gref gen store fid lref0 = do
305+ EnvRef n es0 fs0 <- readIORef lref0
306+ es <- cloneSmallMutableArray es0 0 (sizeofSmallMutableArray es0)
307+ fs <- cloneSmallMutableArray fs0 0 (sizeofSmallMutableArray fs0)
308+ ref <- newIORef $ EnvRef n es fs
309+ modifyIORef' store $ IM.insert (unForkId fid) ref
310+ relinkData gref gen store es fs n
311+ pure ref
294312
295313----------------------------------------
296314
297−-- | Create a local fork of the environment for interpreters.
315+-- | Create a local fork of the environment for handlers.
298316forkEnv :: Env es -> IO (Env es)
299317forkEnv env@(Env NoFork gref gen) = do
300318 size <- sizeEnv env
@@ -320,12 +338,6 @@ sizeEnv (Env (Forks _ baseIx lref _) _ _) = do
320338 pure $ baseIx + n
321339{-# NOINLINE sizeEnv #-}
322340
323−-- | Extract a specific data type from the environment.
324−getEnv :: forall e es i. e :> es => Env es -> IO (i e)
325−getEnv env = do
326− (i, es) <- getLocation @e env
327− fromAny <$> readSmallArray es i
328−
329341-- | Check that the size of the environment is the same as the expected value.
330342checkSizeEnv :: Int -> Env es -> IO ()
331343checkSizeEnv k (Env NoFork ref _) = do
@@ -343,7 +355,11 @@ checkSizeEnv k (Env (Forks _ baseIx lref _) _ _) = do
343355-- Extending and shrinking
344356
345357-- | Extend the environment with a new data type (in place).
346−unsafeConsEnv :: i e -> Relinker i e -> Env es -> IO (Env (e : es))
358+--
359+-- This function is __highly unsafe__ because it renders the input 'Env'
360+-- unusable until the corresponding 'unsafeTailEnv' call is made, but it's not
361+-- checked anywhere.
362+unsafeConsEnv :: adapter e -> Relinker adapter e -> Env es -> IO (Env (e : es))
347363unsafeConsEnv e f (Env fork gref gen) = case fork of
348364 NoFork -> do
349365 extendEnvRef gref
@@ -364,10 +380,10 @@ unsafeConsEnv e f (Env fork gref gen) = case fork of
364380 writeIORef ref $! EnvRef (n + 1) es0 fs0
365381 EQ -> do
366382 let len = doubleCapacity len0
367− es <- newSmallArray len (error "undefined field")
383+ es <- newSmallArray len undefinedData
368384 copySmallMutableArray es 0 es0 0 len0
369385 e `seq` writeSmallArray es n (toAny e)
370− fs <- newSmallArray len (error "undefined field")
386+ fs <- newSmallArray len undefinedData
371387 copySmallMutableArray fs 0 fs0 0 len0
372388 f `seq` writeSmallArray fs n (toAny f)
373389 writeIORef ref $! EnvRef (n + 1) es fs
@@ -378,14 +394,13 @@ unsafeConsEnv e f (Env fork gref gen) = case fork of
378394
379395-- | Shrink the environment by one data type (in place). Makes sure the size of
380396-- the environment is as expected.
381−unsafeTailEnv :: Int -> Env (e : es) -> IO (Env es)
382−unsafeTailEnv len (Env fork gref gen) = case fork of
383− NoFork -> do
384− shrinkEnvRef len gref
385− pure $ Env NoFork gref gen
386− Forks _ baseIx lref _ -> do
387− shrinkEnvRef (len - baseIx) lref
388− pure $ Env fork gref gen
397+--
398+-- This function is __highly unsafe__ because it renders the input 'Env'
399+-- unusable, but it's not checked anywhere.
400+unsafeTailEnv :: Int -> Env (e : es) -> IO ()
401+unsafeTailEnv len (Env fork gref _) = case fork of
402+ NoFork -> shrinkEnvRef len gref
403+ Forks _ baseIx lref _ -> shrinkEnvRef (len - baseIx) lref
389404 where
390405 shrinkEnvRef :: Int -> IORef EnvRef -> IO ()
391406 shrinkEnvRef k ref = do
@@ -393,81 +408,109 @@ unsafeTailEnv len (Env fork gref gen) = case fork of
393408 if k /= n - 1
394409 then error $ "k (" ++ show k ++ ") /= n - 1 (" ++ show (n - 1) ++ ")"
395410 else do
396− writeSmallArray es k (error "undefined field")
397− writeSmallArray fs k (error "undefined field")
411+ writeSmallArray es k undefinedData
412+ writeSmallArray fs k undefinedData
398413 writeIORef ref $! EnvRef k es fs
399414{-# NOINLINE unsafeTailEnv #-}
400415
401416----------------------------------------
402417-- Data retrieval and update
403418
404−-- | Replace the data type in the environment with a new value (in place).
405−unsafePutEnv :: forall e es i. e :> es => i e -> Env es -> IO ()
406−unsafePutEnv e env = do
407− (i, es) <- getLocation @e env
408− e `seq` writeSmallArray es i (toAny e)
419+-- | Extract a specific data type from the environment.
420+getEnv
421+ :: forall e adapter es. e :> es
422+ => Env es
423+ -> IO (adapter e)
424+getEnv env = do
425+ Location i es <- getLocation (reifyIndex @e @es) env
426+ fromAny <$> readSmallArray es i
409427
410−-- | Modify the data type in the environment (in place).
411−unsafeModifyEnv :: forall e es i. e :> es => (i e -> i e) -> Env es -> IO ()
412−unsafeModifyEnv f env = do
413− (i, es) <- getLocation @e env
414− e <- f . fromAny <$> readSmallArray es i
428+-- | Replace the data type in the environment with a new value (in place).
429+putEnv
430+ :: forall e adapter es. e :> es
431+ => Env es
432+ -> adapter e
433+ -> IO ()
434+putEnv env e = do
435+ Location i es <- getLocation (reifyIndex @e @es) env
415436 e `seq` writeSmallArray es i (toAny e)
416437
417438-- | Modify the data type in the environment (in place) and return a value.
418−unsafeStateEnv :: forall e es i a. e :> es => (i e -> (a, i e)) -> Env es -> IO a
419−unsafeStateEnv f env = do
420− (i, es) <- getLocation @e env
439+stateEnv
440+ :: forall e adapter es a. e :> es
441+ => Env es
442+ -> (adapter e -> (a, adapter e))
443+ -> IO a
444+stateEnv env f = do
445+ Location i es <- getLocation (reifyIndex @e @es) env
421446 (a, e) <- f . fromAny <$> readSmallArray es i
422447 e `seq` writeSmallArray es i (toAny e)
423448 pure a
424449
450+-- | Modify the data type in the environment (in place).
451+modifyEnv
452+ :: forall e adapter es. e :> es
453+ => Env es
454+ -> (adapter e -> adapter e)
455+ -> IO ()
456+modifyEnv env f = do
457+ Location i es <- getLocation (reifyIndex @e @es) env
458+ e <- f . fromAny <$> readSmallArray es i
459+ e `seq` writeSmallArray es i (toAny e)
460+
425461----------------------------------------
426462-- Internal helpers
427463
464+undefinedData :: HasCallStack => a
465+undefinedData = error "undefined data"
466+
428467emptyEnvRef :: IO (IORef EnvRef)
429468emptyEnvRef = do
430− es <- newSmallArray 0 (error "undefined field")
431− fs <- newSmallArray 0 (error "undefined field")
469+ es <- newSmallArray 0 undefinedData
470+ fs <- newSmallArray 0 undefinedData
432471 newIORef $ EnvRef 0 es fs
433472
473+-- | Location with unboxed fields to make GHC generate Core without unnecessary
474+-- reboxing around the local 'go' function from 'getLocation'.
475+data Location = Location !Int !(SmallMutableArray RealWorld Any)
476+
434477-- | Determine location of the data type in the environment.
435478getLocation
436− :: forall e es. e :> es
437− => Env es
438− -> IO (Int, SmallMutableArray RealWorld Any)
439−getLocation (Env fork ref _) = do
440− EnvRef n es _ <- readIORef ref
479+ :: Int
480+ -> Env es
481+ -> IO Location
482+getLocation ixE (Env fork ref _) = do
483+ EnvRef n es@(SmallMutableArray es#) _ <- readIORef ref
441484 -- Optimized for the common access pattern:
442485 --
443486 -- - Most application code has no access to forks, in which case we look at
444487 -- the global EnvRef.
445488 --
446− -- - Interpreters will most likely access the newest fork for reinterpretation
447− -- of effects that are there.
489+ -- - Effect handlers will most likely access the newest fork with handler
490+ -- specific effects for reinterpretation.
448491 case fork of
449− NoFork -> pure (ix n, es)
492+ NoFork -> pure $ Location (ix n) es
450493 Forks _ baseIx lref forks -> do
451494 EnvRef ln les _ <- readIORef lref
452495 let i = ix (baseIx + ln)
453496 if i >= baseIx
454− then pure (i - baseIx, les)
455− else go es i forks
497+ then pure $ Location (i - baseIx) les
498+ else go es# i forks
456499 where
457− go :: SmallMutableArray RealWorld Any
500+ go :: SmallMutableArray# RealWorld Any
458501 -> Int
459502 -> Forks
460− -> IO (Int, SmallMutableArray RealWorld Any)
461− go es i = \case
462− NoFork -> pure (i, es)
503+ -> IO Location
504+ go es# i = \case
505+ NoFork -> pure $ Location i (SmallMutableArray es#)
463506 Forks _ baseIx lref forks -> do
464507 EnvRef _ les _ <- readIORef lref
465508 if i >= baseIx
466− then pure (i - baseIx, les)
467− else go es i forks
509+ then pure $ Location (i - baseIx) les
510+ else go es# i forks
468511
469512 ix :: Int -> Int
470− ix n = n - reifyIndex @e @es - 1
513+ ix n = n - ixE - 1
471514
472515toAny :: a -> Any
473516toAny = unsafeCoerce