Support cloning of forked environments

src/Effectful/Internal/Env.hs

@@ -31,6 +31,7 @@ import Control.Monad
 import Control.Monad.Primitive
 import Data.IORef
 import Data.Primitive.SmallArray
+import Data.Word
 import GHC.Exts (Any)
 import GHC.Stack
 import Unsafe.Coerce
@@ -56,17 +57,7 @@ type role Env nominal
 --
 -- - Modification of a specific element: /O(forks)/, usually /O(1)/.
 --
-data Env (es :: [Effect]) = Env Forks (IORef EnvRef)
-
--- | A fork is uniquely determined by its base index and depth.
-data ForkId = ForkId
-  { baseIx :: Int
-  , depth  :: Int
-  } deriving (Eq, Ord, Show)
-
--- | Implicit 'ForkId' of the 'EnvRef' directly available from 'Env'.
-globalFid :: ForkId
-globalFid = ForkId 0 0
+data Env (es :: [Effect]) = Env Forks (IORef EnvRef) UniqueGen
 
 -- | Local forks of the environment.
 data Forks = Forks {-# UNPACK #-} ForkId (IORef EnvRef) Forks | NoFork
@@ -77,6 +68,21 @@ data EnvRef = EnvRef
   , _relinkers :: SmallMutableArray RealWorld Any
   }
 
+----------------------------------------
+-- ForkId
+
+data ForkId = ForkId
+  { baseIx :: Int
+  , unique :: Unique
+  } deriving (Eq, Ord, Show)
+
+-- | Implicit 'ForkId' of the 'EnvRef' directly available from 'Env'.
+globalFid :: ForkId
+globalFid = ForkId 0 (Unique 0)
+
+----------------------------------------
+-- Relinker
+
 newtype Relinker i (e :: Effect) where
   Relinker
     :: ((forall es. Env es -> IO (Env es)) -> i e -> IO (i e))
@@ -85,47 +91,133 @@ newtype Relinker i (e :: Effect) where
 noRelinker :: Relinker i e
 noRelinker = Relinker $ \_ -> pure
 
