SimplM: Create uniques via IO instead of threading

5 files changed, 147 insertions, 70 deletions

diff --git a/compiler/GHC/Builtin/Uniques.hs b/compiler/GHC/Builtin/Uniques.hs
index 038647acc9..4b3f3e00b6 100644
--- a/compiler/GHC/Builtin/Uniques.hs
+++ b/compiler/GHC/Builtin/Uniques.hs
@@ -290,16 +290,13 @@ getTupleDataConName boxity n =
       _          -> panic "getTupleDataConName: impossible"
 
 {-
-************************************************************************
-*                                                                      *
-\subsection[Uniques-prelude]{@Uniques@ for wired-in Prelude things}
-*                                                                      *
-************************************************************************
-
+Note [Uniques for wired-in prelude things and known masks]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Allocation of unique supply characters:
         v,t,u : for renumbering value-, type- and usage- vars.
         B:   builtin
         C-E: pseudo uniques     (used in native-code generator)
+        I:   GHCi evaluation
         X:   uniques from mkLocalUnique
         _:   unifiable tyvars   (above)
         0-9: prelude things below
@@ -308,15 +305,20 @@ Allocation of unique supply characters:
 
         other a-z: lower case chars for unique supplies.  Used so far:
 
+        a       TypeChecking?
+        c       StgToCmm/Renamer
         d       desugarer
         f       AbsC flattener
         g       SimplStg
+        i       TypeChecking interface files
         j       constraint tuple superclass selectors
         k       constraint tuple tycons
         m       constraint tuple datacons
-        n       Native codegen
+        n       Native/LLVM codegen
         r       Hsc name cache
         s       simplifier
+        u       Cmm pipeline
+        y       GHCi bytecode generator
         z       anonymous sums
 -}
 
diff --git a/compiler/GHC/Core/Opt/Pipeline.hs b/compiler/GHC/Core/Opt/Pipeline.hs
index 76d2f3d459..6a21063f22 100644
--- a/compiler/GHC/Core/Opt/Pipeline.hs
+++ b/compiler/GHC/Core/Opt/Pipeline.hs
@@ -69,7 +69,7 @@ import GHC.Types.Basic
 import GHC.Types.Demand ( zapDmdEnvSig )
 import GHC.Types.Var.Set
 import GHC.Types.Var.Env
-import GHC.Types.Unique.Supply ( UniqSupply, mkSplitUniqSupply, splitUniqSupply )
+import GHC.Types.Unique.Supply ( UniqSupply )
 import GHC.Types.Unique.FM
 import GHC.Types.Name.Ppr
 
@@ -634,10 +634,9 @@ simplifyExpr hsc_env expr
                             snd $ ic_instances $ hsc_IC hsc_env )
               simpl_env = simplEnvForGHCi dflags
 
-        ; us <-  mkSplitUniqSupply 's'
         ; let sz = exprSize expr
 
