From 34a3e574208ef5db78d8dab116550c2684e94480 Mon Sep 17 00:00:00 2001
From: Andrzej Rybczak <andrzej@rybczak.net>
Date: Mon, 8 Aug 2022 23:20:27 +0200
Subject: [PATCH] Slightly smaller stateStaticRepM

---
 effectful-core/src/Effectful/Internal/Env.hs         | 8 ++++----
 effectful-core/src/Effectful/Internal/Monad.hs       | 8 +++-----
 effectful-core/src/Effectful/Writer/Static/Local.hs  | 4 ++--
 effectful-core/src/Effectful/Writer/Static/Shared.hs | 2 +-
 4 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/effectful-core/src/Effectful/Internal/Env.hs b/effectful-core/src/Effectful/Internal/Env.hs
index 3e4ea0b..4986d40 100644
--- a/effectful-core/src/Effectful/Internal/Env.hs
+++ b/effectful-core/src/Effectful/Internal/Env.hs
@@ -292,11 +292,11 @@ putEnv env e = do
 stateEnv
   :: forall e es a. e :> es
   => Env es -- ^ The environment.
-  -> (EffectRep (DispatchOf e) e -> (a, EffectRep (DispatchOf e) e))
+  -> (EffectRep (DispatchOf e) e -> IO (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
   e `seq` writeSmallArray es i (toAny e)
   pure a
 
@@ -304,11 +304,11 @@ stateEnv env f = do
 modifyEnv
   :: forall e es. e :> es
   => Env es -- ^ The environment.
-  -> (EffectRep (DispatchOf e) e -> EffectRep (DispatchOf e) e)
+  -> (EffectRep (DispatchOf e) e -> IO (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
   e `seq` writeSmallArray es i (toAny e)
 
 -- | Determine location of the effect in the environment.
diff --git a/effectful-core/src/Effectful/Internal/Monad.hs b/effectful-core/src/Effectful/Internal/Monad.hs
index 286f334..b9cb395 100644
--- a/effectful-core/src/Effectful/Internal/Monad.hs
+++ b/effectful-core/src/Effectful/Internal/Monad.hs
@@ -474,7 +474,7 @@ stateStaticRep
   => (StaticRep e -> (a, StaticRep e))
   -- ^ The function to modify the representation.
   -> Eff es a
-stateStaticRep f = unsafeEff $ \es -> stateEnv es f
+stateStaticRep f = unsafeEff $ \es -> stateEnv es (pure . f)
 
 -- | Apply the monadic function to the current representation of the effect and
 -- return a value.
@@ -484,9 +484,7 @@ stateStaticRepM
   -- ^ The function to modify the representation.
   -> Eff es a
 stateStaticRepM f = unsafeEff $ \es -> E.mask $ \unmask -> do
-  (a, s) <- (\s0 -> unmask $ unEff (f s0) es) =<< getEnv es
-  putEnv es s
-  pure a
+  stateEnv es $ unmask . (`unEff` es) . f
 
 -- | Execute a computation with a temporarily modified representation of the
 -- effect.
@@ -497,6 +495,6 @@ localStaticRep
   -> Eff es a
   -> Eff es a
 localStaticRep f m = unsafeEff $ \es -> do
-  E.bracket (stateEnv es $ \s -> (s, f s))
+  E.bracket (stateEnv es $ \s -> pure (s, f s))
             (\s -> putEnv es s)
             (\_ -> unEff m es)
diff --git a/effectful-core/src/Effectful/Writer/Static/Local.hs b/effectful-core/src/Effectful/Writer/Static/Local.hs
index 976cda9..49943da 100644
--- a/effectful-core/src/Effectful/Writer/Static/Local.hs
+++ b/effectful-core/src/Effectful/Writer/Static/Local.hs
@@ -77,13 +77,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) -> (w, Writer mempty)
+  w0 <- stateEnv es $ \(Writer w) -> pure (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) -> (w1, Writer (w0 <> w1))
+      stateEnv es $ \(Writer w1) -> pure (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
diff --git a/effectful-core/src/Effectful/Writer/Static/Shared.hs b/effectful-core/src/Effectful/Writer/Static/Shared.hs
index 4c17d16..8466ed0 100644
--- a/effectful-core/src/Effectful/Writer/Static/Shared.hs
+++ b/effectful-core/src/Effectful/Writer/Static/Shared.hs
@@ -87,7 +87,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) -> (v, Writer v1)
+    v0 <- stateEnv es $ \(Writer v) -> pure (v, Writer v1)
     a <- unmask (unEff m es) `onException` merge es v0 v1
     (a, ) <$> merge es v0 v1
   where