UniqSM: oneShot-ify

Part of #18202

-------------------------
Metric Decrease:
    T12707
    T3294
-------------------------

1 files changed, 21 insertions, 10 deletions

diff --git a/compiler/GHC/Types/Unique/Supply.hs b/compiler/GHC/Types/Unique/Supply.hs
index 7d6c4914e2..9f4d6744c9 100644
--- a/compiler/GHC/Types/Unique/Supply.hs
+++ b/compiler/GHC/Types/Unique/Supply.hs
@@ -42,7 +42,7 @@ import GHC.Utils.Monad
 import Control.Monad
 import Data.Bits
 import Data.Char
-import GHC.Exts( Ptr(..), noDuplicate# )
+import GHC.Exts( Ptr(..), noDuplicate#, oneShot )
 #if MIN_VERSION_GLASGOW_HASKELL(9,1,0,0)
 import GHC.Exts( Int(..), word2Int#, fetchAddWordAddr#, plusWord#, readWordOffAddr# )
 #if defined(DEBUG)
@@ -297,7 +297,18 @@ pattern UniqResult x y = (# x, y #)
 
 -- | A monad which just gives the ability to obtain 'Unique's
 newtype UniqSM result = USM { unUSM :: UniqSupply -> UniqResult result }
-    deriving (Functor)
+
+-- See Note [The one-shot state monad trick] for why we don't derive this.
+instance Functor UniqSM where
+  fmap f (USM m) = mkUniqSM $ \us ->
+      case m us of
+        (# r, us' #) -> UniqResult (f r) us'
+
+-- | Smart constructor for 'UniqSM', as described in Note [The one-shot state
+-- monad trick].
+mkUniqSM :: (UniqSupply -> UniqResult a) -> UniqSM a
+mkUniqSM f = USM (oneShot f)
+{-# INLINE mkUniqSM #-}
 
 instance Monad UniqSM where
   (>>=) = thenUs
@@ -305,7 +316,7 @@ instance Monad UniqSM where
 
 instance Applicative UniqSM where
     pure = returnUs
-    (USM f) <*> (USM x) = USM $ \us0 -> case f us0 of
+    (USM f) <*> (USM x) = mkUniqSM $ \us0 -> case f us0 of
                             UniqResult ff us1 -> case x us1 of
                               UniqResult xx us2 -> UniqResult (ff xx) us2
     (*>) = thenUs_
@@ -332,22 +343,22 @@ liftUSM :: UniqSM a -> UniqSupply -> (a, UniqSupply)
 liftUSM (USM m) us0 = case m us0 of UniqResult a us1 -> (a, us1)
 
 instance MonadFix UniqSM where
-    mfix m = USM (\us0 -> let (r,us1) = liftUSM (m r) us0 in UniqResult r us1)
+    mfix m = mkUniqSM (\us0 -> let (r,us1) = liftUSM (m r) us0 in UniqResult r us1)
 
 thenUs :: UniqSM a -> (a -> UniqSM b) -> UniqSM b
 thenUs (USM expr) cont
-  = USM (\us0 -> case (expr us0) of
+  = mkUniqSM (\us0 -> case (expr us0) of
                    UniqResult result us1 -> unUSM (cont result) us1)
 
 thenUs_ :: UniqSM a -> UniqSM b -> UniqSM b
 thenUs_ (USM expr) (USM cont)
-  = USM (\us0 -> case (expr us0) of { UniqResult _ us1 -> cont us1 })
+  = mkUniqSM (\us0 -> case (expr us0) of { UniqResult _ us1 -> cont us1 })
 
 returnUs :: a -> UniqSM a
-returnUs result = USM (\us -> UniqResult result us)
+returnUs result = mkUniqSM (\us -> UniqResult result us)
 
 getUs :: UniqSM UniqSupply
-getUs = USM (\us0 -> case splitUniqSupply us0 of (us1,us2) -> UniqResult us1 us2)
+getUs = mkUniqSM (\us0 -> case splitUniqSupply us0 of (us1,us2) -> UniqResult us1 us2)
 
 -- | A monad for generating unique identifiers
 class Monad m => MonadUnique m where
@@ -371,9 +382,9 @@ instance MonadUnique UniqSM where
     getUniquesM = getUniquesUs
 
 getUniqueUs :: UniqSM Unique
-getUniqueUs = USM (\us0 -> case takeUniqFromSupply us0 of
+getUniqueUs = mkUniqSM (\us0 -> case takeUniqFromSupply us0 of
                            (u,us1) -> UniqResult u us1)
 
 getUniquesUs :: UniqSM [Unique]
-getUniquesUs = USM (\us0 -> case splitUniqSupply us0 of
+getUniquesUs = mkUniqSM (\us0 -> case splitUniqSupply us0 of
                             (us1,us2) -> UniqResult (uniqsFromSupply us1) us2)