@@ -31,6 +31,7 @@ import Control.Monad
3131import Control.Monad.Primitive
3232import Data.IORef
3333import Data.Primitive.SmallArray
34+import Data.Word
3435import GHC.Exts (Any)
3536import GHC.Stack
3637import Unsafe.Coerce
@@ -56,17 +57,7 @@ type role Env nominal
5657--
5758-- - Modification of a specific element: /O(forks)/, usually /O(1)/.
5859--
59−data Env (es :: [Effect]) = Env Forks (IORef EnvRef)
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
60+data Env (es :: [Effect]) = Env Forks (IORef EnvRef) UniqueGen
7061
7162-- | Local forks of the environment.
7263data Forks = Forks {-# UNPACK #-} ForkId (IORef EnvRef) Forks | NoFork
@@ -77,6 +68,21 @@ data EnvRef = EnvRef
7768 , _relinkers :: SmallMutableArray RealWorld Any
7869 }
7970
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+
8086newtype Relinker i (e :: Effect) where
8187 Relinker
8288 :: ((forall es. Env es -> IO (Env es)) -> i e -> IO (i e))
@@ -85,47 +91,133 @@ newtype Relinker i (e :: Effect) where
8591noRelinker :: Relinker i e
8692noRelinker = Relinker $ \_ -> pure
8793
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+
88121-- | Create an empty environment.
89122emptyEnv :: IO (Env '[])
90−emptyEnv = Env NoFork <$> emptyEnvRef
123+emptyEnv = Env NoFork <$> emptyEnvRef <*> newUniqueGen
91124
92125-- | Clone the environment.
93126cloneEnv :: HasCallStack => Env es -> IO (Env es)
94−cloneEnv (Env NoFork gref0) = do
127+cloneEnv (Env NoFork gref0 ug0) = do
128+ ug <- cloneUniqueGen ug0
95129 cache <- newIORef M.empty
96− gref <- cloneEnvRef cache globalFid gref0
97− pure $ Env NoFork gref
98−cloneEnv _ = error "Cloning of forked enviroment is not supported yet"
130+ gref <- cloneEnvRef ug cache globalFid gref0
131+ pure $ Env NoFork gref ug
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
99147{-# NOINLINE cloneEnv #-}
100148
101149----------------------------------------
102150-- Utils for cloning
103151
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+
104194type EnvRefCache = IORef (M.Map ForkId (IORef EnvRef))
105195
106196-- | Clone the 'EnvRef' and put it in a cache.
107197cloneEnvRef
108− :: EnvRefCache
198+ :: UniqueGen
199+ -> EnvRefCache
109200 -> ForkId
110201 -> IORef EnvRef
111202 -> IO (IORef EnvRef)
112−cloneEnvRef cache fid ref0 = do
203+cloneEnvRef ug cache fid ref0 = do
113204 EnvRef n es0 fs0 <- readIORef ref0
114205 es <- cloneSmallMutableArray es0 0 (sizeofSmallMutableArray es0)
115206 fs <- cloneSmallMutableArray fs0 0 (sizeofSmallMutableArray fs0)
116207 ref <- newIORef $ EnvRef n es fs
117208 modifyIORef' cache $ M.insert fid ref
118− relinkData cache es fs n
209+ relinkData ug cache es fs n
119210 pure ref
120211
121212-- | Relink local environments hiding in the interpreters.
122213relinkData
123− :: EnvRefCache
214+ :: UniqueGen
215+ -> EnvRefCache
124216 -> SmallMutableArray RealWorld Any
125217 -> SmallMutableArray RealWorld Any
126218 -> Int
127219 -> IO ()
128−relinkData cache es fs = \case
220+relinkData ug cache es fs = \case
129221 0 -> pure ()
130222 n -> do
131223 let i = n - 1
@@ -133,59 +225,62 @@ relinkData cache es fs = \case
133225 readSmallArray es i
134226 >>= f relinkEnv . fromAny
135227 >>= writeSmallArray es i . toAny
136− relinkData cache es fs i
228+ relinkData ug cache es fs i
137229 where
138230 relinkEnv :: Env es -> IO (Env es)
139− relinkEnv (Env forks _) = do
231+ relinkEnv (Env forks _ _) = do
140232 Just gref <- M.lookup globalFid <$> readIORef cache
141− Env <$> relinkForks cache forks <*> pure gref
233+ Env <$> relinkForks ug cache forks <*> pure gref <*> pure ug
142234
143−relinkForks :: EnvRefCache -> Forks -> IO Forks
144−relinkForks cache = \case
235+relinkForks :: UniqueGen -> EnvRefCache -> Forks -> IO Forks
236+relinkForks ug cache = \case
145237 NoFork -> pure NoFork
146− Forks fid lref0 forks -> (M.lookup fid <$> readIORef cache) >>= \case
147− Just lref -> Forks fid <$> pure lref
148− <*> relinkForks cache forks
149− Nothing -> Forks fid <$> cloneEnvRef cache fid lref0
150− <*> relinkForks cache forks
238+ Forks fid lref0 forks -> do
239+ readIORef cache >>= pure . M.lookup fid >>= \case
240+ Just lref -> Forks fid <$> pure lref
241+ <*> relinkForks ug cache forks
242+ Nothing -> Forks fid <$> cloneEnvRef ug cache fid lref0
243+ <*> relinkForks ug cache forks
151244
152245----------------------------------------
153246
154247-- | Create a local fork of the environment for interpreters.
155248forkEnv :: HasCallStack => Env es -> IO (Env es)
156−forkEnv env@(Env NoFork gref) = do
249+forkEnv env@(Env NoFork gref ug) = do
157250 size <- sizeEnv env
251+ uniq <- getUnique ug
158252 let fid = ForkId
159253 { baseIx = size
160− , depth = 1
254+ , unique = uniq
161255 }
162256 Env <$> (Forks fid <$> emptyEnvRef <*> pure NoFork)
163− <*> pure gref
164−forkEnv (Env forks@(Forks fid0 lref0 olderForks) gref) = do
257+ <*> pure gref <*> pure ug
258+forkEnv (Env forks@(Forks fid0 lref0 olderForks) gref ug) = do
165259 EnvRef n _ _ <- readIORef lref0
260+ uniq <- getUnique ug
166261 -- If the fork is empty, replace it as no data is lost.
167262 if n == 0
168263 then do
169264 lref <- emptyEnvRef
170265 let fid = fid0
171− { depth = depth fid0 + 1
266+ { unique = uniq
172267 }
173− pure $ Env (Forks fid lref olderForks) gref
268+ pure $ Env (Forks fid lref olderForks) gref ug
174269 else do
175270 lref <- emptyEnvRef
176271 let fid = fid0
177272 { baseIx = baseIx fid0 + n
178− , depth = depth fid0 + 1
273+ , unique = uniq
179274 }
180− pure $ Env (Forks fid lref forks) gref
275+ pure $ Env (Forks fid lref forks) gref ug
181276{-# NOINLINE forkEnv #-}
182277
183278-- | Get the current size of the environment.
184279sizeEnv :: Env es -> IO Int
185−sizeEnv (Env NoFork ref) = do
280+sizeEnv (Env NoFork ref _) = do
186281 EnvRef n _ _ <- readIORef ref
187282 pure n
188−sizeEnv (Env (Forks fid lref _) _) = do
283+sizeEnv (Env (Forks fid lref _) _ _) = do
189284 EnvRef n _ _ <- readIORef lref
190285 pure $ baseIx fid + n
191286{-# NOINLINE sizeEnv #-}
@@ -198,11 +293,11 @@ getEnv env = do
198293
199294-- | Check that the size of the environment is the same as the expected value.
200295checkSizeEnv :: HasCallStack => Int -> Env es -> IO ()
201−checkSizeEnv k (Env NoFork ref) = do
296+checkSizeEnv k (Env NoFork ref _) = do
202297 EnvRef n _ _ <- readIORef ref
203298 when (k /= n) $ do
204299 error $ "k (" ++ show k ++ ") /= n (" ++ show n ++ ")"
205−checkSizeEnv k (Env (Forks fid lref _) _) = do
300+checkSizeEnv k (Env (Forks fid lref _) _ _) = do
206301 EnvRef n _ _ <- readIORef lref
207302 when (k /= baseIx fid + n) $ do
208303 error $ "k (" ++ show k ++ ") /= baseIx + n (baseIx: "
@@ -214,13 +309,13 @@ checkSizeEnv k (Env (Forks fid lref _) _) = do
214309
215310-- | Extend the environment with a new data type (in place).
216311unsafeConsEnv :: 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
218313 NoFork -> do
219314 extendEnvRef gref
220− pure $ Env NoFork gref
315+ pure $ Env NoFork gref ug
221316 Forks base lref forks -> do
222317 extendEnvRef lref
223− pure $ Env (Forks base lref forks) gref
318+ pure $ Env (Forks base lref forks) gref ug
224319 where
225320 extendEnvRef :: IORef EnvRef -> IO ()
226321 extendEnvRef ref = do
@@ -249,13 +344,13 @@ unsafeConsEnv e f (Env fork gref) = case fork of
249344-- | Shrink the environment by one data type (in place). Makes sure the size of
250345-- the environment is as expected.
251346unsafeTailEnv :: 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
253348 NoFork -> do
254349 shrinkEnvRef len gref
255− pure $ Env NoFork gref
350+ pure $ Env NoFork gref ug
256351 Forks fid lref forks -> do
257352 shrinkEnvRef (len - baseIx fid) lref
258− pure $ Env (Forks fid lref forks) gref
353+ pure $ Env (Forks fid lref forks) gref ug
259354 where
260355 shrinkEnvRef :: Int -> IORef EnvRef -> IO ()
261356 shrinkEnvRef k ref = do
@@ -317,7 +412,7 @@ getLocation
317412 :: forall e es. e :> es
318413 => Env es
319414 -> IO (Int, SmallMutableArray RealWorld Any)
320−getLocation (Env fork ref) = do
415+getLocation (Env fork ref _) = do
321416 EnvRef n es _ <- readIORef ref
322417 -- Optimized for the common access pattern:
323418 --