-        ; (expr', counts) <- initSmpl dflags rule_env fi_env us sz $
+        ; (expr', counts) <- initSmpl dflags rule_env fi_env sz $
                              simplExprGently simpl_env expr
 
         ; Err.dumpIfSet dflags (dopt Opt_D_dump_simpl_stats dflags)
@@ -685,27 +684,25 @@ simplExprGently env expr = do
 simplifyPgm :: CoreToDo -> ModGuts -> CoreM ModGuts
 simplifyPgm pass guts
   = do { hsc_env <- getHscEnv
-       ; us <- getUniqueSupplyM
        ; rb <- getRuleBase
        ; liftIOWithCount $
-         simplifyPgmIO pass hsc_env us rb guts }
+         simplifyPgmIO pass hsc_env rb guts }
 
 simplifyPgmIO :: CoreToDo
               -> HscEnv
-              -> UniqSupply
               -> RuleBase
               -> ModGuts
               -> IO (SimplCount, ModGuts)  -- New bindings
 
 simplifyPgmIO pass@(CoreDoSimplify max_iterations mode)
-              hsc_env us hpt_rule_base
+              hsc_env hpt_rule_base
               guts@(ModGuts { mg_module = this_mod
                             , mg_rdr_env = rdr_env
                             , mg_deps = deps
                             , mg_binds = binds, mg_rules = rules
                             , mg_fam_inst_env = fam_inst_env })
   = do { (termination_msg, it_count, counts_out, guts')
-           <- do_iteration us 1 [] binds rules
+           <- do_iteration 1 [] binds rules
 
         ; Err.dumpIfSet dflags (dopt Opt_D_verbose_core2core dflags &&
                                 dopt Opt_D_dump_simpl_stats  dflags)
@@ -724,14 +721,14 @@ simplifyPgmIO pass@(CoreDoSimplify max_iterations mode)
     active_rule  = activeRule mode
     active_unf   = activeUnfolding mode
 
-    do_iteration :: UniqSupply
-                 -> Int          -- Counts iterations
+    do_iteration :: Int --UniqSupply
+                --  -> Int          -- Counts iterations
                  -> [SimplCount] -- Counts from earlier iterations, reversed
                  -> CoreProgram  -- Bindings in
                  -> [CoreRule]   -- and orphan rules
                  -> IO (String, Int, SimplCount, ModGuts)
 
-    do_iteration us iteration_no counts_so_far binds rules
+    do_iteration iteration_no counts_so_far binds rules
         -- iteration_no is the number of the iteration we are
         -- about to begin, with '1' for the first
       | iteration_no > max_iterations   -- Stop if we've run out of iterations
@@ -776,7 +773,7 @@ simplifyPgmIO pass@(CoreDoSimplify max_iterations mode)
 
                 -- Simplify the program
            ((binds1, rules1), counts1) <-
-             initSmpl dflags (mkRuleEnv rule_base2 vis_orphs) fam_envs us1 sz $
+             initSmpl dflags (mkRuleEnv rule_base2 vis_orphs) fam_envs sz $
                do { (floats, env1) <- {-# SCC "SimplTopBinds" #-}
                                       simplTopBinds simpl_env tagged_binds
 
@@ -810,20 +807,18 @@ simplifyPgmIO pass@(CoreDoSimplify max_iterations mode)
            lintPassResult hsc_env pass binds2 ;
 
                 -- Loop
-           do_iteration us2 (iteration_no + 1) (counts1:counts_so_far) binds2 rules1
+           do_iteration (iteration_no + 1) (counts1:counts_so_far) binds2 rules1
            } }
 #if __GLASGOW_HASKELL__ <= 810
       | otherwise = panic "do_iteration"
 #endif
       where
-        (us1, us2) = splitUniqSupply us
-
         -- Remember the counts_so_far are reversed
         totalise :: [SimplCount] -> SimplCount
         totalise = foldr (\c acc -> acc `plusSimplCount` c)
                          (zeroSimplCount dflags)
 
-simplifyPgmIO _ _ _ _ _ = panic "simplifyPgmIO"
+simplifyPgmIO _ _ _ _ = panic "simplifyPgmIO"
 
 -------------------
 dump_end_iteration :: DynFlags -> PrintUnqualified -> Int
diff --git a/compiler/GHC/Core/Opt/Simplify/Monad.hs b/compiler/GHC/Core/Opt/Simplify/Monad.hs
index 620db9da22..4af454e381 100644
--- a/compiler/GHC/Core/Opt/Simplify/Monad.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Monad.hs
@@ -60,15 +60,14 @@ For the simplifier monad, we want to {\em thread} a unique supply and a counter.
 
 newtype SimplM result
   =  SM'  { unSM :: SimplTopEnv  -- Envt that does not change much
-                 -> UniqSupply   -- We thread the unique supply because
-                                 -- constantly splitting it is rather expensive
                  -> SimplCount
-                 -> IO (result, UniqSupply, SimplCount)}
-    -- We only need IO here for dump output
+                 -> IO (result, SimplCount)}
+    -- We only need IO here for dump output, but since we already have it
+    -- we might as well use it for uniques.
     deriving (Functor)
 
-pattern SM :: (SimplTopEnv -> UniqSupply -> SimplCount
-               -> IO (result, UniqSupply, SimplCount))
+pattern SM :: (SimplTopEnv -> SimplCount
+               -> IO (result, SimplCount))
           -> SimplM result
 -- This pattern synonym makes the simplifier monad eta-expand,
 -- which as a very beneficial effect on compiler performance
@@ -89,14 +88,15 @@ data SimplTopEnv
         }
 
 initSmpl :: DynFlags -> RuleEnv -> (FamInstEnv, FamInstEnv)
