GHC.Tc.Solver.Rewrite: oneShot-ify

Following the example of Note [The one-shot state monad trick].
c.f. #18202.

1 files changed, 14 insertions, 7 deletions

diff --git a/compiler/GHC/Tc/Solver/Rewrite.hs b/compiler/GHC/Tc/Solver/Rewrite.hs
index 78b32bec15..76500f0519 100644
--- a/compiler/GHC/Tc/Solver/Rewrite.hs
+++ b/compiler/GHC/Tc/Solver/Rewrite.hs
@@ -32,6 +32,7 @@ import GHC.Tc.Solver.Monad as TcS
 
 import GHC.Utils.Misc
 import GHC.Data.Maybe
+import GHC.Exts (oneShot)
 import Control.Monad
 import GHC.Utils.Monad ( zipWith3M )
 import Data.List.NonEmpty ( NonEmpty(..) )
@@ -58,13 +59,19 @@ newtype RewriteM a
   = RewriteM { runRewriteM :: RewriteEnv -> TcS a }
   deriving (Functor)
 
+-- | Smart constructor for 'RewriteM', as describe in Note [The one-shot state
+-- monad trick] in "GHC.Utils.Monad".
+mkRewriteM :: (RewriteEnv -> TcS a) -> RewriteM a
+mkRewriteM f = RewriteM (oneShot f)
+{-# INLINE mkRewriteM #-}
+
 instance Monad RewriteM where
-  m >>= k  = RewriteM $ \env ->
+  m >>= k  = mkRewriteM $ \env ->
              do { a  <- runRewriteM m env
                 ; runRewriteM (k a) env }
 
 instance Applicative RewriteM where
-  pure x = RewriteM $ const (pure x)
+  pure x = mkRewriteM $ \_ -> pure x
   (<*>) = ap
 
 instance HasDynFlags RewriteM where
@@ -72,7 +79,7 @@ instance HasDynFlags RewriteM where
 
 liftTcS :: TcS a -> RewriteM a
 liftTcS thing_inside
-  = RewriteM $ const thing_inside
+  = mkRewriteM $ \_ -> thing_inside
 
 -- convenient wrapper when you have a CtEvidence describing
 -- the rewriting operation
@@ -95,7 +102,7 @@ traceRewriteM herald doc = liftTcS $ traceTcS herald doc
 
 getRewriteEnvField :: (RewriteEnv -> a) -> RewriteM a
 getRewriteEnvField accessor
-  = RewriteM $ \env -> return (accessor env)
+  = mkRewriteM $ \env -> return (accessor env)
 
 getEqRel :: RewriteM EqRel
 getEqRel = getRewriteEnvField fe_eq_rel
@@ -123,7 +130,7 @@ checkStackDepth ty
 -- | Change the 'EqRel' in a 'RewriteM'.
 setEqRel :: EqRel -> RewriteM a -> RewriteM a
 setEqRel new_eq_rel thing_inside
-  = RewriteM $ \env ->
+  = mkRewriteM $ \env ->
     if new_eq_rel == fe_eq_rel env
     then runRewriteM thing_inside env
     else runRewriteM thing_inside (env { fe_eq_rel = new_eq_rel })
@@ -134,7 +141,7 @@ setEqRel new_eq_rel thing_inside
 -- Note [No derived kind equalities]
 noBogusCoercions :: RewriteM a -> RewriteM a
 noBogusCoercions thing_inside
-  = RewriteM $ \env ->
+  = mkRewriteM $ \env ->
     -- No new thunk is made if the flavour hasn't changed (note the bang).
     let !env' = case fe_flavour env of
           Derived -> env { fe_flavour = Wanted WDeriv }
@@ -144,7 +151,7 @@ noBogusCoercions thing_inside
 
 bumpDepth :: RewriteM a -> RewriteM a
 bumpDepth (RewriteM thing_inside)
-  = RewriteM $ \env -> do
+  = mkRewriteM $ \env -> do
       -- bumpDepth can be called a lot during rewriting so we force the
       -- new env to avoid accumulating thunks.
       { let !env' = env { fe_loc = bumpCtLocDepth (fe_loc env) }