@@ -1,14 +1,18 @@
1+{-# LANGUAGE AllowAmbiguousTypes #-}
12{-# OPTIONS_HADDOCK not-home #-}
2−-- | The enviroment for 'Effectful.Internal.Monad.Eff'.
3+-- | The environment for 'Effectful.Internal.Monad.Eff'.
34--
45-- This module is intended for internal use only, and may change without warning
56-- in subsequent releases.
67module Effectful.Internal.Env
78 ( Env
9+ , Relinker(..)
10+ , noRelinker
811
912 -- * Safe operations
1013 , emptyEnv
1114 , cloneEnv
15+ , forkEnv
1216 , sizeEnv
1317 , getEnv
1418 , checkSizeEnv
@@ -30,8 +34,9 @@ import Data.Primitive.SmallArray
3034import GHC.Exts (Any)
3135import GHC.Stack
3236import Unsafe.Coerce
37+import qualified Data.Map.Strict as M
3338
34−import Effectful.Internal.Has
39+import Effectful.Internal.Effect
3540
3641type role Env nominal
3742
@@ -48,61 +53,192 @@ type role Env nominal
4853--
4954-- - Modification of a specific element: /O(1)/.
5055--
51−newtype Env (es :: [Effect]) = Env (IORef EnvRef)
56+data Env (es :: [Effect]) = Env Forks (IORef EnvRef)
5257
53−data EnvRef = EnvRef Int (SmallMutableArray RealWorld Any)
58+-- | A fork is uniquely determined by its base index and depth.
59+data ForkId = ForkId
60+ { baseIx :: Int
61+ , depth :: Int
62+ } deriving (Eq, Ord, Show)
63+
64+-- | Implicit 'ForkId' of the 'EnvRef' directly available from 'Env'.
65+globalFid :: ForkId
66+globalFid = ForkId 0 0
67+
68+-- | Local forks of the environment.
69+data Forks = Forks ForkId (IORef EnvRef) Forks | NoFork
70+
71+data EnvRef = EnvRef
72+ { _capacity :: Int
73+ , _data :: SmallMutableArray RealWorld Any
74+ , _relinkers :: SmallMutableArray RealWorld Any
75+ }
76+
77+newtype Relinker i (e :: Effect) where
78+ Relinker
79+ :: ((forall es. Env es -> IO (Env es)) -> i e -> IO (i e))
80+ -> Relinker i e
81+
82+noRelinker :: Relinker i e
83+noRelinker = Relinker $ \_ -> pure
5484
5585-- | Create an empty environment.
5686emptyEnv :: IO (Env '[])
57−emptyEnv = fmap Env . newIORef . EnvRef 0
58− =<< newSmallArray 0 (error "undefined field")
87+emptyEnv = Env NoFork <$> emptyEnvRef
5988
6089-- | Clone the environment.
61−cloneEnv :: Env es -> IO (Env es)
62−cloneEnv (Env ref) = do
63− EnvRef n es <- readIORef ref
64− fmap Env . newIORef . EnvRef n
65− =<< cloneSmallMutableArray es 0 (sizeofSmallMutableArray es)
90+cloneEnv :: HasCallStack => Env es -> IO (Env es)
91+cloneEnv (Env NoFork gref0) = do
92+ cache <- newIORef M.empty
93+ gref <- cloneEnvRef cache globalFid gref0
94+ pure $ Env NoFork gref
95+cloneEnv _ = error "Cloning of forked enviroment is not supported yet"
96+{-# NOINLINE cloneEnv #-}
97+
98+----------------------------------------
99+-- Utils for cloning
100+
101+type EnvRefCache = IORef (M.Map ForkId (IORef EnvRef))
102+
103+-- | Clone the 'EnvRef' and put it in a cache.
104+cloneEnvRef
105+ :: EnvRefCache
106+ -> ForkId
107+ -> IORef EnvRef
108+ -> IO (IORef EnvRef)
109+cloneEnvRef cache fid ref0 = do
110+ EnvRef n es0 fs0 <- readIORef ref0
111+ es <- cloneSmallMutableArray es0 0 (sizeofSmallMutableArray es0)
112+ fs <- cloneSmallMutableArray fs0 0 (sizeofSmallMutableArray fs0)
113+ ref <- newIORef $ EnvRef n es fs
114+ modifyIORef' cache $ M.insert fid ref
115+ relinkData cache es fs n
116+ pure ref
117+
118+-- | Relink local environments hiding in the interpreters.
119+relinkData
120+ :: EnvRefCache
121+ -> SmallMutableArray RealWorld Any
122+ -> SmallMutableArray RealWorld Any
123+ -> Int
124+ -> IO ()
125+relinkData cache es fs = \case
126+ 0 -> pure ()
127+ n -> do
128+ let i = n - 1
129+ Relinker f <- fromAny <$> readSmallArray fs i
130+ readSmallArray es i
131+ >>= f relinkEnv . fromAny
132+ >>= writeSmallArray es i . toAny
133+ relinkData cache es fs i
134+ where
135+ relinkEnv :: Env es -> IO (Env es)
136+ relinkEnv (Env forks _) = do
137+ Just gref <- M.lookup globalFid <$> readIORef cache
138+ Env <$> relinkForks cache forks <*> pure gref
139+
140+relinkForks :: EnvRefCache -> Forks -> IO Forks
141+relinkForks cache = \case
142+ NoFork -> pure NoFork
143+ Forks fid lref0 forks -> (M.lookup fid <$> readIORef cache) >>= \case
144+ Just lref -> Forks fid <$> pure lref
145+ <*> relinkForks cache forks
146+ Nothing -> Forks fid <$> cloneEnvRef cache fid lref0
147+ <*> relinkForks cache forks
148+
149+----------------------------------------
150+
151+-- | Create a local fork of the environment for interpreters.
152+forkEnv :: HasCallStack => Env es -> IO (Env es)
153+forkEnv env@(Env NoFork gref) = do
154+ size <- sizeEnv env
155+ let fid = ForkId
156+ { baseIx = size
157+ , depth = 1
158+ }
159+ Env <$> (Forks fid <$> emptyEnvRef <*> pure NoFork)
160+ <*> pure gref
161+forkEnv (Env forks@(Forks fid0 lref0 olderForks) gref) = do
162+ EnvRef n _ _ <- readIORef lref0
163+ -- If the fork is empty, replace it as no data is lost.
164+ if n == 0
165+ then do
166+ lref <- emptyEnvRef
167+ let fid = fid0
168+ { depth = depth fid0 + 1
169+ }
170+ pure $ Env (Forks fid lref olderForks) gref
171+ else do
172+ lref <- emptyEnvRef
173+ let fid = fid0
174+ { baseIx = baseIx fid0 + n
175+ , depth = depth fid0 + 1
176+ }
177+ pure $ Env (Forks fid lref forks) gref
178+{-# NOINLINE forkEnv #-}
66179
67180-- | Get the current size of the environment.
68181sizeEnv :: Env es -> IO Int
69−sizeEnv (Env ref) = do
70− EnvRef n _ <- readIORef ref
182+sizeEnv (Env NoFork ref) = do
183+ EnvRef n _ _ <- readIORef ref
71184 pure n
185+sizeEnv (Env (Forks fid lref _) _) = do
186+ EnvRef n _ _ <- readIORef lref
187+ pure $ baseIx fid + n
188+{-# NOINLINE sizeEnv #-}
72189
73190-- | Extract a specific data type from the environment.
74−getEnv :: forall e es. (HasCallStack, e :> es) => Env es -> IO e
75−getEnv (Env ref) = do
76− EnvRef n es <- readIORef ref
77− fromAny <$> readSmallArray es (ixEnv @e @es n)
191+getEnv :: forall e es i. (HasCallStack, e :> es) => Env es -> IO (i e)
192+getEnv env = do
193+ (i, es) <- getLocation @e env
194+ fromAny <$> readSmallArray es i
78195
79196-- | Check that the size of the environment is the same as the expected value.
80197checkSizeEnv :: HasCallStack => Int -> Env es -> IO ()
81−checkSizeEnv k (Env ref) = do
82− EnvRef n _ <- readIORef ref
83− when (k /= n) . error $ "k (" ++ show k ++ ") /= n (" ++ show n ++ ")"
198+checkSizeEnv k (Env NoFork ref) = do
199+ EnvRef n _ _ <- readIORef ref
200+ when (k /= n) $ do
201+ error $ "k (" ++ show k ++ ") /= n (" ++ show n ++ ")"
202+checkSizeEnv k (Env (Forks fid lref _) _) = do
203+ EnvRef n _ _ <- readIORef lref
204+ when (k /= baseIx fid + n) $ do
205+ error $ "k (" ++ show k ++ ") /= baseIx + n (baseIx: "
206+ ++ show (baseIx fid) ++ ", n: " ++ show n ++ ")"
207+{-# NOINLINE checkSizeEnv #-}
84208
85209----------------------------------------
86210-- Extending and shrinking
87211
88212-- | Extend the environment with a new data type (in place).
89−unsafeConsEnv :: HasCallStack => e -> Env es -> IO (Env (e : es))
90−unsafeConsEnv e (Env ref) = do
91− EnvRef n es0 <- readIORef ref
92− let len0 = sizeofSmallMutableArray es0
93− case n `compare` len0 of
94− GT -> error $ "n (" ++ show n ++ ") > len0 (" ++ show len0 ++ ")"
95− LT -> do
96− e `seq` writeSmallArray es0 n (toAny e)
97− writeIORef ref $! EnvRef (n + 1) es0
98− EQ -> do
99− let len = doubleCapacity len0
100− es <- newSmallArray len (error "undefined field")
101− copySmallMutableArray es 0 es0 0 len0
102− e `seq` writeSmallArray es n (toAny e)
103− writeIORef ref $! EnvRef (n + 1) es
104− pure $ Env ref
213+unsafeConsEnv :: HasCallStack => i e -> Relinker i e -> Env es -> IO (Env (e : es))
214+unsafeConsEnv e f (Env fork gref) = case fork of
215+ NoFork -> do
216+ extendEnvRef gref
217+ pure $ Env NoFork gref
218+ Forks base lref forks -> do
219+ extendEnvRef lref
220+ pure $ Env (Forks base lref forks) gref
105221 where
222+ extendEnvRef :: IORef EnvRef -> IO ()
223+ extendEnvRef ref = do
224+ EnvRef n es0 fs0 <- readIORef ref
225+ let len0 = sizeofSmallMutableArray es0
226+ case n `compare` len0 of
227+ GT -> error $ "n (" ++ show n ++ ") > len0 (" ++ show len0 ++ ")"
228+ LT -> do
229+ e `seq` writeSmallArray es0 n (toAny e)
230+ f `seq` writeSmallArray fs0 n (toAny f)
231+ writeIORef ref $! EnvRef (n + 1) es0 fs0
232+ EQ -> do
233+ let len = doubleCapacity len0
234+ es <- newSmallArray len (error "undefined field")
235+ copySmallMutableArray es 0 es0 0 len0
236+ e `seq` writeSmallArray es n (toAny e)
237+ fs <- newSmallArray len (error "undefined field")
238+ copySmallMutableArray fs 0 fs0 0 len0
239+ f `seq` writeSmallArray fs n (toAny f)
240+ writeIORef ref $! EnvRef (n + 1) es fs
241+
106242 doubleCapacity :: Int -> Int
107243 doubleCapacity n = max 1 n * 2
108244{-# NOINLINE unsafeConsEnv #-}
@@ -110,14 +246,23 @@ unsafeConsEnv e (Env ref) = do
110246-- | Shrink the environment by one data type (in place). Makes sure the size of
111247-- the environment is as expected.
112248unsafeTailEnv :: HasCallStack => Int -> Env (e : es) -> IO (Env es)
113−unsafeTailEnv k (Env ref) = do
114− EnvRef n es <- readIORef ref
115− if k /= n - 1
116− then error $ "k (" ++ show k ++ ") /= n - 1 (" ++ show (n - 1) ++ ")"
117− else do
118− writeSmallArray es k (error "undefined field")
119− writeIORef ref $! EnvRef k es
120− pure $ Env ref
249+unsafeTailEnv len (Env fork gref) = case fork of
250+ NoFork -> do
251+ shrinkEnvRef len gref
252+ pure $ Env NoFork gref
253+ Forks fid lref forks -> do
254+ shrinkEnvRef (len - baseIx fid) lref
255+ pure $ Env (Forks fid lref forks) gref
256+ where
257+ shrinkEnvRef :: Int -> IORef EnvRef -> IO ()
258+ shrinkEnvRef k ref = do
259+ EnvRef n es fs <- readIORef ref
260+ if k /= n - 1
261+ then error $ "k (" ++ show k ++ ") /= n - 1 (" ++ show (n - 1) ++ ")"
262+ else do
263+ writeSmallArray es k (error "undefined field")
264+ writeSmallArray fs k (error "undefined field")
265+ writeIORef ref $! EnvRef k es fs
121266{-# NOINLINE unsafeTailEnv #-}
122267
123268----------------------------------------
@@ -125,34 +270,32 @@ unsafeTailEnv k (Env ref) = do
125270
126271-- | Replace the data type in the environment with a new value (in place).
127272unsafePutEnv
128− :: forall e es. (HasCallStack, e :> es)
129− => e
273+ :: forall e es i. (HasCallStack, e :> es)
274+ => i e
130275 -> Env es
131276 -> IO ()
132−unsafePutEnv e (Env ref) = do
133− EnvRef n es <- readIORef ref
134− e `seq` writeSmallArray es (ixEnv @e @es n) (toAny e)
277+unsafePutEnv e env = do
278+ (i, es) <- getLocation @e env
279+ e `seq` writeSmallArray es i (toAny e)
135280
136281-- | Modify the data type in the environment (in place).
137282unsafeModifyEnv
138− :: forall e es. (HasCallStack, e :> es)
139− => (e -> e)
283+ :: forall e es i. (HasCallStack, e :> es)
284+ => (i e -> i e)
140285 -> Env es
141286 -> IO ()
142−unsafeModifyEnv f (Env ref) = do
143− EnvRef n es <- readIORef ref
144− let i = ixEnv @e @es n
287+unsafeModifyEnv f env = do
288+ (i, es) <- getLocation @e env
145289 e <- f . fromAny <$> readSmallArray es i
146290 e `seq` writeSmallArray es i (toAny e)
147291
148−-- | Modify the data type in the enviroment (in place) and return a value.
292+-- | Modify the data type in the environment (in place) and return a value.
149293unsafeStateEnv
150− :: forall e es a. (HasCallStack, e :> es)
151− => (e -> (a, e))
294+ :: forall e es i a. (HasCallStack, e :> es)
295+ => (i e -> (a, i e))
152296 -> Env es -> IO a
153−unsafeStateEnv f (Env ref) = do
154− EnvRef n es <- readIORef ref
155− let i = ixEnv @e @es n
297+unsafeStateEnv f env = do
298+ (i, es) <- getLocation @e env
156299 (a, e) <- f . fromAny <$> readSmallArray es i
157300 e `seq` writeSmallArray es i (toAny e)
158301 pure a
@@ -160,6 +303,50 @@ unsafeStateEnv f (Env ref) = do
160303----------------------------------------
161304-- Internal helpers
162305
306+emptyEnvRef :: IO (IORef EnvRef)
307+emptyEnvRef = do
308+ es <- newSmallArray 0 (error "undefined field")
309+ fs <- newSmallArray 0 (error "undefined field")
310+ newIORef $ EnvRef 0 es fs
311+
312+-- | Determine location of the data type in the environment.
313+getLocation
314+ :: forall e es. e :> es
315+ => Env es
316+ -> IO (Int, SmallMutableArray RealWorld Any)
317+getLocation (Env fork ref) = do
318+ EnvRef n es _ <- readIORef ref
319+ -- Optimized for the common access pattern:
320+ --
321+ -- - Most application code has no access to forks, in which case we look at
322+ -- the global EnvRef.
323+ --
324+ -- - Interpreters will most likely access the newest fork for reinterpretation
325+ -- of effects that are there.
326+ case fork of
327+ NoFork -> pure (ix n, es)
328+ Forks fid lref forks -> do
329+ EnvRef ln les _ <- readIORef lref
330+ let i = ix (baseIx fid + ln)
331+ if i >= baseIx fid
332+ then pure (i - baseIx fid, les)
333+ else go es i forks
334+ where
335+ go :: SmallMutableArray RealWorld Any
336+ -> Int
337+ -> Forks
338+ -> IO (Int, SmallMutableArray RealWorld Any)
339+ go es i = \case
340+ NoFork -> pure (i, es)
341+ Forks fid lref forks -> do
342+ EnvRef _ les _ <- readIORef lref
343+ if i >= baseIx fid
344+ then pure (i - baseIx fid, les)
345+ else go es i forks
346+
347+ ix :: Int -> Int
348+ ix n = n - reifyIndex @e @es - 1
349+
163350toAny :: a -> Any
164351toAny = unsafeCoerce
165352