+----------------------------------------
+-- UniqueGen
+
+newtype Unique = Unique Word64
+  deriving (Eq, Ord, Show)
+
+-- | Uniques for 'ForkId'S.
+newtype UniqueGen = UniqueGen (IORef Unique)
+
+-- | Create a new thread local unique generator.
+newUniqueGen :: IO UniqueGen
+newUniqueGen = UniqueGen <$> newIORef (Unique 1)
+
+-- | Clone the unique generator for use in a different thread.
+cloneUniqueGen :: UniqueGen -> IO UniqueGen
+cloneUniqueGen (UniqueGen ref) = fmap UniqueGen . newIORef =<< readIORef ref
+
+-- | Get a 'Unique' from the generator.
+getUnique :: UniqueGen -> IO Unique
+getUnique (UniqueGen ref) = do
+  unique@(Unique n) <- readIORef ref
+  writeIORef ref $! Unique (n + 1)
+  pure unique
+
+----------------------------------------
+-- Operations
+
 -- | Create an empty environment.
 emptyEnv :: IO (Env '[])
-emptyEnv = Env NoFork <$> emptyEnvRef
+emptyEnv = Env NoFork <$> emptyEnvRef <*> newUniqueGen
 
 -- | Clone the environment.
 cloneEnv :: HasCallStack => Env es -> IO (Env es)
-cloneEnv (Env NoFork gref0) = do
+cloneEnv (Env NoFork gref0 ug0) = do
+  ug <- cloneUniqueGen ug0
   cache <- newIORef M.empty
-  gref <- cloneEnvRef cache globalFid gref0
-  pure $ Env NoFork gref
-cloneEnv _ = error "Cloning of forked enviroment is not supported yet"
+  gref <- cloneEnvRef ug cache globalFid gref0
+  pure $ Env NoFork gref ug
+cloneEnv (Env forks@(Forks fid lref0 _) gref0 ug0) = do
+  EnvRef _ es0 fs0 <- readIORef gref0
+  EnvRef n _   _   <- readIORef lref0
+  let len = baseIx fid + n
+  es  <- newSmallArray len (error "undefined field")
+  fs  <- newSmallArray len (error "undefined field")
+  baseN <- copyForks es fs len forks
+  copySmallMutableArray es 0 es0 0 baseN
+  copySmallMutableArray fs 0 fs0 0 baseN
+  ug <- cloneUniqueGen ug0
+  gref <- newIORef $ EnvRef len es fs
+  cache <- newIORef $ M.singleton globalFid gref
+  relinkData ug cache es fs n
+  -- The forked environment is flattened and becomes the global one.
+  pure $ Env NoFork gref ug
 {-# NOINLINE cloneEnv #-}
 
 ----------------------------------------
 -- Utils for cloning
 
+-- | Let's say that the forked environment looks like this:
+--
+-- [0,1][2,3,4,5][6,7,8,9][10,11]
+--
+-- Then forks will look like this:
+--
+-- Fork (baseIx: 10, unique: ...) [10,11]
+--   (Fork (baseIx: 6, unique: ...) [6,7,8,9]
+--     (Fork (baseIx: 2, unique: ...) [2,3,4,5]
+--       NoFork))
+--
+-- and elements [0,1] are taken from the global environment.
+--
+-- We start with len: 12, we subtract baseIx: 10 and get n: 2, the number of
+-- elements to copy from the fork. We copy them, then call recursively with
+-- len-n (10).
+--
+-- Then len: 10, we subtract baseIx: 6 and get n: 4. the number of elements to
+-- copy from the fork (note that we can't use the capacity from EnvRef, because
+-- it might be longer than 4 if the fork was locally extended). We copy them,
+-- then call recursively with len-n (6).
+--
+-- Then len: 6, we subtract baseIx: 2 and get n: 4. We copy the elements, then
+-- call recursively with len-n (2).
+--
+-- Now len: 2, but there are no more forks left, so we return len (2), the
+-- number of elements to copy from the global environment.
+copyForks
+  :: SmallMutableArray RealWorld Any
+  -> SmallMutableArray RealWorld Any
+  -> Int
+  -> Forks
+  -> IO Int
+copyForks es fs len = \case
+  NoFork -> pure len
+  Forks fid lref0 forks -> do
+    EnvRef _ es0 fs0 <- readIORef lref0
+    let n = len - baseIx fid
+    copySmallMutableArray es (baseIx fid) es0 0 n
+    copySmallMutableArray fs (baseIx fid) fs0 0 n
+    copyForks es fs (len - n) forks
+
 type EnvRefCache = IORef (M.Map ForkId (IORef EnvRef))
 
 -- | Clone the 'EnvRef' and put it in a cache.
 cloneEnvRef
-  :: EnvRefCache
+  :: UniqueGen
+  -> EnvRefCache
   -> ForkId
   -> IORef EnvRef
   -> IO (IORef EnvRef)
-cloneEnvRef cache fid ref0 = do
+cloneEnvRef ug cache fid ref0 = do
   EnvRef n es0 fs0 <- readIORef ref0
   es  <- cloneSmallMutableArray es0 0 (sizeofSmallMutableArray es0)
   fs  <- cloneSmallMutableArray fs0 0 (sizeofSmallMutableArray fs0)
   ref <- newIORef $ EnvRef n es fs
   modifyIORef' cache $ M.insert fid ref
-  relinkData cache es fs n
+  relinkData ug cache es fs n
   pure ref
 
 -- | Relink local environments hiding in the interpreters.
 relinkData
-  :: EnvRefCache
+  :: UniqueGen
+  -> EnvRefCache
   -> SmallMutableArray RealWorld Any
   -> SmallMutableArray RealWorld Any
   -> Int
   -> IO ()
-relinkData cache es fs = \case
+relinkData ug cache es fs = \case
   0 -> pure ()
   n -> do
     let i = n - 1
@@ -133,59 +225,62 @@ relinkData cache es fs = \case
     readSmallArray es i
       >>= f relinkEnv . fromAny
       >>= writeSmallArray es i . toAny
-    relinkData cache es fs i
+    relinkData ug cache es fs i
   where
     relinkEnv :: Env es -> IO (Env es)
-    relinkEnv (Env forks _) = do
+    relinkEnv (Env forks _ _) = do
       Just gref <- M.lookup globalFid <$> readIORef cache
-      Env <$> relinkForks cache forks <*> pure gref
+      Env <$> relinkForks ug cache forks <*> pure gref <*> pure ug
 
-relinkForks :: EnvRefCache -> Forks -> IO Forks
-relinkForks cache = \case
+relinkForks :: UniqueGen -> EnvRefCache -> Forks -> IO Forks
+relinkForks ug cache = \case
   NoFork -> pure NoFork
-  Forks fid lref0 forks -> (M.lookup fid <$> readIORef cache) >>= \case
-    Just lref -> Forks fid <$> pure lref
-                           <*> relinkForks cache forks
-    Nothing   -> Forks fid <$> cloneEnvRef cache fid lref0
-                           <*> relinkForks cache forks
+  Forks fid lref0 forks -> do
+    readIORef cache >>= pure . M.lookup fid >>= \case
+      Just lref -> Forks fid <$> pure lref
+                             <*> relinkForks ug cache forks
+      Nothing   -> Forks fid <$> cloneEnvRef ug cache fid lref0
+                             <*> relinkForks ug cache forks
 
 ----------------------------------------
 
 -- | Create a local fork of the environment for interpreters.
 forkEnv :: HasCallStack => Env es -> IO (Env es)
-forkEnv env@(Env NoFork gref) = do
+forkEnv env@(Env NoFork gref ug) = do
   size <- sizeEnv env
+  uniq <- getUnique ug
   let fid = ForkId
         { baseIx = size
-        , depth  = 1
+        , unique = uniq
         }
   Env <$> (Forks fid <$> emptyEnvRef <*> pure NoFork)
-      <*> pure gref
-forkEnv (Env forks@(Forks fid0 lref0 olderForks) gref) = do
+      <*> pure gref <*> pure ug
+forkEnv (Env forks@(Forks fid0 lref0 olderForks) gref ug) = do
   EnvRef n _ _ <- readIORef lref0
+  uniq <- getUnique ug
   -- If the fork is empty, replace it as no data is lost.
   if n == 0
     then do
       lref <- emptyEnvRef
       let fid = fid0
-            { depth = depth fid0 + 1
+            { unique = uniq
             }
-      pure $ Env (Forks fid lref olderForks) gref
+      pure $ Env (Forks fid lref olderForks) gref ug
     else do
       lref <- emptyEnvRef
       let fid = fid0
             { baseIx = baseIx fid0 + n
-            , depth  = depth  fid0 + 1
+            , unique = uniq
             }
-      pure $ Env (Forks fid lref forks) gref
+      pure $ Env (Forks fid lref forks) gref ug
 {-# NOINLINE forkEnv #-}
 
 -- | Get the current size of the environment.
 sizeEnv :: Env es -> IO Int
-sizeEnv (Env NoFork ref) = do
+sizeEnv (Env NoFork ref _) = do
   EnvRef n _ _ <- readIORef ref
   pure n
-sizeEnv (Env (Forks fid lref _) _) = do
+sizeEnv (Env (Forks fid lref _) _ _) = do
   EnvRef n _ _ <- readIORef lref
   pure $ baseIx fid + n
 {-# NOINLINE sizeEnv #-}
@@ -198,11 +293,11 @@ getEnv env = do
 
 -- | Check that the size of the environment is the same as the expected value.
 checkSizeEnv :: HasCallStack => Int -> Env es -> IO ()
-checkSizeEnv k (Env NoFork ref) = do
+checkSizeEnv k (Env NoFork ref _) = do
   EnvRef n _ _ <- readIORef ref
   when (k /= n) $ do
     error $ "k (" ++ show k ++ ") /= n (" ++ show n ++ ")"
-checkSizeEnv k (Env (Forks fid lref _) _) = do
+checkSizeEnv k (Env (Forks fid lref _) _ _) = do
   EnvRef n _ _ <- readIORef lref
   when (k /= baseIx fid + n) $ do
     error $ "k (" ++ show k ++ ") /= baseIx + n (baseIx: "
@@ -214,13 +309,13 @@ checkSizeEnv k (Env (Forks fid lref _) _) = do
 
 -- | Extend the environment with a new data type (in place).
 unsafeConsEnv :: HasCallStack => i e -> Relinker i e -> Env es -> IO (Env (e : es))
-unsafeConsEnv e f (Env fork gref) = case fork of
+unsafeConsEnv e f (Env fork gref ug) = case fork of
   NoFork -> do
     extendEnvRef gref
-    pure $ Env NoFork gref
+    pure $ Env NoFork gref ug
   Forks base lref forks -> do
     extendEnvRef lref
-    pure $ Env (Forks base lref forks) gref
+    pure $ Env (Forks base lref forks) gref ug
   where
     extendEnvRef :: IORef EnvRef -> IO ()
     extendEnvRef ref = do
@@ -249,13 +344,13 @@ unsafeConsEnv e f (Env fork gref) = case fork of
 -- | Shrink the environment by one data type (in place). Makes sure the size of
 -- the environment is as expected.
 unsafeTailEnv :: HasCallStack => Int -> Env (e : es) -> IO (Env es)
-unsafeTailEnv len (Env fork gref) = case fork of
+unsafeTailEnv len (Env fork gref ug) = case fork of
   NoFork -> do
     shrinkEnvRef len gref
-    pure $ Env NoFork gref
+    pure $ Env NoFork gref ug
   Forks fid lref forks -> do
     shrinkEnvRef (len - baseIx fid) lref
-    pure $ Env (Forks fid lref forks) gref
+    pure $ Env (Forks fid lref forks) gref ug
   where
     shrinkEnvRef :: Int -> IORef EnvRef -> IO ()
     shrinkEnvRef k ref = do
@@ -317,7 +412,7 @@ getLocation
   :: forall e es. e :> es
   => Env es
   -> IO (Int, SmallMutableArray RealWorld Any)
-getLocation (Env fork ref) = do
+getLocation (Env fork ref _) = do
   EnvRef n es _ <- readIORef ref
   -- Optimized for the common access pattern:
   --