SimplfyMonad: Mark all lambdas as oneShotwip/simplify-oneshot

The eta expansion previously performed in
3d345c9680ab3d766ef43dd8389ccc1eaeca066c only marked the *first*
(`SimplTopEnv`) argument, since `oneShot` only applies to the lambda at
the head of the enclosed expression. Fix this.

2 files changed, 12 insertions, 1 deletions

diff --git a/compiler/GHC/Core/Opt/Simplify/Monad.hs b/compiler/GHC/Core/Opt/Simplify/Monad.hs
index 620db9da22..b9c1cb98ba 100644
--- a/compiler/GHC/Core/Opt/Simplify/Monad.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Monad.hs
@@ -76,7 +76,7 @@ pattern SM :: (SimplTopEnv -> UniqSupply -> SimplCount
 -- See Note [The one-shot state monad trick] in GHC.Utils.Monad
 pattern SM m <- SM' m
   where
-    SM m = SM' (oneShot m)
+    SM m = SM' $ oneShot $ \env -> oneShot $ \us -> oneShot $ \count -> m env us count
 
 data SimplTopEnv
   = STE { st_flags     :: DynFlags
diff --git a/compiler/GHC/Utils/Monad.hs b/compiler/GHC/Utils/Monad.hs
index da415ba44c..d775f30530 100644
--- a/compiler/GHC/Utils/Monad.hs
+++ b/compiler/GHC/Utils/Monad.hs
@@ -322,6 +322,17 @@ The trick is very similar to the built-in "state hack"
 (see Note [The state-transformer hack] in "GHC.Core.Opt.Arity") but is
 applicable on a monad-by-monad basis under programmer control.
 
+Do note, however, that each 'oneShot' only applies to a single lambda.
+Consequently, it may be necessary to apply it more than once. For instance, in
+the case of a monad that takes multiple arguments, e.g.
+
+    newtype M a =  M (Env -> UniqSupply -> IO (a, UniqSupply))
+
+we have would have the following smart constructor,
+
+    mkM :: Env -> UniqSupply -> IO (a, UniqSupply)
+    mkM env us = M $ oneShot $ \env -> oneShot $ \us -> f env us
+
 Using pattern synonyms
 ~~~~~~~~~~~~~~~~~~~~~~
 Using a smart constructor is fine, but there is no way to check that we