CmmToAsm.Reg.Linear: oneShot-ify RegMwip/ncg-perf

-------------------------
Metric Decrease:
    T783
    T4801
    T12707
    T13379
    T3294
    T4801
    T5321FD
-------------------------

1 files changed, 22 insertions, 16 deletions

diff --git a/compiler/GHC/CmmToAsm/Reg/Linear/State.hs b/compiler/GHC/CmmToAsm/Reg/Linear/State.hs
index b36270f3bc..4fdc5c96cf 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear/State.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear/State.hs
@@ -50,6 +50,7 @@ import GHC.Cmm.BlockId
 import GHC.Platform
 import GHC.Types.Unique
 import GHC.Types.Unique.Supply
+import GHC.Exts (oneShot)
 
 import Control.Monad (ap)
 
@@ -64,16 +65,21 @@ newtype RegM freeRegs a
         = RegM { unReg :: RA_State freeRegs -> RA_Result freeRegs a }
         deriving (Functor)
 
+-- | Smart constructor for 'RegM', as described in Note [The one-shot state
+-- monad trick] in GHC.Utils.Monad.
+mkRegM :: (RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
+mkRegM f = RegM (oneShot f)
+
 instance Applicative (RegM freeRegs) where
-      pure a  =  RegM $ \s -> RA_Result s a
+      pure a  =  mkRegM $ \s -> RA_Result s a
       (<*>) = ap
 
 instance Monad (RegM freeRegs) where
-  m >>= k   =  RegM $ \s -> case unReg m s of { RA_Result s a -> unReg (k a) s }
+  m >>= k   =  mkRegM $ \s -> case unReg m s of { RA_Result s a -> unReg (k a) s }
 
 -- | Get native code generator configuration
 getConfig :: RegM a NCGConfig
-getConfig = RegM $ \s -> RA_Result s (ra_config s)
+getConfig = mkRegM $ \s -> RA_Result s (ra_config s)
 
 -- | Get target platform from native code generator configuration
 getPlatform :: RegM a Platform
@@ -117,7 +123,7 @@ makeRAStats state
 spillR :: Instruction instr
        => Reg -> Unique -> RegM freeRegs (instr, Int)
 
-spillR reg temp = RegM $ \s ->
+spillR reg temp = mkRegM $ \s ->
   let (stack1,slot) = getStackSlotFor (ra_stack s) temp
       instr  = mkSpillInstr (ra_config s) reg (ra_delta s) slot
   in
@@ -127,42 +133,42 @@ spillR reg temp = RegM $ \s ->
 loadR :: Instruction instr
       => Reg -> Int -> RegM freeRegs instr
 
-loadR reg slot = RegM $ \s ->
+loadR reg slot = mkRegM $ \s ->
   RA_Result s (mkLoadInstr (ra_config s) reg (ra_delta s) slot)
 
 getFreeRegsR :: RegM freeRegs freeRegs
-getFreeRegsR = RegM $ \ s@RA_State{ra_freeregs = freeregs} ->
+getFreeRegsR = mkRegM $ \ s@RA_State{ra_freeregs = freeregs} ->
   RA_Result s freeregs
 
 setFreeRegsR :: freeRegs -> RegM freeRegs ()
-setFreeRegsR regs = RegM $ \ s ->
+setFreeRegsR regs = mkRegM $ \ s ->
   RA_Result s{ra_freeregs = regs} ()
 
 getAssigR :: RegM freeRegs (RegMap Loc)
-getAssigR = RegM $ \ s@RA_State{ra_assig = assig} ->
+getAssigR = mkRegM $ \ s@RA_State{ra_assig = assig} ->
   RA_Result s assig
 
 setAssigR :: RegMap Loc -> RegM freeRegs ()
-setAssigR assig = RegM $ \ s ->
+setAssigR assig = mkRegM $ \ s ->
   RA_Result s{ra_assig=assig} ()
 
 getBlockAssigR :: RegM freeRegs (BlockAssignment freeRegs)
-getBlockAssigR = RegM $ \ s@RA_State{ra_blockassig = assig} ->
+getBlockAssigR = mkRegM $ \ s@RA_State{ra_blockassig = assig} ->
   RA_Result s assig
 
 setBlockAssigR :: BlockAssignment freeRegs -> RegM freeRegs ()
-setBlockAssigR assig = RegM $ \ s ->
+setBlockAssigR assig = mkRegM $ \ s ->
   RA_Result s{ra_blockassig = assig} ()
 
 setDeltaR :: Int -> RegM freeRegs ()
-setDeltaR n = RegM $ \ s ->
+setDeltaR n = mkRegM $ \ s ->
   RA_Result s{ra_delta = n} ()
 
 getDeltaR :: RegM freeRegs Int
-getDeltaR = RegM $ \s -> RA_Result s (ra_delta s)
+getDeltaR = mkRegM $ \s -> RA_Result s (ra_delta s)
 
 getUniqueR :: RegM freeRegs Unique
-getUniqueR = RegM $ \s ->
+getUniqueR = mkRegM $ \s ->
   case takeUniqFromSupply (ra_us s) of
     (uniq, us) -> RA_Result s{ra_us = us} uniq
 
@@ -170,9 +176,9 @@ getUniqueR = RegM $ \s ->
 -- | Record that a spill instruction was inserted, for profiling.
 recordSpill :: SpillReason -> RegM freeRegs ()
 recordSpill spill
-    = RegM $ \s -> RA_Result (s { ra_spills = spill : ra_spills s }) ()
+    = mkRegM $ \s -> RA_Result (s { ra_spills = spill : ra_spills s }) ()
 
 -- | Record a created fixup block
 recordFixupBlock :: BlockId -> BlockId -> BlockId -> RegM freeRegs ()
 recordFixupBlock from between to
-    = RegM $ \s -> RA_Result (s { ra_fixups = (from,between,to) : ra_fixups s }) ()
+    = mkRegM $ \s -> RA_Result (s { ra_fixups = (from,between,to) : ra_fixups s }) ()