@@ -1,50 +1,62 @@
1−{-# LANGUAGE AllowAmbiguousTypes #-}
1+{-# LANGUAGE MagicHash #-}
22{-# OPTIONS_HADDOCK not-home #-}
3−-- | The environment for 'Effectful.Internal.Monad.Eff'.
3+-- | 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
4153type role Env nominal
4254
43−-- | A mutable, strict (elements are kept in weak head normal form), extensible
44−-- record indexed by effect data types.
55+-- | A strict (WHNF), thread local, mutable, extensible record indexed by types
56+-- 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--
@@ -54,41 +66,90 @@ type role Env nominal
5466--
5567-- - Indexing via '(:>)': /O(forks)/, usually /O(1)/ (amortized).
5668--
57−-- - Modification of a specific element: /O(forks)/, usually /O(1)/.
69+-- - Modification of a specific element: /O(1)/.
70+--
71+-- Here's an example of how the environment might look:
72+--
73+-- @
74+-- e00 - e01 - e05 - e07 (*)
75+-- | | |
76+-- | e06 e08 - e09
77+-- |
78+-- e02 - e03
79+-- |
80+-- e04
81+-- @
82+--
83+-- The point of execution is currently at (*), i.e. the mainline. Moreover,
84+-- currently (brackets represent forks):
85+--
86+-- - Mainline sees @[e00, e01, e05, e07]@.
87+--
88+-- - @e01@ is an interpreted effect, its handler sees @[e00][e02, e03]@.
89+--
90+-- - @e03@ is an interpreted effect, its handler sees @[e00][e02][e04]@.
91+--
92+-- - @e05@ is an interpreted effect, its handler sees @[e00, e01][e06]@.
93+--
94+-- - @e07@ is an interpreted effect, its handler sees @[e00, e01, e05][e08,
95+-- e09]@.
96+--
97+-- If an operation from @e01@ is invoked, the environment in the middle of
98+-- handling it might look like this:
99+--
100+-- @
101+-- e00 - e01 - e05 - e07
102+-- | | |
103+-- | e06 e08 - e09
104+-- |
105+-- e02 - e03 - e10 (*)
106+-- | |
107+-- e04 e11
108+-- |
109+-- e12
110+-- @
111+--
112+-- The point of execution is at (*), i.e. inside the handler of @e01@. The
113+-- handler needed to introduce @e10@ (which introduced @e11@, which in turn
114+-- introduced @e12@), so its local environment was temporarily extended with
115+-- it. Moreover, currently:
116+--
117+-- - Handler of @e01@ sees @[e00][e02, e03, e10]@.
118+--
119+-- - Handler of @e10@ sees @[e00][e02, e03][e11]@.
120+--
121+-- - Handler of @e11@ sees @[e00][e02, e03][e12]@.
58122--
59123data Env (es :: [Effect]) = Env
60− { _forks :: Forks
61− , _globalRef :: IORef EnvRef
62− , _forkIdGen :: ForkIdGen
124+ { _forks :: !Forks
125+ , _globalRef :: !(IORef EnvRef)
126+ , _forkIdGen :: !ForkIdGen
63127 }
64128
65129-- | Local forks of the environment.
66−data Forks = Forks ForkId Int (IORef EnvRef) Forks | NoFork
130+data Forks = Forks !ForkId !Int !(IORef EnvRef) Forks | NoFork
67131
132+-- | Data held in the environment.
68133data EnvRef = EnvRef
69− { _capacity :: Int
70− , _data :: SmallMutableArray RealWorld Any
71− , _relinkers :: SmallMutableArray RealWorld Any
134+ { _size :: !Int
135+ , _data :: !(SmallMutableArray RealWorld Any)
136+ , _relinkers :: !(SmallMutableArray RealWorld Any)
72137 }
73138
74139----------------------------------------
75140-- ForkId
76141
77−newtype ForkId = ForkId Word64
78− deriving (Eq, Ord, Show)
79−
80−-- | Implicit 'ForkId' of the 'EnvRef' directly available from 'Env'.
81−globalFid :: ForkId
82−globalFid = ForkId 0
142+-- | Internal id of the fork.
143+newtype ForkId = ForkId { unForkId :: Int }
83144
84145-- | 'ForkId' generation.
85146newtype ForkIdGen = ForkIdGen (IORef ForkId)
86147
87−-- | Create a new thread local unique generator.
148+-- | Create a new thread local 'ForkId' generator.
88149newForkIdGen :: IO ForkIdGen
89−newForkIdGen = ForkIdGen <$> newIORef (ForkId 1)
150+newForkIdGen = ForkIdGen <$> newIORef (ForkId 0)
90151
91−-- | Clone the unique generator for use in a different thread.
152+-- | Clone the 'ForkId' generator for use in a different thread.
92153cloneForkIdGen :: ForkIdGen -> IO ForkIdGen
93154cloneForkIdGen (ForkIdGen ref) = fmap ForkIdGen . newIORef =<< readIORef ref
94155
@@ -102,12 +163,15 @@ newForkId (ForkIdGen ref) = do
102163----------------------------------------
103164-- Relinker
104165
105−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
106169 Relinker
107− :: ((forall es. Env es -> IO (Env es)) -> i e -> IO (i e))
108− -> Relinker i e
170+ :: ((forall es. Env es -> IO (Env es)) -> adapter e -> IO (adapter e))
171+ -> Relinker adapter e
109172
110−noRelinker :: Relinker i e
173+-- | A dummy 'Relinker' that does nothing.
174+noRelinker :: Relinker adapter e
111175noRelinker = Relinker $ \_ -> pure
112176
113177----------------------------------------
@@ -120,23 +184,27 @@ emptyEnv = Env NoFork <$> emptyEnvRef <*> newForkIdGen
120184-- | Clone the environment.
121185cloneEnv :: Env es -> IO (Env es)
122186cloneEnv (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)
123190 gen <- cloneForkIdGen gen0
124− cache <- newIORef M.empty
125− 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
126194 pure $ Env NoFork gref gen
127195cloneEnv (Env forks@(Forks _ baseIx lref0 _) gref0 gen0) = do
128196 EnvRef _ es0 fs0 <- readIORef gref0
129197 EnvRef n _ _ <- readIORef lref0
130198 let len = baseIx + n
131− es <- newSmallArray len (error "undefined field")
132− fs <- newSmallArray len (error "undefined field")
199+ es <- newSmallArray len undefinedData
200+ fs <- newSmallArray len undefinedData
133201 baseN <- copyForks es fs len forks
134202 copySmallMutableArray es 0 es0 0 baseN
135203 copySmallMutableArray fs 0 fs0 0 baseN
136204 gen <- cloneForkIdGen gen0
137205 gref <- newIORef $ EnvRef len es fs
138− cache <- newIORef $ M.singleton globalFid gref
139− relinkData gen cache es fs n
206+ store <- newIORef IM.empty
207+ relinkData gref gen store es fs n
140208 -- The forked environment is flattened and becomes the global one.
141209 pure $ Env NoFork gref gen
142210{-# NOINLINE cloneEnv #-}
@@ -186,33 +254,18 @@ copyForks es fs len = \case
186254 copySmallMutableArray fs baseIx fs0 0 n
187255 copyForks es fs baseIx forks
188256
189−type EnvRefCache = IORef (M.Map ForkId (IORef EnvRef))
257+type EnvRefStore = IORef (IM.IntMap (IORef EnvRef))
190258
191−-- | Clone the 'EnvRef' and put it in a cache.
192−cloneEnvRef
193− :: ForkIdGen
194− -> EnvRefCache
195− -> ForkId
196− -> IORef EnvRef
197− -> IO (IORef EnvRef)
198−cloneEnvRef gen cache fid ref0 = do
199− EnvRef n es0 fs0 <- readIORef ref0
200− es <- cloneSmallMutableArray es0 0 (sizeofSmallMutableArray es0)
201− fs <- cloneSmallMutableArray fs0 0 (sizeofSmallMutableArray fs0)
202− ref <- newIORef $ EnvRef n es fs
203− modifyIORef' cache $ M.insert fid ref
204− relinkData gen cache es fs n
205− pure ref
206−
207−-- | Relink local environments hiding in the interpreters.
259+-- | Relink local environments hiding in the handlers.
208260relinkData
209− :: ForkIdGen
210− -> EnvRefCache
261+ :: IORef EnvRef
262+ -> ForkIdGen
263+ -> EnvRefStore
211264 -> SmallMutableArray RealWorld Any
212265 -> SmallMutableArray RealWorld Any
213266 -> Int
214267 -> IO ()
215−relinkData gen cache es fs = \case
268+relinkData gref gen store es fs = \case
216269 0 -> pure ()
217270 n -> do
218271 let i = n - 1
@@ -220,29 +273,46 @@ relinkData gen cache es fs = \case
220273 readSmallArray es i
221274 >>= f relinkEnv . fromAny
222275 >>= writeSmallArray es i . toAny
223− relinkData gen cache es fs i
276+ relinkData gref gen store es fs i
224277 where
225278 relinkEnv :: Env es -> IO (Env es)
226− relinkEnv (Env forks _ _) = do
227− Just gref <- M.lookup globalFid <$> readIORef cache
228− 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
229282
230−relinkForks :: ForkIdGen -> EnvRefCache -> Forks -> IO Forks
231−relinkForks gen cache = \case
283+relinkForks :: IORef EnvRef -> ForkIdGen -> EnvRefStore -> Forks -> IO Forks
284+relinkForks gref gen store = \case
232285 NoFork -> pure NoFork
233286 Forks fid baseIx lref0 forks -> do
234287 -- A specific IORef EnvRef can be held by more than one local environment
235288 -- and we need to replace all its occurences with the same, new value
236289 -- containing its clone.
237− readIORef cache >>= pure . M.lookup fid >>= \case
290+ readIORef store >>= pure . IM.lookup (unForkId fid) >>= \case
238291 Just lref -> Forks fid baseIx <$> pure lref
239− <*> relinkForks gen cache forks
240− Nothing -> Forks fid baseIx <$> cloneEnvRef gen cache fid lref0
241− <*> 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
242312
243313----------------------------------------
244314
245−-- | Create a local fork of the environment for interpreters.
315+-- | Create a local fork of the environment for handlers.
246316forkEnv :: Env es -> IO (Env es)
247317forkEnv env@(Env NoFork gref gen) = do
248318 size <- sizeEnv env
@@ -252,7 +322,7 @@ forkEnv env@(Env NoFork gref gen) = do
252322forkEnv (Env forks@(Forks _ baseIx lref0 olderForks) gref gen) = do
253323 fid <- newForkId gen
254324 lref <- emptyEnvRef
255− (_capacity <$> readIORef lref0) >>= \case
325+ (_size <$> readIORef lref0) >>= \case
256326 -- If the fork is empty, replace it as no data is lost.
257327 0 -> pure $ Env (Forks fid baseIx lref olderForks) gref gen
258328 _ -> pure $ Env (Forks fid baseIx lref forks) gref gen
@@ -268,12 +338,6 @@ sizeEnv (Env (Forks _ baseIx lref _) _ _) = do
268338 pure $ baseIx + n
269339{-# NOINLINE sizeEnv #-}
270340
271−-- | Extract a specific data type from the environment.
272−getEnv :: forall e es i. e :> es => Env es -> IO (i e)
273−getEnv env = do
274− (i, es) <- getLocation @e env
275− fromAny <$> readSmallArray es i
276−
277341-- | Check that the size of the environment is the same as the expected value.
278342checkSizeEnv :: Int -> Env es -> IO ()
279343checkSizeEnv k (Env NoFork ref _) = do
@@ -291,7 +355,11 @@ checkSizeEnv k (Env (Forks _ baseIx lref _) _ _) = do
291355-- Extending and shrinking
292356
293357-- | Extend the environment with a new data type (in place).
294−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))
295363unsafeConsEnv e f (Env fork gref gen) = case fork of
296364 NoFork -> do
297365 extendEnvRef gref
@@ -312,10 +380,10 @@ unsafeConsEnv e f (Env fork gref gen) = case fork of
312380 writeIORef ref $! EnvRef (n + 1) es0 fs0
313381 EQ -> do
314382 let len = doubleCapacity len0
315− es <- newSmallArray len (error "undefined field")
383+ es <- newSmallArray len undefinedData
316384 copySmallMutableArray es 0 es0 0 len0
317385 e `seq` writeSmallArray es n (toAny e)
318− fs <- newSmallArray len (error "undefined field")
386+ fs <- newSmallArray len undefinedData
319387 copySmallMutableArray fs 0 fs0 0 len0
320388 f `seq` writeSmallArray fs n (toAny f)
321389 writeIORef ref $! EnvRef (n + 1) es fs
@@ -326,14 +394,13 @@ unsafeConsEnv e f (Env fork gref gen) = case fork of
326394
327395-- | Shrink the environment by one data type (in place). Makes sure the size of
328396-- the environment is as expected.
329−unsafeTailEnv :: Int -> Env (e : es) -> IO (Env es)
330−unsafeTailEnv len (Env fork gref gen) = case fork of
331− NoFork -> do
332− shrinkEnvRef len gref
333− pure $ Env NoFork gref gen
334− Forks _ baseIx lref _ -> do
335− shrinkEnvRef (len - baseIx) lref
336− 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
337404 where
338405 shrinkEnvRef :: Int -> IORef EnvRef -> IO ()
339406 shrinkEnvRef k ref = do
@@ -341,82 +408,109 @@ unsafeTailEnv len (Env fork gref gen) = case fork of
341408 if k /= n - 1
342409 then error $ "k (" ++ show k ++ ") /= n - 1 (" ++ show (n - 1) ++ ")"
343410 else do
344− writeSmallArray es k (error "undefined field")
345− writeSmallArray fs k (error "undefined field")
411+ writeSmallArray es k undefinedData
412+ writeSmallArray fs k undefinedData
346413 writeIORef ref $! EnvRef k es fs
347414{-# NOINLINE unsafeTailEnv #-}
348415
349416----------------------------------------
350417-- Data retrieval and update
351418
352−-- | Replace the data type in the environment with a new value (in place).
353−unsafePutEnv :: forall e es i. e :> es => i e -> Env es -> IO ()
354−unsafePutEnv e env = do
355− (i, es) <- getLocation @e env
356− 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
357427
358−-- | Modify the data type in the environment (in place).
359−unsafeModifyEnv :: forall e es i. e :> es => (i e -> i e) -> Env es -> IO ()
360−unsafeModifyEnv f env = do
361− (i, es) <- getLocation @e env
362− 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
363436 e `seq` writeSmallArray es i (toAny e)
364437
365438-- | Modify the data type in the environment (in place) and return a value.
366−unsafeStateEnv :: forall e es i a. e :> es => (i e -> (a, i e)) -> Env es -> IO a
367−unsafeStateEnv f env = do
368− (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
369446 (a, e) <- f . fromAny <$> readSmallArray es i
370447 e `seq` writeSmallArray es i (toAny e)
371448 pure a
372449
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+
373461----------------------------------------
374462-- Internal helpers
375463
464+undefinedData :: HasCallStack => a
465+undefinedData = error "undefined data"
466+
376467emptyEnvRef :: IO (IORef EnvRef)
377468emptyEnvRef = do
378− es <- newSmallArray 0 (error "undefined field")
379− fs <- newSmallArray 0 (error "undefined field")
469+ es <- newSmallArray 0 undefinedData
470+ fs <- newSmallArray 0 undefinedData
380471 newIORef $ EnvRef 0 es fs
381472
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+
382477-- | Determine location of the data type in the environment.
383478getLocation
384− :: forall e es. e :> es
385− => Env es
386− -> IO (Int, SmallMutableArray RealWorld Any)
387−getLocation (Env fork ref _) = do
388− 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
389484 -- Optimized for the common access pattern:
390485 --
391486 -- - Most application code has no access to forks, in which case we look at
392487 -- the global EnvRef.
393488 --
394− -- - Interpreters will most likely access the newest fork for reinterpretation
395− -- of effects that are there.
489+ -- - Effect handlers will most likely access the newest fork with handler
490+ -- specific effects for reinterpretation.
396491 case fork of
397− NoFork -> pure (ix n, es)
492+ NoFork -> pure $ Location (ix n) es
398493 Forks _ baseIx lref forks -> do
399494 EnvRef ln les _ <- readIORef lref
400495 let i = ix (baseIx + ln)
401496 if i >= baseIx
402− then pure (i - baseIx, les)
403− else go es i forks
497+ then pure $ Location (i - baseIx) les
498+ else go es# i forks
404499 where
405− go :: SmallMutableArray RealWorld Any
500+ go :: SmallMutableArray# RealWorld Any
406501 -> Int
407502 -> Forks
408− -> IO (Int, SmallMutableArray RealWorld Any)
409− go es i = \case
410− NoFork -> pure (i, es)
503+ -> IO Location
504+ go es# i = \case
505+ NoFork -> pure $ Location i (SmallMutableArray es#)
411506 Forks _ baseIx lref forks -> do
412507 EnvRef _ les _ <- readIORef lref
413508 if i >= baseIx
414− then pure (i - baseIx, les)
415− else go es i forks
509+ then pure $ Location (i - baseIx) les
510+ else go es# i forks
416511
417512 ix :: Int -> Int
418− ix n = n - reifyIndex @e @es - 1
419−{-# INLINE getLocation #-}
513+ ix n = n - ixE - 1
420514
421515toAny :: a -> Any
422516toAny = unsafeCoerce