-         -> UniqSupply          -- No init count; set to 0
          -> Int                 -- Size of the bindings, used to limit
                                 -- the number of ticks we allow
          -> SimplM a
          -> IO (a, SimplCount)
 
-initSmpl dflags rules fam_envs us size m
-  = do (result, _, count) <- unSM m env us (zeroSimplCount dflags)
+initSmpl dflags rules fam_envs size m
+  = do -- No init count; set to 0
+       let simplCount = zeroSimplCount dflags
+       (result, count) <- unSM m env simplCount
        return (result, count)
   where
     env = STE { st_flags = dflags
@@ -141,20 +141,20 @@ instance Monad SimplM where
    (>>=)  = thenSmpl
 
 returnSmpl :: a -> SimplM a
-returnSmpl e = SM (\_st_env us sc -> return (e, us, sc))
+returnSmpl e = SM (\_st_env sc -> return (e, sc))
 
 thenSmpl  :: SimplM a -> (a -> SimplM b) -> SimplM b
 thenSmpl_ :: SimplM a -> SimplM b -> SimplM b
 
 thenSmpl m k
-  = SM $ \st_env us0 sc0 -> do
-      (m_result, us1, sc1) <- unSM m st_env us0 sc0
-      unSM (k m_result) st_env us1 sc1
+  = SM $ \st_env sc0 -> do
+      (m_result, sc1) <- unSM m st_env sc0
+      unSM (k m_result) st_env sc1
 
 thenSmpl_ m k
-  = SM $ \st_env us0 sc0 -> do
-      (_, us1, sc1) <- unSM m st_env us0 sc0
-      unSM k st_env us1 sc1
+  = SM $ \st_env sc0 -> do
+      (_, sc1) <- unSM m st_env sc0
+      unSM k st_env sc1
 
 -- TODO: this specializing is not allowed
 -- {-# SPECIALIZE mapM         :: (a -> SimplM b) -> [a] -> SimplM [b] #-}
@@ -177,35 +177,30 @@ traceSmpl herald doc
 ************************************************************************
 -}
 
-instance MonadUnique SimplM where
-    getUniqueSupplyM
-       = SM (\_st_env us sc -> case splitUniqSupply us of
-                                (us1, us2) -> return (us1, us2, sc))
-
-    getUniqueM
-       = SM (\_st_env us sc -> case takeUniqFromSupply us of
-                                (u, us') -> return (u, us', sc))
+-- See Note [Uniques for wired-in prelude things and known masks] in GHC.Builtin.Uniques
+simplMask :: Char
+simplMask = 's'
 
-    getUniquesM
-        = SM (\_st_env us sc -> case splitUniqSupply us of
-                                (us1, us2) -> return (uniqsFromSupply us1, us2, sc))
+instance MonadUnique SimplM where
+    getUniqueSupplyM = liftIO $ mkSplitUniqSupply simplMask
+    getUniqueM = liftIO $ uniqFromMask simplMask
 
 instance HasDynFlags SimplM where
-    getDynFlags = SM (\st_env us sc -> return (st_flags st_env, us, sc))
+    getDynFlags = SM (\st_env sc -> return (st_flags st_env, sc))
 
 instance MonadIO SimplM where
-    liftIO m = SM $ \_ us sc -> do
+    liftIO m = SM $ \_ sc -> do
       x <- m
-      return (x, us, sc)
+      return (x, sc)
 
 getSimplRules :: SimplM RuleEnv
-getSimplRules = SM (\st_env us sc -> return (st_rules st_env, us, sc))
+getSimplRules = SM (\st_env sc -> return (st_rules st_env, sc))
 
 getFamEnvs :: SimplM (FamInstEnv, FamInstEnv)
-getFamEnvs = SM (\st_env us sc -> return (st_fams st_env, us, sc))
+getFamEnvs = SM (\st_env sc -> return (st_fams st_env, sc))
 
 getOptCoercionOpts :: SimplM OptCoercionOpts
-getOptCoercionOpts = SM (\st_env us sc -> return (st_co_opt_opts st_env, us, sc))
+getOptCoercionOpts = SM (\st_env sc -> return (st_co_opt_opts st_env, sc))
 
 newId :: FastString -> Mult -> Type -> SimplM Id
 newId fs w ty = do uniq <- getUniqueM
@@ -234,21 +229,21 @@ newJoinId bndrs body_ty
 -}
 
 getSimplCount :: SimplM SimplCount
