git.iko.soy/effectful
on master help sign in
Support cloning of forked environments
67c1a67Andrzej Rybczak
src/Effectful/Internal/Env.hs+145−50
compare with master
unifiedsplit
@@ -31,6 +31,7 @@ import Control.Monad
31import Control.Monad.Primitive
31import Control.Monad.Primitive
32import Data.IORef
32import Data.IORef
33import Data.Primitive.SmallArray
33import Data.Primitive.SmallArray
34+import Data.Word
34import GHC.Exts (Any)
35import GHC.Exts (Any)
35import GHC.Stack
36import GHC.Stack
36import Unsafe.Coerce
37import Unsafe.Coerce
@@ -56,17 +57,7 @@ type role Env nominal
56--
57--
57-- - Modification of a specific element: /O(forks)/, usually /O(1)/.
58-- - Modification of a specific element: /O(forks)/, usually /O(1)/.
58--
59--
59−data Env (es :: [Effect]) = Env Forks (IORef EnvRef)
60+data Env (es :: [Effect]) = Env Forks (IORef EnvRef) UniqueGen
60−
61−-- | A fork is uniquely determined by its base index and depth.
62−data ForkId = ForkId
63− { baseIx :: Int
64− , depth :: Int
65− } deriving (Eq, Ord, Show)
66−
67−-- | Implicit 'ForkId' of the 'EnvRef' directly available from 'Env'.
68−globalFid :: ForkId
69−globalFid = ForkId 0 0
70
61
71-- | Local forks of the environment.
62-- | Local forks of the environment.
72data Forks = Forks {-# UNPACK #-} ForkId (IORef EnvRef) Forks | NoFork
63data Forks = Forks {-# UNPACK #-} ForkId (IORef EnvRef) Forks | NoFork
@@ -77,6 +68,21 @@ data EnvRef = EnvRef
77 , _relinkers :: SmallMutableArray RealWorld Any
68 , _relinkers :: SmallMutableArray RealWorld Any
78 }
69 }
79
70
71+----------------------------------------
72+-- ForkId
73+
74+data ForkId = ForkId
75+ { baseIx :: Int
76+ , unique :: Unique
77+ } deriving (Eq, Ord, Show)
78+
79+-- | Implicit 'ForkId' of the 'EnvRef' directly available from 'Env'.
80+globalFid :: ForkId
81+globalFid = ForkId 0 (Unique 0)
82+
83+----------------------------------------
84+-- Relinker
85+
80newtype Relinker i (e :: Effect) where
86newtype Relinker i (e :: Effect) where
81 Relinker
87 Relinker
82 :: ((forall es. Env es -> IO (Env es)) -> i e -> IO (i e))
88 :: ((forall es. Env es -> IO (Env es)) -> i e -> IO (i e))
@@ -85,47 +91,133 @@ newtype Relinker i (e :: Effect) where
85noRelinker :: Relinker i e
91noRelinker :: Relinker i e
86noRelinker = Relinker $ \_ -> pure
92noRelinker = Relinker $ \_ -> pure
87
93
94+----------------------------------------
95+-- UniqueGen
96+
97+newtype Unique = Unique Word64
98+ deriving (Eq, Ord, Show)
99+
100+-- | Uniques for 'ForkId'S.
101+newtype UniqueGen = UniqueGen (IORef Unique)
102+
103+-- | Create a new thread local unique generator.
104+newUniqueGen :: IO UniqueGen
105+newUniqueGen = UniqueGen <$> newIORef (Unique 1)
106+
107+-- | Clone the unique generator for use in a different thread.
108+cloneUniqueGen :: UniqueGen -> IO UniqueGen
109+cloneUniqueGen (UniqueGen ref) = fmap UniqueGen . newIORef =<< readIORef ref
110+
111+-- | Get a 'Unique' from the generator.
112+getUnique :: UniqueGen -> IO Unique
113+getUnique (UniqueGen ref) = do
114+ unique@(Unique n) <- readIORef ref
115+ writeIORef ref $! Unique (n + 1)
116+ pure unique
117+
118+----------------------------------------
119+-- Operations
120+
88-- | Create an empty environment.
121-- | Create an empty environment.
89emptyEnv :: IO (Env '[])
122emptyEnv :: IO (Env '[])
90−emptyEnv = Env NoFork <$> emptyEnvRef
123+emptyEnv = Env NoFork <$> emptyEnvRef <*> newUniqueGen
91
124
92-- | Clone the environment.
125-- | Clone the environment.
93cloneEnv :: HasCallStack => Env es -> IO (Env es)
126cloneEnv :: HasCallStack => Env es -> IO (Env es)
94−cloneEnv (Env NoFork gref0) = do
127+cloneEnv (Env NoFork gref0 ug0) = do
128+ ug <- cloneUniqueGen ug0
95 cache <- newIORef M.empty
129 cache <- newIORef M.empty
96− gref <- cloneEnvRef cache globalFid gref0
130+ gref <- cloneEnvRef ug cache globalFid gref0
97− pure $ Env NoFork gref
131+ pure $ Env NoFork gref ug
98−cloneEnv _ = error "Cloning of forked enviroment is not supported yet"
132+cloneEnv (Env forks@(Forks fid lref0 _) gref0 ug0) = do
133+ EnvRef _ es0 fs0 <- readIORef gref0
134+ EnvRef n _ _ <- readIORef lref0
135+ let len = baseIx fid + n
136+ es <- newSmallArray len (error "undefined field")
137+ fs <- newSmallArray len (error "undefined field")
138+ baseN <- copyForks es fs len forks
139+ copySmallMutableArray es 0 es0 0 baseN
140+ copySmallMutableArray fs 0 fs0 0 baseN
141+ ug <- cloneUniqueGen ug0
142+ gref <- newIORef $ EnvRef len es fs
143+ cache <- newIORef $ M.singleton globalFid gref
144+ relinkData ug cache es fs n
145+ -- The forked environment is flattened and becomes the global one.
146+ pure $ Env NoFork gref ug
99{-# NOINLINE cloneEnv #-}
147{-# NOINLINE cloneEnv #-}
100
148
101----------------------------------------
149----------------------------------------
102-- Utils for cloning
150-- Utils for cloning
103
151
152+-- | Let's say that the forked environment looks like this:
153+--
154+-- [0,1][2,3,4,5][6,7,8,9][10,11]
155+--
156+-- Then forks will look like this:
157+--
158+-- Fork (baseIx: 10, unique: ...) [10,11]
159+-- (Fork (baseIx: 6, unique: ...) [6,7,8,9]
160+-- (Fork (baseIx: 2, unique: ...) [2,3,4,5]
161+-- NoFork))
162+--
163+-- and elements [0,1] are taken from the global environment.
164+--
165+-- We start with len: 12, we subtract baseIx: 10 and get n: 2, the number of
166+-- elements to copy from the fork. We copy them, then call recursively with
167+-- len-n (10).
168+--
169+-- Then len: 10, we subtract baseIx: 6 and get n: 4. the number of elements to
170+-- copy from the fork (note that we can't use the capacity from EnvRef, because
171+-- it might be longer than 4 if the fork was locally extended). We copy them,
172+-- then call recursively with len-n (6).
173+--
174+-- Then len: 6, we subtract baseIx: 2 and get n: 4. We copy the elements, then
175+-- call recursively with len-n (2).
176+--
177+-- Now len: 2, but there are no more forks left, so we return len (2), the
178+-- number of elements to copy from the global environment.
179+copyForks
180+ :: SmallMutableArray RealWorld Any
181+ -> SmallMutableArray RealWorld Any
182+ -> Int
183+ -> Forks
184+ -> IO Int
185+copyForks es fs len = \case
186+ NoFork -> pure len
187+ Forks fid lref0 forks -> do
188+ EnvRef _ es0 fs0 <- readIORef lref0
189+ let n = len - baseIx fid
190+ copySmallMutableArray es (baseIx fid) es0 0 n
191+ copySmallMutableArray fs (baseIx fid) fs0 0 n
192+ copyForks es fs (len - n) forks
193+
104type EnvRefCache = IORef (M.Map ForkId (IORef EnvRef))
194type EnvRefCache = IORef (M.Map ForkId (IORef EnvRef))
105
195
106-- | Clone the 'EnvRef' and put it in a cache.
196-- | Clone the 'EnvRef' and put it in a cache.
107cloneEnvRef
197cloneEnvRef
198+ :: UniqueGen
108− :: EnvRefCache
199+ -> EnvRefCache
109 -> ForkId
200 -> ForkId
110 -> IORef EnvRef
201 -> IORef EnvRef
111 -> IO (IORef EnvRef)
202 -> IO (IORef EnvRef)
112−cloneEnvRef cache fid ref0 = do
203+cloneEnvRef ug cache fid ref0 = do
113 EnvRef n es0 fs0 <- readIORef ref0
204 EnvRef n es0 fs0 <- readIORef ref0
114 es <- cloneSmallMutableArray es0 0 (sizeofSmallMutableArray es0)
205 es <- cloneSmallMutableArray es0 0 (sizeofSmallMutableArray es0)
115 fs <- cloneSmallMutableArray fs0 0 (sizeofSmallMutableArray fs0)
206 fs <- cloneSmallMutableArray fs0 0 (sizeofSmallMutableArray fs0)
116 ref <- newIORef $ EnvRef n es fs
207 ref <- newIORef $ EnvRef n es fs
117 modifyIORef' cache $ M.insert fid ref
208 modifyIORef' cache $ M.insert fid ref
118− relinkData cache es fs n
209+ relinkData ug cache es fs n
119 pure ref
210 pure ref
120
211
121-- | Relink local environments hiding in the interpreters.
212-- | Relink local environments hiding in the interpreters.
122relinkData
213relinkData
214+ :: UniqueGen
123− :: EnvRefCache
215+ -> EnvRefCache
124 -> SmallMutableArray RealWorld Any
216 -> SmallMutableArray RealWorld Any
125 -> SmallMutableArray RealWorld Any
217 -> SmallMutableArray RealWorld Any
126 -> Int
218 -> Int
127 -> IO ()
219 -> IO ()
128−relinkData cache es fs = \case
220+relinkData ug cache es fs = \case
129 0 -> pure ()
221 0 -> pure ()
130 n -> do
222 n -> do
131 let i = n - 1
223 let i = n - 1
@@ -133,59 +225,62 @@ relinkData cache es fs = \case
133 readSmallArray es i
225 readSmallArray es i
134 >>= f relinkEnv . fromAny
226 >>= f relinkEnv . fromAny
135 >>= writeSmallArray es i . toAny
227 >>= writeSmallArray es i . toAny
136− relinkData cache es fs i
228+ relinkData ug cache es fs i
137 where
229 where
138 relinkEnv :: Env es -> IO (Env es)
230 relinkEnv :: Env es -> IO (Env es)
139− relinkEnv (Env forks _) = do
231+ relinkEnv (Env forks _ _) = do
140 Just gref <- M.lookup globalFid <$> readIORef cache
232 Just gref <- M.lookup globalFid <$> readIORef cache
141− Env <$> relinkForks cache forks <*> pure gref
233+ Env <$> relinkForks ug cache forks <*> pure gref <*> pure ug
142
234
143−relinkForks :: EnvRefCache -> Forks -> IO Forks
235+relinkForks :: UniqueGen -> EnvRefCache -> Forks -> IO Forks
144−relinkForks cache = \case
236+relinkForks ug cache = \case
145 NoFork -> pure NoFork
237 NoFork -> pure NoFork
146− Forks fid lref0 forks -> (M.lookup fid <$> readIORef cache) >>= \case
238+ Forks fid lref0 forks -> do
239+ readIORef cache >>= pure . M.lookup fid >>= \case
147− Just lref -> Forks fid <$> pure lref
240+ Just lref -> Forks fid <$> pure lref
148− <*> relinkForks cache forks
241+ <*> relinkForks ug cache forks
149− Nothing -> Forks fid <$> cloneEnvRef cache fid lref0
242+ Nothing -> Forks fid <$> cloneEnvRef ug cache fid lref0
150− <*> relinkForks cache forks
243+ <*> relinkForks ug cache forks
151
244
152----------------------------------------
245----------------------------------------
153
246
154-- | Create a local fork of the environment for interpreters.
247-- | Create a local fork of the environment for interpreters.
155forkEnv :: HasCallStack => Env es -> IO (Env es)
248forkEnv :: HasCallStack => Env es -> IO (Env es)
156−forkEnv env@(Env NoFork gref) = do
249+forkEnv env@(Env NoFork gref ug) = do
157 size <- sizeEnv env
250 size <- sizeEnv env
251+ uniq <- getUnique ug
158 let fid = ForkId
252 let fid = ForkId
159 { baseIx = size
253 { baseIx = size
160− , depth = 1
254+ , unique = uniq
161 }
255 }
162 Env <$> (Forks fid <$> emptyEnvRef <*> pure NoFork)
256 Env <$> (Forks fid <$> emptyEnvRef <*> pure NoFork)
163− <*> pure gref
257+ <*> pure gref <*> pure ug
164−forkEnv (Env forks@(Forks fid0 lref0 olderForks) gref) = do
258+forkEnv (Env forks@(Forks fid0 lref0 olderForks) gref ug) = do
165 EnvRef n _ _ <- readIORef lref0
259 EnvRef n _ _ <- readIORef lref0
260+ uniq <- getUnique ug
166 -- If the fork is empty, replace it as no data is lost.
261 -- If the fork is empty, replace it as no data is lost.
167 if n == 0
262 if n == 0
168 then do
263 then do
169 lref <- emptyEnvRef
264 lref <- emptyEnvRef
170 let fid = fid0
265 let fid = fid0
171− { depth = depth fid0 + 1
266+ { unique = uniq
172 }
267 }
173− pure $ Env (Forks fid lref olderForks) gref
268+ pure $ Env (Forks fid lref olderForks) gref ug
174 else do
269 else do
175 lref <- emptyEnvRef
270 lref <- emptyEnvRef
176 let fid = fid0
271 let fid = fid0
177 { baseIx = baseIx fid0 + n
272 { baseIx = baseIx fid0 + n
178− , depth = depth fid0 + 1
273+ , unique = uniq
179 }
274 }
180− pure $ Env (Forks fid lref forks) gref
275+ pure $ Env (Forks fid lref forks) gref ug
181{-# NOINLINE forkEnv #-}
276{-# NOINLINE forkEnv #-}
182
277
183-- | Get the current size of the environment.
278-- | Get the current size of the environment.
184sizeEnv :: Env es -> IO Int
279sizeEnv :: Env es -> IO Int
185−sizeEnv (Env NoFork ref) = do
280+sizeEnv (Env NoFork ref _) = do
186 EnvRef n _ _ <- readIORef ref
281 EnvRef n _ _ <- readIORef ref
187 pure n
282 pure n
188−sizeEnv (Env (Forks fid lref _) _) = do
283+sizeEnv (Env (Forks fid lref _) _ _) = do
189 EnvRef n _ _ <- readIORef lref
284 EnvRef n _ _ <- readIORef lref
190 pure $ baseIx fid + n
285 pure $ baseIx fid + n
191{-# NOINLINE sizeEnv #-}
286{-# NOINLINE sizeEnv #-}
@@ -198,11 +293,11 @@ getEnv env = do
198
293
199-- | Check that the size of the environment is the same as the expected value.
294-- | Check that the size of the environment is the same as the expected value.
200checkSizeEnv :: HasCallStack => Int -> Env es -> IO ()
295checkSizeEnv :: HasCallStack => Int -> Env es -> IO ()
201−checkSizeEnv k (Env NoFork ref) = do
296+checkSizeEnv k (Env NoFork ref _) = do
202 EnvRef n _ _ <- readIORef ref
297 EnvRef n _ _ <- readIORef ref
203 when (k /= n) $ do
298 when (k /= n) $ do
204 error $ "k (" ++ show k ++ ") /= n (" ++ show n ++ ")"
299 error $ "k (" ++ show k ++ ") /= n (" ++ show n ++ ")"
205−checkSizeEnv k (Env (Forks fid lref _) _) = do
300+checkSizeEnv k (Env (Forks fid lref _) _ _) = do
206 EnvRef n _ _ <- readIORef lref
301 EnvRef n _ _ <- readIORef lref
207 when (k /= baseIx fid + n) $ do
302 when (k /= baseIx fid + n) $ do
208 error $ "k (" ++ show k ++ ") /= baseIx + n (baseIx: "
303 error $ "k (" ++ show k ++ ") /= baseIx + n (baseIx: "
@@ -214,13 +309,13 @@ checkSizeEnv k (Env (Forks fid lref _) _) = do
214
309
215-- | Extend the environment with a new data type (in place).
310-- | Extend the environment with a new data type (in place).
216unsafeConsEnv :: HasCallStack => i e -> Relinker i e -> Env es -> IO (Env (e : es))
311unsafeConsEnv :: HasCallStack => i e -> Relinker i e -> Env es -> IO (Env (e : es))
217−unsafeConsEnv e f (Env fork gref) = case fork of
312+unsafeConsEnv e f (Env fork gref ug) = case fork of
218 NoFork -> do
313 NoFork -> do
219 extendEnvRef gref
314 extendEnvRef gref
220− pure $ Env NoFork gref
315+ pure $ Env NoFork gref ug
221 Forks base lref forks -> do
316 Forks base lref forks -> do
222 extendEnvRef lref
317 extendEnvRef lref
223− pure $ Env (Forks base lref forks) gref
318+ pure $ Env (Forks base lref forks) gref ug
224 where
319 where
225 extendEnvRef :: IORef EnvRef -> IO ()
320 extendEnvRef :: IORef EnvRef -> IO ()
226 extendEnvRef ref = do
321 extendEnvRef ref = do
@@ -249,13 +344,13 @@ unsafeConsEnv e f (Env fork gref) = case fork of
249-- | Shrink the environment by one data type (in place). Makes sure the size of
344-- | Shrink the environment by one data type (in place). Makes sure the size of
250-- the environment is as expected.
345-- the environment is as expected.
251unsafeTailEnv :: HasCallStack => Int -> Env (e : es) -> IO (Env es)
346unsafeTailEnv :: HasCallStack => Int -> Env (e : es) -> IO (Env es)
252−unsafeTailEnv len (Env fork gref) = case fork of
347+unsafeTailEnv len (Env fork gref ug) = case fork of
253 NoFork -> do
348 NoFork -> do
254 shrinkEnvRef len gref
349 shrinkEnvRef len gref
255− pure $ Env NoFork gref
350+ pure $ Env NoFork gref ug
256 Forks fid lref forks -> do
351 Forks fid lref forks -> do
257 shrinkEnvRef (len - baseIx fid) lref
352 shrinkEnvRef (len - baseIx fid) lref
258− pure $ Env (Forks fid lref forks) gref
353+ pure $ Env (Forks fid lref forks) gref ug
259 where
354 where
260 shrinkEnvRef :: Int -> IORef EnvRef -> IO ()
355 shrinkEnvRef :: Int -> IORef EnvRef -> IO ()
261 shrinkEnvRef k ref = do
356 shrinkEnvRef k ref = do
@@ -317,7 +412,7 @@ getLocation
317 :: forall e es. e :> es
412 :: forall e es. e :> es
318 => Env es
413 => Env es
319 -> IO (Int, SmallMutableArray RealWorld Any)
414 -> IO (Int, SmallMutableArray RealWorld Any)
320−getLocation (Env fork ref) = do
415+getLocation (Env fork ref _) = do
321 EnvRef n es _ <- readIORef ref
416 EnvRef n es _ <- readIORef ref
322 -- Optimized for the common access pattern:
417 -- Optimized for the common access pattern:
323 --
418 --