Try eta expanding FCode (See #18202)

Also updates the note with the case of multi-argument lambdas.

Seems slightly beneficial based on the Cabal test:

-O0: -1MB allocations (out of 50GB)
-O : -1MB allocations (out of ~200GB)

2 files changed, 33 insertions, 2 deletions

diff --git a/compiler/GHC/StgToCmm/Monad.hs b/compiler/GHC/StgToCmm/Monad.hs
index 6731ad2d9a..3066609d7e 100644
--- a/compiler/GHC/StgToCmm/Monad.hs
+++ b/compiler/GHC/StgToCmm/Monad.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE PatternSynonyms #-}
 
 
 -----------------------------------------------------------------------------
@@ -87,6 +88,7 @@ import GHC.Data.FastString
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Utils.Misc
+import GHC.Exts (oneShot)
 
 import Control.Monad
 import Data.List (mapAccumL)
@@ -119,8 +121,26 @@ import Data.List (mapAccumL)
 
 --------------------------------------------------------
 
-newtype FCode a = FCode { doFCode :: CgInfoDownwards -> CgState -> (a, CgState) }
-    deriving (Functor)
+newtype FCode a = FCode' { doFCode :: CgInfoDownwards -> CgState -> (a, CgState) }
+
+-- Not derived because of #18202.
+-- See Note [The one-shot state monad trick] in GHC.Utils.Monad
+instance Functor FCode where
+  fmap f (FCode m) =
+    FCode $ \info_down state ->
+      case m info_down state of
+        (x, state') -> (f x, state')
+
+-- This pattern synonym makes the simplifier monad eta-expand,
+-- which as a very beneficial effect on compiler performance
+-- See #18202.
+-- See Note [The one-shot state monad trick] in GHC.Utils.Monad
+{-# COMPLETE FCode #-}
+pattern FCode :: (CgInfoDownwards -> CgState -> (a, CgState))
+              -> FCode a
+pattern FCode m <- FCode' m
+  where
+    FCode m = FCode' $ oneShot (\cgInfoDown -> oneShot (\state ->m cgInfoDown state))
 
 instance Applicative FCode where
     pure val = FCode (\_info_down state -> (val, state))
diff --git a/compiler/GHC/Utils/Monad.hs b/compiler/GHC/Utils/Monad.hs
index da415ba44c..a65947e59e 100644
--- a/compiler/GHC/Utils/Monad.hs
+++ b/compiler/GHC/Utils/Monad.hs
@@ -344,6 +344,17 @@ it is more elaborate.
 
 The pattern synonym approach is due to Sebastian Graaf (#18238)
 
+Do note that for monads for multiple arguments more than one oneShot
+function might be required. For example in FCode we use:
+
+    newtype FCode a = FCode' { doFCode :: CgInfoDownwards -> CgState -> (a, CgState) }
+
+    pattern FCode :: (CgInfoDownwards -> CgState -> (a, CgState))
+                  -> FCode a
+    pattern FCode m <- FCode' m
+      where
+        FCode m = FCode' $ oneShot (\cgInfoDown -> oneShot (\state ->m cgInfoDown state))
+
 Derived instances
 ~~~~~~~~~~~~~~~~~
 One caveat of both approaches is that derived instances don't use the smart