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>

6 files changed, 26 insertions, 9 deletions

diff --git a/compiler/basicTypes/UniqSupply.lhs b/compiler/basicTypes/UniqSupply.lhs
index fb07e73824..0c6007a4f7 100644
--- a/compiler/basicTypes/UniqSupply.lhs
+++ b/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
diff --git a/compiler/codeGen/StgCmmExpr.hs b/compiler/codeGen/StgCmmExpr.hs
index d7edf8e193..3d60def450 100644
--- a/compiler/codeGen/StgCmmExpr.hs
+++ b/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
diff --git a/compiler/codeGen/StgCmmMonad.hs b/compiler/codeGen/StgCmmMonad.hs
index 3f361e3f51..251b679078 100644
--- a/compiler/codeGen/StgCmmMonad.hs
+++ b/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
diff --git a/compiler/simplCore/CoreMonad.lhs b/compiler/simplCore/CoreMonad.lhs
index 2aa42cc9ad..04cdc36b28 100644
--- a/compiler/simplCore/CoreMonad.lhs
+++ b/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
diff --git a/compiler/simplCore/SimplMonad.lhs b/compiler/simplCore/SimplMonad.lhs
index a5eb116d82..4c3c72d301 100644
--- a/compiler/simplCore/SimplMonad.lhs
+++ b/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
diff --git a/compiler/specialise/Specialise.lhs b/compiler/specialise/Specialise.lhs
index a161444d7b..bf73bec240 100644
--- a/compiler/specialise/Specialise.lhs
+++ b/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