Avoid needlessly splitting a UniqSupply when extracting a Unique (#8041)

In many places, 'splitUniqSupply' + 'uniqFromSupply' is used to split a
UniqSupply into a Unique and a new UniqSupply. In such places we should
instead use the more efficient and more appropriate
'takeUniqFromSupply' (or equivalent).

Not only is the former method slower, it also generates and throws away
an extra Unique.
Signed-off-by: Austin Seipp <aseipp@pobox.com>

compiler/basicTypes/UniqSupply.lhs

@@ -176,6 +176,10 @@ class Monad m => MonadUnique m where
     -- | Get an infinite list of new unique identifiers
     getUniquesM :: m [Unique]
 
+    -- This default definition of getUniqueM, while correct, is not as
+    -- efficient as it could be since it needlessly generates and throws away
+    -- an extra Unique. For your instances consider providing an explicit
+    -- definition for 'getUniqueM' which uses 'takeUniqFromSupply' directly.
     getUniqueM  = liftM uniqFromSupply  getUniqueSupplyM
     getUniquesM = liftM uniqsFromSupply getUniqueSupplyM
 
@@ -185,8 +189,8 @@ instance MonadUnique UniqSM where
     getUniquesM = getUniquesUs
 
 getUniqueUs :: UniqSM Unique
-getUniqueUs = USM (\us -> case splitUniqSupply us of
-                          (us1,us2) -> (# uniqFromSupply us1, us2 #))
+getUniqueUs = USM (\us -> case takeUniqFromSupply us of
+                          (u,us') -> (# u, us' #))
 
 getUniquesUs :: UniqSM [Unique]
 getUniquesUs = USM (\us -> case splitUniqSupply us of

compiler/codeGen/StgCmmExpr.hs

@@ -43,7 +43,6 @@ import Maybes
 import Util
 import FastString
 import Outputable
-import UniqSupply
 
 import Control.Monad (when,void)
 
@@ -70,8 +69,8 @@ cgExpr (StgLit lit)       = do cmm_lit <- cgLit lit
 
 cgExpr (StgLet binds expr)             = do { cgBind binds;     cgExpr expr }
 cgExpr (StgLetNoEscape _ _ binds expr) =
-  do { us <- newUniqSupply
-     ; let join_id = mkBlockId (uniqFromSupply us)
+  do { u <- newUnique
+     ; let join_id = mkBlockId u
      ; cgLneBinds join_id binds
      ; r <- cgExpr expr
      ; emitLabel join_id

compiler/codeGen/StgCmmMonad.hs

@@ -446,8 +446,10 @@ newUniqSupply = do
 
 newUnique :: FCode Unique
 newUnique = do
-        us <- newUniqSupply
-        return (uniqFromSupply us)
+        state <- getState
+        let (u,us') = takeUniqFromSupply (cgs_uniqs state)
+        setState $ state { cgs_uniqs = us' }
+        return u
 
 ------------------
 getInfoDown :: FCode CgInfoDownwards

compiler/simplCore/CoreMonad.lhs

@@ -783,6 +783,12 @@ instance MonadUnique CoreM where
         modifyS (\s -> s { cs_uniq_supply = us2 })
         return us1
 
+    getUniqueM = do
+        us <- getS cs_uniq_supply
+        let (u,us') = takeUniqFromSupply us
+        modifyS (\s -> s { cs_uniq_supply = us' })
+        return u
+
 runCoreM :: HscEnv
          -> RuleBase
          -> UniqSupply

compiler/simplCore/SimplMonad.lhs

@@ -145,8 +145,8 @@ instance MonadUnique SimplM where
                                 (us1, us2) -> return (us1, us2, sc))
 
     getUniqueM
-       = SM (\_st_env us sc -> case splitUniqSupply us of
-                                (us1, us2) -> return (uniqFromSupply us1, us2, sc))
+       = SM (\_st_env us sc -> case takeUniqFromSupply us of
+                                (u, us') -> return (u, us', sc))
 
     getUniquesM
         = SM (\_st_env us sc -> case splitUniqSupply us of

compiler/specialise/Specialise.lhs

@@ -1882,6 +1882,12 @@ instance MonadUnique SpecM where
                      put $ st { spec_uniq_supply = us2 }
                      return us1
 
+    getUniqueM
+        = SpecM $ do st <- get
+                     let (u,us') = takeUniqFromSupply $ spec_uniq_supply st
+                     put $ st { spec_uniq_supply = us' }
+                     return u
+
 instance HasDynFlags SpecM where
     getDynFlags = SpecM $ liftM spec_dflags get