unsafe{Put,State,Modify}Env are safe

effectful-core/src/Effectful/Internal/Env.hs

@@ -9,22 +9,22 @@ module Effectful.Internal.Env
   , Relinker(..)
   , noRelinker
 
-  -- * Safe operations
+  -- * Operations
   , emptyEnv
   , cloneEnv
   , forkEnv
   , sizeEnv
-  , getEnv
   , checkSizeEnv
 
-  -- * Extending and shrinking
+  -- ** Extending and shrinking
   , unsafeConsEnv
   , unsafeTailEnv
 
-  -- * Data retrieval and update
+  -- ** Data retrieval and update
-  , unsafePutEnv
+  , getEnv
-  , unsafeModifyEnv
+  , putEnv
-  , unsafeStateEnv
+  , stateEnv
+  , modifyEnv
   ) where
 
 import Control.Monad
@@ -322,12 +322,6 @@ sizeEnv (Env (Forks _ baseIx lref _) _ _) = do
   pure $ baseIx + n
 {-# NOINLINE sizeEnv #-}
 
--- | Extract a specific data type from the environment.
-getEnv :: forall e es handler. e :> es => Env es -> IO (handler e)
-getEnv env = do
-  (i, es) <- getLocation @e env
-  fromAny <$> readSmallArray es i
-
 -- | Check that the size of the environment is the same as the expected value.
 checkSizeEnv :: Int -> Env es -> IO ()
 checkSizeEnv k (Env NoFork ref _) = do
@@ -345,6 +339,9 @@ checkSizeEnv k (Env (Forks _ baseIx lref _) _ _) = do
 -- Extending and shrinking
 
 -- | Extend the environment with a new data type (in place).
+--
+-- /Note:/ this function is __unsafe__ because it renders the input 'Env'
+-- unusable, but it's not checked anywhere.
 unsafeConsEnv :: handler e -> Relinker handler e -> Env es -> IO (Env (e : es))
 unsafeConsEnv e f (Env fork gref gen) = case fork of
   NoFork -> do
@@ -380,6 +377,9 @@ unsafeConsEnv e f (Env fork gref gen) = case fork of
 
 -- | Shrink the environment by one data type (in place). Makes sure the size of
 -- the environment is as expected.
+--
+-- /Note:/ this function is __unsafe__ because it renders the input 'Env'
+-- unusable, but it's not checked anywhere.
 unsafeTailEnv :: Int -> Env (e : es) -> IO (Env es)
 unsafeTailEnv len (Env fork gref gen) = case fork of
   NoFork -> do
@@ -403,39 +403,45 @@ unsafeTailEnv len (Env fork gref gen) = case fork of
 ----------------------------------------
 -- Data retrieval and update
 
+-- | Extract a specific data type from the environment.
+getEnv :: forall e es handler. e :> es => Env es -> IO (handler e)
+getEnv env = do
+  (i, es) <- getLocation @e env
+  fromAny <$> readSmallArray es i
+
 -- | Replace the data type in the environment with a new value (in place).
-unsafePutEnv
+putEnv
   :: forall e es handler. e :> es
   => Env es
   -> handler e
   -> IO ()
-unsafePutEnv env e = do
+putEnv env e = do
   (i, es) <- getLocation @e env
   e `seq` writeSmallArray es i (toAny e)
 
--- | Modify the data type in the environment (in place).
-unsafeModifyEnv
-  :: forall e es handler. e :> es
-  => Env es
-  -> (handler e -> handler e)
-  -> IO ()
-unsafeModifyEnv env f = do
-  (i, es) <- getLocation @e env
-  e <- f . fromAny <$> readSmallArray es i
-  e `seq` writeSmallArray es i (toAny e)
-
 -- | Modify the data type in the environment (in place) and return a value.
-unsafeStateEnv
+stateEnv
   :: forall e es handler a. e :> es
   => Env es
   -> (handler e -> (a, handler e))
   -> IO a
-unsafeStateEnv env f = do
+stateEnv env f = do
   (i, es) <- getLocation @e env
   (a, e) <- f . fromAny <$> readSmallArray es i
   e `seq` writeSmallArray es i (toAny e)
   pure a
 
+-- | Modify the data type in the environment (in place).
+modifyEnv
+  :: forall e es handler. e :> es
+  => Env es
+  -> (handler e -> handler e)
+  -> IO ()
+modifyEnv env f = do
+  (i, es) <- getLocation @e env
+  e <- f . fromAny <$> readSmallArray es i
+  e `seq` writeSmallArray es i (toAny e)
+
 ----------------------------------------
 -- Internal helpers
 

effectful-core/src/Effectful/Internal/Monad.hs

@@ -432,19 +432,19 @@ getEffect :: e :> es => Eff es (handler e)
 getEffect = unsafeEff $ \es -> getEnv es
 
 putEffect :: e :> es => handler e -> Eff es ()
-putEffect e = unsafeEff $ \es -> unsafePutEnv es e
+putEffect e = unsafeEff $ \es -> putEnv es e
 
 stateEffect :: e :> es => (handler e -> (a, handler e)) -> Eff es a
-stateEffect f = unsafeEff $ \es -> unsafeStateEnv es f
+stateEffect f = unsafeEff $ \es -> stateEnv es f
 
 stateEffectM :: e :> es => (handler e -> Eff es (a, handler e)) -> Eff es a
 stateEffectM f = unsafeEff $ \es -> mask $ \release -> do
   (a, e) <- (\e -> release $ unEff (f e) es) =<< getEnv es
-  unsafePutEnv es e
+  putEnv es e
   pure a
 
 localEffect :: e :> es => (handler e -> handler e) -> Eff es a -> Eff es a
 localEffect f m = unsafeEff $ \es -> do
-  bracket (unsafeStateEnv es $ \e -> (e, f e))
+  bracket (stateEnv es $ \e -> (e, f e))
-          (\e -> unsafePutEnv es e)
+          (\e -> putEnv es e)
           (\_ -> unEff m es)

effectful-core/src/Effectful/Writer/Local.hs

@@ -34,13 +34,13 @@ tell w = stateEffect $ \(IdE (Writer w0)) -> ((), IdE (Writer (w0 <> w)))
 
 listen :: (Writer w :> es, Monoid w) => Eff es a -> Eff es (a, w)
 listen m = unsafeEff $ \es -> mask $ \restore -> do
-  w0 <- unsafeStateEnv es $ \(IdE (Writer w)) -> (w, IdE (Writer mempty))
+  w0 <- stateEnv es $ \(IdE (Writer w)) -> (w, IdE (Writer mempty))
   a <- restore (unEff m es) `onException` merge es w0
   (a, ) <$> merge es w0
   where
     merge es w0 =
       -- If an exception is thrown, restore w0 and keep parts of w1.
-      unsafeStateEnv es $ \(IdE (Writer w1)) -> (w1, IdE (Writer (w0 <> w1)))
+      stateEnv es $ \(IdE (Writer w1)) -> (w1, IdE (Writer (w0 <> w1)))
 
 listens :: (Writer w :> es, Monoid w) => (w -> b) -> Eff es a -> Eff es (a, b)
 listens f m = do

effectful-core/src/Effectful/Writer/Shared.hs

@@ -43,14 +43,14 @@ listen m = unsafeEff $ \es -> do
   uninterruptibleMask $ \restore -> do
     v1 <- newMVar mempty
     -- Replace thread local MVar with a fresh one for isolated listening.
-    v0 <- unsafeStateEnv es $ \(IdE (Writer v)) -> (v, IdE (Writer v1))
+    v0 <- stateEnv es $ \(IdE (Writer v)) -> (v, IdE (Writer v1))
     a <- restore (unEff m es) `onException` merge es v0 v1
     (a, ) <$> merge es v0 v1
   where
     -- Merge results accumulated in the local MVar with the mainline. If an
     -- exception was received while listening, merge results recorded so far.
     merge es v0 v1 = do
-      unsafePutEnv es $ IdE (Writer v0)
+      putEnv es $ IdE (Writer v0)
       w1 <- readMVar v1
       modifyMVar_ v0 $ \w0 -> let w = w0 <> w1 in w `seq` pure w
       pure w1