Optimize MonadUnique instances based on IO (#16843)

Metric Decrease:
    T14683

11 files changed, 79 insertions, 121 deletions

diff --git a/compiler/basicTypes/UniqSupply.hs b/compiler/basicTypes/UniqSupply.hs
index 7e87315212..fb321454aa 100644
--- a/compiler/basicTypes/UniqSupply.hs
+++ b/compiler/basicTypes/UniqSupply.hs
@@ -18,7 +18,7 @@ module UniqSupply (
 
         -- ** Operations on supplies
         uniqFromSupply, uniqsFromSupply, -- basic ops
-        takeUniqFromSupply,
+        takeUniqFromSupply, uniqFromMask,
 
         mkSplitUniqSupply,
         splitUniqSupply, listSplitUniqSupply,
@@ -84,6 +84,11 @@ uniqsFromSupply :: UniqSupply -> [Unique] -- Infinite
 takeUniqFromSupply :: UniqSupply -> (Unique, UniqSupply)
 -- ^ Obtain the 'Unique' from this particular 'UniqSupply', and a new supply
 
+uniqFromMask :: Char -> IO Unique
+uniqFromMask mask
+  = do { uqNum <- genSym
+       ; return $! mkUnique mask uqNum }
+
 mkSplitUniqSupply c
   = case ord c `shiftL` uNIQUE_BITS of
      !mask -> let
diff --git a/compiler/coreSyn/CoreLint.hs b/compiler/coreSyn/CoreLint.hs
index 56921ac434..def51f5010 100644
--- a/compiler/coreSyn/CoreLint.hs
+++ b/compiler/coreSyn/CoreLint.hs
@@ -63,7 +63,6 @@ import FastString
 import Util
 import InstEnv     ( instanceDFunId )
 import OptCoercion ( checkAxInstCo )
-import UniqSupply
 import CoreArity ( typeArity )
 import Demand ( splitStrictSig, isBotRes )
 
@@ -2778,8 +2777,9 @@ withoutAnnots pass guts = do
   dflags <- getDynFlags
   let removeFlag env = env{ hsc_dflags = dflags{ debugLevel = 0} }
       withoutFlag corem =
+          -- TODO: supply tag here as well ?
         liftIO =<< runCoreM <$> fmap removeFlag getHscEnv <*> getRuleBase <*>
-                                getUniqueSupplyM <*> getModule <*>
+                                getUniqMask <*> getModule <*>
                                 getVisibleOrphanMods <*>
                                 getPrintUnqualified <*> getSrcSpanM <*>
                                 pure corem
diff --git a/compiler/llvmGen/LlvmCodeGen.hs b/compiler/llvmGen/LlvmCodeGen.hs
index ad6aebeb5a..b463693a82 100644
--- a/compiler/llvmGen/LlvmCodeGen.hs
+++ b/compiler/llvmGen/LlvmCodeGen.hs
@@ -42,10 +42,10 @@ import System.IO
 -- -----------------------------------------------------------------------------
 -- | Top-level of the LLVM Code generator
 --
-llvmCodeGen :: DynFlags -> Handle -> UniqSupply
+llvmCodeGen :: DynFlags -> Handle
                -> Stream.Stream IO RawCmmGroup a
                -> IO a
-llvmCodeGen dflags h us cmm_stream
+llvmCodeGen dflags h cmm_stream
   = withTiming dflags (text "LLVM CodeGen") (const ()) $ do
        bufh <- newBufHandle h
 
@@ -72,7 +72,7 @@ llvmCodeGen dflags h us cmm_stream
            "You are using LLVM version: " <> text (llvmVersionStr ver)
 
        -- run code generation
-       a <- runLlvm dflags (fromMaybe supportedLlvmVersion mb_ver) bufh us $
+       a <- runLlvm dflags (fromMaybe supportedLlvmVersion mb_ver) bufh $
          llvmCodeGen' (liftStream cmm_stream)
 
        bFlush bufh
diff --git a/compiler/llvmGen/LlvmCodeGen/Base.hs b/compiler/llvmGen/LlvmCodeGen/Base.hs
index b132a1b023..eaa49fc50e 100644
--- a/compiler/llvmGen/LlvmCodeGen/Base.hs
+++ b/compiler/llvmGen/LlvmCodeGen/Base.hs
@@ -218,7 +218,7 @@ data LlvmEnv = LlvmEnv
   { envVersion :: LlvmVersion      -- ^ LLVM version
   , envDynFlags :: DynFlags        -- ^ Dynamic flags
   , envOutput :: BufHandle         -- ^ Output buffer
-  , envUniq :: UniqSupply          -- ^ Supply of unique values
+  , envMask :: !Char               -- ^ Mask for creating unique values
   , envFreshMeta :: MetaId         -- ^ Supply of fresh metadata IDs
   , envUniqMeta :: UniqFM MetaId   -- ^ Global metadata nodes
   , envFunMap :: LlvmEnvMap        -- ^ Global functions so far, with type
@@ -249,16 +249,12 @@ instance HasDynFlags LlvmM where
 
 instance MonadUnique LlvmM where
     getUniqueSupplyM = do
-        us <- getEnv envUniq
-        let (us1, us2) = splitUniqSupply us
-        modifyEnv (\s -> s { envUniq = us2 })
-        return us1
+        mask <- getEnv envMask
+        liftIO $! mkSplitUniqSupply mask
 
     getUniqueM = do
-        us <- getEnv envUniq
-        let (u,us') = takeUniqFromSupply us
-        modifyEnv (\s -> s { envUniq = us' })
-        return u
+        mask <- getEnv envMask
+        liftIO $! uniqFromMask mask
 
 -- | Lifting of IO actions. Not exported, as we want to encapsulate IO.
 liftIO :: IO a -> LlvmM a
@@ -266,8 +262,8 @@ liftIO m = LlvmM $ \env -> do x <- m
                               return (x, env)
 
 -- | Get initial Llvm environment.
-runLlvm :: DynFlags -> LlvmVersion -> BufHandle -> UniqSupply -> LlvmM a -> IO a
-runLlvm dflags ver out us m = do
+runLlvm :: DynFlags -> LlvmVersion -> BufHandle -> LlvmM a -> IO a
+runLlvm dflags ver out m = do
     (a, _) <- runLlvmM m env
     return a
   where env = LlvmEnv { envFunMap = emptyUFM
@@ -278,7 +274,7 @@ runLlvm dflags ver out us m = do
                       , envVersion = ver
                       , envDynFlags = dflags
                       , envOutput = out
-                      , envUniq = us
+                      , envMask = 'n'
                       , envFreshMeta = MetaId 0
                       , envUniqMeta = emptyUFM
                       }
diff --git a/compiler/main/CodeOutput.hs b/compiler/main/CodeOutput.hs
index 01d714d57a..6b70366f45 100644
--- a/compiler/main/CodeOutput.hs
+++ b/compiler/main/CodeOutput.hs
@@ -176,11 +176,9 @@ outputAsm dflags this_mod location filenm cmm_stream
 
 outputLlvm :: DynFlags -> FilePath -> Stream IO RawCmmGroup a -> IO a
 outputLlvm dflags filenm cmm_stream
-  = do ncg_uniqs <- mkSplitUniqSupply 'n'
-
-       {-# SCC "llvm_output" #-} doOutput filenm $
+  = do {-# SCC "llvm_output" #-} doOutput filenm $
            \f -> {-# SCC "llvm_CodeGen" #-}
-                 llvmCodeGen dflags f ncg_uniqs cmm_stream
+                 llvmCodeGen dflags f cmm_stream
 
 {-
 ************************************************************************
@@ -262,4 +260,3 @@ outputForeignStubs_help _fname ""      _header _footer = return False
 outputForeignStubs_help fname doc_str header footer
    = do writeFile fname (header ++ doc_str ++ '\n':footer ++ "\n")
         return True
-
diff --git a/compiler/simplCore/CoreMonad.hs b/compiler/simplCore/CoreMonad.hs
index d2918a263f..fde925063b 100644
--- a/compiler/simplCore/CoreMonad.hs
+++ b/compiler/simplCore/CoreMonad.hs
@@ -28,7 +28,7 @@ module CoreMonad (
     -- ** Reading from the monad
     getHscEnv, getRuleBase, getModule,
     getDynFlags, getOrigNameCache, getPackageFamInstEnv,
-    getVisibleOrphanMods,
+    getVisibleOrphanMods, getUniqMask,
     getPrintUnqualified, getSrcSpanM,
 
     -- ** Writing to the monad
@@ -546,10 +546,6 @@ cmpEqTick _                             _                               = EQ
 ************************************************************************
 -}
 
-newtype CoreState = CoreState {
-        cs_uniq_supply :: UniqSupply
-}
-
 data CoreReader = CoreReader {
         cr_hsc_env             :: HscEnv,
         cr_rule_base           :: RuleBase,
@@ -557,7 +553,8 @@ data CoreReader = CoreReader {
         cr_print_unqual        :: PrintUnqualified,
         cr_loc                 :: SrcSpan,   -- Use this for log/error messages so they
                                              -- are at least tagged with the right source file
-        cr_visible_orphan_mods :: !ModuleSet
+        cr_visible_orphan_mods :: !ModuleSet,
+        cr_uniq_mask           :: !Char      -- Mask for creating unique values
 }
 
 -- Note: CoreWriter used to be defined with data, rather than newtype.  If it
@@ -579,55 +576,51 @@ plusWriter w1 w2 = CoreWriter {
 
 type CoreIOEnv = IOEnv CoreReader
 
--- | The monad used by Core-to-Core passes to access common state, register simplification
--- statistics and so on
-newtype CoreM a = CoreM { unCoreM :: CoreState -> CoreIOEnv (a, CoreState, CoreWriter) }
+-- | The monad used by Core-to-Core passes to register simplification statistics.
+--  Also used to have common state (in the form of UniqueSupply) for generating Uniques.
+newtype CoreM a = CoreM { unCoreM :: CoreIOEnv (a, CoreWriter) }
     deriving (Functor)
 
 instance Monad CoreM where
-    mx >>= f = CoreM $ \s -> do
-            (x, s', w1) <- unCoreM mx s
-            (y, s'', w2) <- unCoreM (f x) s'
+    mx >>= f = CoreM $ do
+            (x, w1) <- unCoreM mx
+            (y, w2) <- unCoreM (f x)
             let w = w1 `plusWriter` w2
-            return $ seq w (y, s'', w)
+            return $ seq w (y, w)
             -- forcing w before building the tuple avoids a space leak
             -- (#7702)
 
 instance Applicative CoreM where
-    pure x = CoreM $ \s -> nop s x
+    pure x = CoreM $ nop x
     (<*>) = ap
     m *> k = m >>= \_ -> k
 
 instance Alternative CoreM where
-    empty   = CoreM (const Control.Applicative.empty)
-    m <|> n = CoreM (\rs -> unCoreM m rs <|> unCoreM n rs)
+    empty   = CoreM Control.Applicative.empty
+    m <|> n = CoreM (unCoreM m <|> unCoreM n)
 
 instance MonadPlus CoreM
 
 instance MonadUnique CoreM where
     getUniqueSupplyM = do
-        us <- getS cs_uniq_supply
-        let (us1, us2) = splitUniqSupply us
-        modifyS (\s -> s { cs_uniq_supply = us2 })
-        return us1
+        mask <- read cr_uniq_mask
+        liftIO $! mkSplitUniqSupply mask
 
     getUniqueM = do
-        us <- getS cs_uniq_supply
-        let (u,us') = takeUniqFromSupply us
-        modifyS (\s -> s { cs_uniq_supply = us' })
-        return u
+        mask <- read cr_uniq_mask
+        liftIO $! uniqFromMask mask
 
 runCoreM :: HscEnv
          -> RuleBase
-         -> UniqSupply
+         -> Char -- ^ Mask
          -> Module
          -> ModuleSet
          -> PrintUnqualified
          -> SrcSpan
          -> CoreM a
          -> IO (a, SimplCount)
-runCoreM hsc_env rule_base us mod orph_imps print_unqual loc m
-  = liftM extract $ runIOEnv reader $ unCoreM m state
+runCoreM hsc_env rule_base mask mod orph_imps print_unqual loc m
+  = liftM extract $ runIOEnv reader $ unCoreM m
   where
     reader = CoreReader {
             cr_hsc_env = hsc_env,
@@ -635,14 +628,12 @@ runCoreM hsc_env rule_base us mod orph_imps print_unqual loc m
             cr_module = mod,
             cr_visible_orphan_mods = orph_imps,
             cr_print_unqual = print_unqual,
-            cr_loc = loc
-        }
-    state = CoreState {
-            cs_uniq_supply = us
+            cr_loc = loc,
+            cr_uniq_mask = mask
         }
 
-    extract :: (a, CoreState, CoreWriter) -> (a, SimplCount)
-    extract (value, _, writer) = (value, cw_simpl_count writer)
+    extract :: (a, CoreWriter) -> (a, SimplCount)
+    extract (value, writer) = (value, cw_simpl_count writer)
 
 {-
 ************************************************************************
@@ -652,28 +643,22 @@ runCoreM hsc_env rule_base us mod orph_imps print_unqual loc m
 ************************************************************************
 -}
 
-nop :: CoreState -> a -> CoreIOEnv (a, CoreState, CoreWriter)
-nop s x = do
+nop :: a -> CoreIOEnv (a, CoreWriter)
+nop x = do
     r <- getEnv
-    return (x, s, emptyWriter $ (hsc_dflags . cr_hsc_env) r)
+    return (x, emptyWriter $ (hsc_dflags . cr_hsc_env) r)
 
 read :: (CoreReader -> a) -> CoreM a
-read f = CoreM (\s -> getEnv >>= (\r -> nop s (f r)))
-
-getS :: (CoreState -> a) -> CoreM a
-getS f = CoreM (\s -> nop s (f s))
-
-modifyS :: (CoreState -> CoreState) -> CoreM ()
-modifyS f = CoreM (\s -> nop (f s) ())
+read f = CoreM $ getEnv >>= (\r -> nop (f r))
 
 write :: CoreWriter -> CoreM ()
-write w = CoreM (\s -> return ((), s, w))
+write w = CoreM $ return ((), w)
 
 -- \subsection{Lifting IO into the monad}
 
 -- | Lift an 'IOEnv' operation into 'CoreM'
 liftIOEnv :: CoreIOEnv a -> CoreM a
-liftIOEnv mx = CoreM (\s -> mx >>= (\x -> nop s x))
+liftIOEnv mx = CoreM (mx >>= (\x -> nop x))
 
 instance MonadIO CoreM where
     liftIO = liftIOEnv . IOEnv.liftIO
@@ -708,6 +693,9 @@ getSrcSpanM = read cr_loc
 addSimplCount :: SimplCount -> CoreM ()
 addSimplCount count = write (CoreWriter { cw_simpl_count = count })
 
+getUniqMask :: CoreM Char
+getUniqMask = read cr_uniq_mask
+
 -- Convenience accessors for useful fields of HscEnv
 
 instance HasDynFlags CoreM where
diff --git a/compiler/simplCore/CoreMonad.hs-boot b/compiler/simplCore/CoreMonad.hs-boot
index 206675e5e2..74c21e8216 100644
--- a/compiler/simplCore/CoreMonad.hs-boot
+++ b/compiler/simplCore/CoreMonad.hs-boot
@@ -12,11 +12,6 @@ module CoreMonad ( CoreToDo, CoreM ) where
 import GhcPrelude
 
 import IOEnv ( IOEnv )
-import UniqSupply ( UniqSupply )
-
-newtype CoreState = CoreState {
-        cs_uniq_supply :: UniqSupply
-}
 
 type CoreIOEnv = IOEnv CoreReader
 
@@ -28,9 +23,7 @@ newtype CoreWriter = CoreWriter {
 
 data SimplCount
 
-newtype CoreM a
-          = CoreM { unCoreM :: CoreState
-                                 -> CoreIOEnv (a, CoreState, CoreWriter) }
+newtype CoreM a = CoreM { unCoreM :: CoreIOEnv (a, CoreWriter) }
 
 instance Monad CoreM
 
diff --git a/compiler/simplCore/SimplCore.hs b/compiler/simplCore/SimplCore.hs
index cbfa757552..149a079a0a 100644
--- a/compiler/simplCore/SimplCore.hs
+++ b/compiler/simplCore/SimplCore.hs
@@ -72,13 +72,13 @@ core2core hsc_env guts@(ModGuts { mg_module  = mod
                                 , mg_loc     = loc
                                 , mg_deps    = deps
                                 , mg_rdr_env = rdr_env })
-  = do { us <- mkSplitUniqSupply 's'
-       -- make sure all plugins are loaded
+  = do { -- make sure all plugins are loaded
 
        ; let builtin_passes = getCoreToDo dflags
              orph_mods = mkModuleSet (mod : dep_orphs deps)
+             uniq_mask = 's'
        ;
-       ; (guts2, stats) <- runCoreM hsc_env hpt_rule_base us mod
+       ; (guts2, stats) <- runCoreM hsc_env hpt_rule_base uniq_mask mod
                                     orph_mods print_unqual loc $
                            do { hsc_env' <- getHscEnv
                               ; dflags' <- liftIO $ initializePlugins hsc_env'
diff --git a/compiler/simplStg/SimplStg.hs b/compiler/simplStg/SimplStg.hs
index c2f145df11..89b7d4205e 100644
--- a/compiler/simplStg/SimplStg.hs
+++ b/compiler/simplStg/SimplStg.hs
@@ -32,15 +32,17 @@ import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.Trans.State.Strict
 
-newtype StgM a = StgM { _unStgM :: StateT UniqSupply IO a }
+newtype StgM a = StgM { _unStgM :: StateT Char IO a }
   deriving (Functor, Applicative, Monad, MonadIO)
 
 instance MonadUnique StgM where
-  getUniqueSupplyM = StgM (state splitUniqSupply)
-  getUniqueM = StgM (state takeUniqFromSupply)
+  getUniqueSupplyM = StgM $ do { mask <- get
+                               ; liftIO $! mkSplitUniqSupply mask}
+  getUniqueM = StgM $ do { mask <- get
+                         ; liftIO $! uniqFromMask mask}
 
-runStgM :: UniqSupply -> StgM a -> IO a
-runStgM us (StgM m) = evalStateT m us
+runStgM :: Char -> StgM a -> IO a
+runStgM mask (StgM m) = evalStateT m mask
 
 stg2stg :: DynFlags                  -- includes spec of what stg-to-stg passes to do
         -> Module                    -- module being compiled
@@ -50,10 +52,8 @@ stg2stg :: DynFlags                  -- includes spec of what stg-to-stg passes
 stg2stg dflags this_mod binds
   = do  { dump_when Opt_D_dump_stg "STG:" binds
         ; showPass dflags "Stg2Stg"
-        ; us <- mkSplitUniqSupply 'g'
-
         -- Do the main business!
-        ; binds' <- runStgM us $
+        ; binds' <- runStgM 'g' $
             foldM do_stg_pass binds (getStgToDo dflags)
 
         ; dump_when Opt_D_dump_stg_final "Final STG:" binds'
diff --git a/compiler/typecheck/TcRnMonad.hs b/compiler/typecheck/TcRnMonad.hs
index c820eb3c20..3442e8729a 100644
--- a/compiler/typecheck/TcRnMonad.hs
+++ b/compiler/typecheck/TcRnMonad.hs
@@ -397,17 +397,14 @@ an actual crash (attempting to look up the Integer type).
 ************************************************************************
 -}
 
-initTcRnIf :: Char              -- Tag for unique supply
+initTcRnIf :: Char              -- ^ Mask for unique supply
            -> HscEnv
            -> gbl -> lcl
            -> TcRnIf gbl lcl a
            -> IO a
-initTcRnIf uniq_tag hsc_env gbl_env lcl_env thing_inside
-   = do { us     <- mkSplitUniqSupply uniq_tag ;
-        ; us_var <- newIORef us ;
-
-        ; let { env = Env { env_top = hsc_env,
-                            env_us  = us_var,
+initTcRnIf uniq_mask hsc_env gbl_env lcl_env thing_inside
+   = do { let { env = Env { env_top = hsc_env,
+                            env_um  = uniq_mask,
                             env_gbl = gbl_env,
                             env_lcl = lcl_env} }
 
@@ -595,27 +592,15 @@ escapeArrowScope
 
 newUnique :: TcRnIf gbl lcl Unique
 newUnique
- = do { env <- getEnv ;
-        let { u_var = env_us env } ;
-        us <- readMutVar u_var ;
-        case takeUniqFromSupply us of { (uniq, us') -> do {
-        writeMutVar u_var us' ;
-        return $! uniq }}}
-   -- NOTE 1: we strictly split the supply, to avoid the possibility of leaving
-   -- a chain of unevaluated supplies behind.
-   -- NOTE 2: we use the uniq in the supply from the MutVar directly, and
-   -- throw away one half of the new split supply.  This is safe because this
-   -- is the only place we use that unique.  Using the other half of the split
-   -- supply is safer, but slower.
+ = do { env <- getEnv
+      ; let mask = env_um env
+      ; liftIO $! uniqFromMask mask }
 
 newUniqueSupply :: TcRnIf gbl lcl UniqSupply
 newUniqueSupply
- = do { env <- getEnv ;
-        let { u_var = env_us env } ;
-        us <- readMutVar u_var ;
-        case splitUniqSupply us of { (us1,us2) -> do {
-        writeMutVar u_var us1 ;
-        return us2 }}}
+ = do { env <- getEnv
+      ; let mask = env_um env
+      ; liftIO $! mkSplitUniqSupply mask }
 
 cloneLocalName :: Name -> TcM Name
 -- Make a fresh Internal name with the same OccName and SrcSpan
@@ -1944,12 +1929,8 @@ forkM_maybe :: SDoc -> IfL a -> IfL (Maybe a)
 -- signatures, which is pretty benign
 
 forkM_maybe doc thing_inside
- -- NB: Don't share the mutable env_us with the interleaved thread since env_us
- --     does not get updated atomically (e.g. in newUnique and newUniqueSupply).
- = do { child_us <- newUniqueSupply
-      ; child_env_us <- newMutVar child_us
-        -- see Note [Masking exceptions in forkM_maybe]
-      ; unsafeInterleaveM $ uninterruptibleMaskM_ $ updEnv (\env -> env { env_us = child_env_us }) $
+ = do { -- see Note [Masking exceptions in forkM_maybe]
+      ; unsafeInterleaveM $ uninterruptibleMaskM_ $
         do { traceIf (text "Starting fork {" <+> doc)
            ; mb_res <- tryM $
                        updLclEnv (\env -> env { if_loc = if_loc env $$ doc }) $
diff --git a/compiler/typecheck/TcRnTypes.hs b/compiler/typecheck/TcRnTypes.hs
index 8fa12b28b1..3445d5b793 100644
--- a/compiler/typecheck/TcRnTypes.hs
+++ b/compiler/typecheck/TcRnTypes.hs
@@ -113,7 +113,6 @@ import SrcLoc
 import VarSet
 import ErrUtils
 import UniqFM
-import UniqSupply
 import BasicTypes
 import Bag
 import DynFlags
@@ -209,8 +208,7 @@ data Env gbl lcl
                              -- Includes all info about imported things
                              -- BangPattern is to fix leak, see #15111
 
-        env_us   :: {-# UNPACK #-} !(IORef UniqSupply),
-                             -- Unique supply for local variables
+        env_um   :: !Char,   -- Mask for Uniques
 
         env_gbl  :: gbl,     -- Info about things defined at the top level
                              -- of the module being compiled