Consider array reallocations when doing monadic state updates (#238)

Fixes #237.

effectful-core/CHANGELOG.md

@@ -16,6 +16,8 @@
     `throwError_` for them.
   - `ThrowError` operation from the dynamic version of the `Error` effect was
     replaced with `ThrowErrorWith`.
+  - `stateEnv` and `modifyEnv` now take pure modification functions. Use a
+    combination of `getEnv` and `putEnv` for forward compatibility.
 
 # effectful-core-2.3.1.0 (2024-06-07)
 * Drop support for GHC 8.8.

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

@@ -321,23 +321,23 @@ putEnv env e = do
 stateEnv
   :: forall e es a. e :> es
   => Env es -- ^ The environment.
-  -> (EffectRep (DispatchOf e) e -> IO (a, EffectRep (DispatchOf e) e))
+  -> (EffectRep (DispatchOf e) e -> (a, EffectRep (DispatchOf e) e))
   -> IO a
 stateEnv env f = do
   (i, es) <- getLocation @e env
-  (a, e) <- f . fromAny =<< readSmallArray es i
+  (a, e) <- f . fromAny <$> readSmallArray es i
   writeSmallArray' es i (toAny e)
   pure a
 
 -- | Modify the data type in the environment (in place).
 modifyEnv
   :: forall e es. e :> es
   => Env es -- ^ The environment.
-  -> (EffectRep (DispatchOf e) e -> IO (EffectRep (DispatchOf e) e))
+  -> (EffectRep (DispatchOf e) e -> (EffectRep (DispatchOf e) e))
   -> IO ()
 modifyEnv env f = do
   (i, es) <- getLocation @e env
-  e <- f . fromAny =<< readSmallArray es i
+  e <- f . fromAny <$> readSmallArray es i
   writeSmallArray' es i (toAny e)
 
 -- | Determine location of the effect in the environment.

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

@@ -604,7 +604,7 @@ stateStaticRep
   => (StaticRep e -> (a, StaticRep e))
   -- ^ The function to modify the representation.
   -> Eff es a
-stateStaticRep f = unsafeEff $ \es -> stateEnv es (pure . f)
+stateStaticRep f = unsafeEff $ \es -> stateEnv es f
 
 -- | Apply the monadic function to the current representation of the effect and
 -- return a value.
@@ -614,7 +614,9 @@ stateStaticRepM
   -- ^ The function to modify the representation.
   -> Eff es a
 stateStaticRepM f = unsafeEff $ \es -> E.mask $ \unmask -> do
-  stateEnv es $ unmask . (`unEff` es) . f
+  (a, e) <- unmask . (`unEff` es) . f =<< getEnv es
+  putEnv es e
+  pure a
 
 -- | Execute a computation with a temporarily modified representation of the
 -- effect.
@@ -626,6 +628,6 @@ localStaticRep
   -> Eff es a
 localStaticRep f m = unsafeEff $ \es -> do
   inlineBracket
-    (stateEnv es $ \s -> pure (s, f s))
+    (stateEnv es $ \s -> (s, f s))
     (\s -> putEnv es s)
     (\_ -> unEff m es)

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

@@ -78,13 +78,13 @@ tell w = stateStaticRep $ \(Writer w0) -> ((), Writer (w0 <> w))
 -- "Hi there!"
 listen :: (Writer w :> es, Monoid w) => Eff es a -> Eff es (a, w)
 listen m = unsafeEff $ \es -> mask $ \unmask -> do
-  w0 <- stateEnv es $ \(Writer w) -> pure (w, Writer mempty)
+  w0 <- stateEnv es $ \(Writer w) -> (w, Writer mempty)
   a <- unmask (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.
-      stateEnv es $ \(Writer w1) -> pure (w1, Writer (w0 <> w1))
+      stateEnv es $ \(Writer w1) -> (w1, Writer (w0 <> w1))
 
 -- | Execute an action and append its output to the overall output of the
 -- 'Writer', then return the final value along with a function of the recorded

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

@@ -88,7 +88,7 @@ listen m = unsafeEff $ \es -> do
   uninterruptibleMask $ \unmask -> do
     v1 <- newMVar' mempty
     -- Replace thread local MVar with a fresh one for isolated listening.
-    v0 <- stateEnv es $ \(Writer v) -> pure (v, Writer v1)
+    v0 <- stateEnv es $ \(Writer v) -> (v, Writer v1)
     a <- unmask (unEff m es) `onException` merge es v0 v1
     (a, ) <$> merge es v0 v1
   where

effectful/tests/StateTests.hs

@@ -1,13 +1,17 @@
 module StateTests (stateTests) where
 
 import Control.Exception.Lifted qualified as LE
+import Control.Monad
 import Control.Monad.Catch qualified as E
 import Test.Tasty
 import Test.Tasty.HUnit
 import UnliftIO.Exception qualified as UE
 
 import Effectful
 import Effectful.Dispatch.Dynamic
+import Effectful.Dispatch.Static
+import Effectful.Internal.Env
+import Effectful.Internal.Utils
 import Effectful.State.Static.Local
 import Utils qualified as U
 
@@ -34,9 +38,19 @@ test_evalState = runEff $ do
 
 test_stateM :: Assertion
 test_stateM = runEff $ do
-  (a, b) <- runState "hi" . stateM $ \s -> pure (s, s ++ "!!!")
+  (a, b) <- runState "hi" $ do
+    stateM $ \s -> do
+      effs0 <- getEffectReps
+      -- Trigger reallocation of the internal array in Storage.
+      _ <- evalState () $ pure ()
+      effs1 <- getEffectReps
+      when (effs0 == effs1) $ do
+        U.assertFailure "Internal array was not reallocated"
+      pure (s, s ++ "!!!")
   U.assertEqual "correct a" "hi"    a
   U.assertEqual "correct b" "hi!!!" b
+  where
+    getEffectReps = unsafeEff $ \es -> stEffects <$> readIORef' (envStorage es)
 
 test_deepStack :: Assertion
 test_deepStack = runEff $ do