-getSimplCount = SM (\_st_env us sc -> return (sc, us, sc))
+getSimplCount = SM (\_st_env sc -> return (sc, sc))
 
 tick :: Tick -> SimplM ()
-tick t = SM (\st_env us sc -> let sc' = doSimplTick (st_flags st_env) t sc
-                              in sc' `seq` return ((), us, sc'))
+tick t = SM (\st_env sc -> let sc' = doSimplTick (st_flags st_env) t sc
+                              in sc' `seq` return ((), sc'))
 
 checkedTick :: Tick -> SimplM ()
 -- Try to take a tick, but fail if too many
 checkedTick t
-  = SM (\st_env us sc ->
+  = SM (\st_env sc ->
            if st_max_ticks st_env <= mkIntWithInf (simplCountN sc)
            then throwGhcExceptionIO $
                   PprProgramError "Simplifier ticks exhausted" (msg sc)
            else let sc' = doSimplTick (st_flags st_env) t sc
-                in sc' `seq` return ((), us, sc'))
+                in sc' `seq` return ((), sc'))
   where
     msg sc = vcat
       [ text "When trying" <+> ppr t
@@ -276,5 +271,5 @@ freeTick :: Tick -> SimplM ()
 -- Record a tick, but don't add to the total tick count, which is
 -- used to decide when nothing further has happened
 freeTick t
-   = SM (\_st_env us sc -> let sc' = doFreeSimplTick t sc
-                           in sc' `seq` return ((), us, sc'))
+   = SM (\_st_env sc -> let sc' = doFreeSimplTick t sc
+                           in sc' `seq` return ((), sc'))
diff --git a/compiler/GHC/Types/Unique.hs b/compiler/GHC/Types/Unique.hs
index cc013dc9a1..4d19ba7f66 100644
--- a/compiler/GHC/Types/Unique.hs
+++ b/compiler/GHC/Types/Unique.hs
@@ -68,8 +68,20 @@ import Data.Bits
 *                                                                      *
 ************************************************************************
 
-The @Chars@ are ``tag letters'' that identify the @UniqueSupply@.
-Fast comparison is everything on @Uniques@:
+Note [Uniques and masks]
+~~~~~~~~~~~~~~~~~~~~~~~~
+A `Unique` in GHC is a Word-sized value composed of two pieces:
+* A "mask", of width `UNIQUE_TAG_BITS`, in the high order bits
+* A number, of width `uNIQUE_BITS`, which fills up the remainder of the Word
+
+The mask is typically an ASCII character.  It is typically used to make it easier
+to distinguish uniques constructed by different parts of the compiler.
+There is a (potentially incomplete) list of unique masks used given in
+GHC.Builtin.Uniques. See Note [Uniques-prelude - Uniques for wired-in Prelude things]
+
+`mkUnique` constructs a `Unique` from its pieces
+  mkUnique :: Char -> Int -> Unique
+
 -}
 
 -- | Unique identifier.
diff --git a/compiler/GHC/Types/Unique/Supply.hs b/compiler/GHC/Types/Unique/Supply.hs
index 4b146edd9f..0a10fde9b3 100644
--- a/compiler/GHC/Types/Unique/Supply.hs
+++ b/compiler/GHC/Types/Unique/Supply.hs
@@ -86,9 +86,35 @@ lazily-evaluated infinite tree.
      * The fresh node
      * A thunk for each sub-tree
 
