diff --git a/src/Effectful/Internal/Env.hs b/src/Effectful/Internal/Env.hs index 5df078f..edc9f32 100644 --- a/src/Effectful/Internal/Env.hs +++ b/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: --