-Note [Optimising the unique supply]
+Note [How unique supplies are used]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The general design (used throughout GHC) is to:
+
+* For creating new uniques either a UniqSupply is used and threaded through
+  or for monadic code a MonadUnique instance might conjure up uniques using
+  `uniqFromMask`.
+* Different parts of the compiler will use a UniqSupply or MonadUnique instance
+  with a specific mask. This way the different parts of the compiler will
+  generate uniques with different masks.
+
+If different code shares the same mask then care has to be taken that all uniques
+still get distinct numbers. Usually this is done by relying on genSym which
+has *one* counter per GHC invocation that is relied on by all calls to it.
+But using something like the address for pinned objects works as well and in fact is done
+for fast strings.
+
+This is important for example in the simplifier. Most passes of the simplifier use
+the same mask 's'. However in some places we create a unique supply using `mkSplitUniqSupply`
+and thread it through the code, while in GHC.Core.Opt.Simplify.Monad  we use the
+`instance MonadUnique SimplM`, which uses `mkSplitUniqSupply` in getUniqueSupplyM
+and `uniqFromMask` in getUniqeM.
+
+Ultimately all these boil down to each new unique consisting of the mask and the result from
+a call to `genSym`. The later producing a distinct number for each invocation ensuring
+uniques are distinct.
 
+Note [Optimising the unique supply]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 The inner loop of mkSplitUniqSupply is a function closure
 
      mk_supply s0 =
@@ -117,6 +143,46 @@ result. It was very brittle and required enabling -fno-state-hack globally. So
 it has been rewritten using lower level constructs to explicitly state what we
 want.
 
+Note [Optimising use of unique supplies]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+When it comes to having a way to generate new Uniques
+there are generally three ways to deal with this:
+
+For pure code the only good approach is to take an UniqSupply
+as argument. Then  thread it through the code splitting it
+for sub-passes or when creating uniques.
+The code for this is about as optimized as it gets, but we can't
+get around the need to allocate one `UniqSupply` for each Unique
+we need.
+
+For code in IO we can improve on this by threading only the *mask*
+we are going to use for Uniques. Using `uniqFromMask` to
+generate uniques as needed. This gets rid of the overhead of
+allocating a new UniqSupply for each unique generated. It also avoids
+frequent state updates when the Unique/Mask is part of the state in a
+state monad.
+
+For monadic code in IO which always uses the same mask we can go further
+and hardcode the mask into the MonadUnique instance. On top of all the
+benefits of threading the mask this *also* has the benefit of avoiding
+the mask getting captured in thunks, or being passed around at runtime.
+It does however come at the cost of having to use a fixed Mask for all
+code run in this Monad. But rememeber, the Mask is purely cosmetic:
+See Note [Uniques and masks].
+
+NB: It's *not* an optimization to pass around the UniqSupply inside an
+IORef instead of the mask. While this would avoid frequent state updates
+it still requires allocating one UniqSupply per Unique. On top of some
+overhead for reading/writing to/from the IORef.
+
+All of this hinges on the assumption that UniqSupply and
+uniqFromMask use the same source of distinct numbers (`genSym`) which
+allows both to be used at the same time, with the same mask, while still
+ensuring distinct uniques.
+One might consider this fact to be an "accident". But GHC worked like this
+as far back as source control history goes. It also allows the later two
+optimizations to be used. So it seems safe to depend on this fact.
+
 -}
 
 
@@ -132,9 +198,16 @@ data UniqSupply
                                 -- when split => these two supplies
 
 mkSplitUniqSupply :: Char -> IO UniqSupply
--- ^ Create a unique supply out of thin air. The character given must
--- be distinct from those of all calls to this function in the compiler
--- for the values generated to be truly unique.
+-- ^ Create a unique supply out of thin air.
+-- The "mask" (Char) supplied is purely cosmetic, making it easier
+-- to figure out where a Unique was born. See
+-- Note [Uniques and masks].
+--
+-- The payload part of the Uniques allocated from this UniqSupply are
+-- guaranteed distinct wrt all other supplies, regardless of their "mask".
+-- This is achieved by allocating the payload part from
+-- a single source of Uniques, namely `genSym`, shared across
+-- all UniqSupply's.
 
 -- See Note [How the unique supply works]
 -- See Note [Optimising the unique supply]
@@ -187,7 +260,7 @@ initUniqSupply counter inc = do
     poke ghc_unique_inc     inc
 
 uniqFromMask :: Char -> IO Unique
-uniqFromMask mask
+uniqFromMask !mask
   = do { uqNum <- genSym
        ; return $! mkUnique mask uqNum }