Improve efficiency of extending a RuleEnv with a new RuleBasewip/lazy-unfolding-typecheck

Essentially we apply the identity:

> lookupNameEnv n (plusNameEnv_C (++) rb1 rb2)
>   = lookupNameEnv n rb1 ++ lookupNameEnv n rb2

The latter being more efficient as we don't construct an intermediate
map.

This is now quite important as each time we try and apply rules we need
to combine the current EPS RuleBase with the HPT and ModGuts rule bases.

4 files changed, 35 insertions, 14 deletions

diff --git a/compiler/GHC/Core.hs b/compiler/GHC/Core.hs
index 4fac9ce8e1..2216c65591 100644
--- a/compiler/GHC/Core.hs
+++ b/compiler/GHC/Core.hs
@@ -101,7 +101,7 @@ import GHC.Core.Type
 import GHC.Core.Coercion
 import GHC.Types.Name
 import GHC.Types.Name.Set
-import GHC.Types.Name.Env( NameEnv, emptyNameEnv )
+import GHC.Types.Name.Env( NameEnv )
 import GHC.Types.Literal
 import GHC.Types.Tickish
 import GHC.Core.DataCon
@@ -1047,15 +1047,40 @@ type RuleBase = NameEnv [CoreRule]
 -- but it also includes the set of visible orphans we use to filter out orphan
 -- rules which are not visible (even though we can see them...)
 data RuleEnv
-    = RuleEnv { re_base          :: RuleBase
+    = RuleEnv { re_base          :: [RuleBase] -- See Note [Why re_base is a list]
               , re_visible_orphs :: ModuleSet
               }
 
 mkRuleEnv :: RuleBase -> [Module] -> RuleEnv
-mkRuleEnv rules vis_orphs = RuleEnv rules (mkModuleSet vis_orphs)
+mkRuleEnv rules vis_orphs = RuleEnv [rules] (mkModuleSet vis_orphs)
 
 emptyRuleEnv :: RuleEnv
-emptyRuleEnv = RuleEnv emptyNameEnv emptyModuleSet
+emptyRuleEnv = RuleEnv [] emptyModuleSet
+
+{-
+Note [Why re_base is a list]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+In Note [Overall plumbing for rules], it is explained that the final
+RuleBase which we must consider is combined from 4 different sources.
+
+During simplifier runs, the fourth source of rules is constantly being updated
+as new interfaces are loaded into the EPS. Therefore just before we check to see
+if any rules match we get the EPS RuleBase and combine it with the existing RuleBase
+and then perform exactly 1 lookup into the new map.
+
+It is more efficient to avoid combining the environments and store the uncombined
+environments as we can instead perform 1 lookup into each environment and then combine
+the results.
+
+Essentially we use the identity:
+
+> lookupNameEnv n (plusNameEnv_C (++) rb1 rb2)
+>   = lookupNameEnv n rb1 ++ lookupNameEnv n rb2
+
+The latter being more efficient as we don't construct an intermediate
+map.
+-}
 
 -- | A 'CoreRule' is:
 --
diff --git a/compiler/GHC/Core/Opt/Pipeline.hs b/compiler/GHC/Core/Opt/Pipeline.hs
index a6170e185b..45f5b3a550 100644
--- a/compiler/GHC/Core/Opt/Pipeline.hs
+++ b/compiler/GHC/Core/Opt/Pipeline.hs
@@ -567,7 +567,7 @@ ruleCheckPass current_phase pat guts = do
                 (const ()) $ do
         rb <- getRuleBase
         vis_orphs <- getVisibleOrphanMods
-        let rule_fn fn = getRules (RuleEnv rb vis_orphs) fn
+        let rule_fn fn = getRules (RuleEnv [rb] vis_orphs) fn
                           ++ (mg_rules guts)
         let ropts = initRuleOpts dflags
         liftIO $ logDumpMsg logger "Rule check"
diff --git a/compiler/GHC/Core/Opt/Specialise.hs b/compiler/GHC/Core/Opt/Specialise.hs
index 25e4859300..cbf3a4e10e 100644
--- a/compiler/GHC/Core/Opt/Specialise.hs
+++ b/compiler/GHC/Core/Opt/Specialise.hs
@@ -739,8 +739,7 @@ spec_import top_env callers rb dict_binds cis@(CIS fn _)
        ; hsc_env <- getHscEnv
        ; eps <- liftIO $ hscEPS hsc_env
        ; vis_orphs <- getVisibleOrphanMods
-       ; let full_rb = unionRuleBase rb (eps_rule_base eps)
-             rules_for_fn = getRules (RuleEnv full_rb vis_orphs) fn
+       ; let rules_for_fn = getRules (RuleEnv [rb, eps_rule_base eps] vis_orphs) fn
 
        ; (rules1, spec_pairs, MkUD { ud_binds = dict_binds1, ud_calls = new_calls })
             <- -- debugTraceMsg (text "specImport1" <+> vcat [ppr fn, ppr good_calls, ppr rhs]) >>
diff --git a/compiler/GHC/Core/Rules.hs b/compiler/GHC/Core/Rules.hs
index 3db004ae9d..01c14db17a 100644
--- a/compiler/GHC/Core/Rules.hs
+++ b/compiler/GHC/Core/Rules.hs
@@ -10,7 +10,7 @@
 module GHC.Core.Rules (
         -- ** Constructing
         emptyRuleBase, mkRuleBase, extendRuleBaseList,
-        unionRuleBase, pprRuleBase, extendRuleEnv,
+        pprRuleBase, extendRuleEnv,
 
         -- ** Checking rule applications
         ruleCheckProgram,
@@ -317,9 +317,9 @@ rulesOfBinds binds = concatMap (concatMap idCoreRules . bindersOf) binds
 getRules :: RuleEnv -> Id -> [CoreRule]
 -- See Note [Where rules are found]
 getRules (RuleEnv { re_base = rule_base, re_visible_orphs = orphs }) fn
-  = idCoreRules fn ++ filter (ruleIsVisible orphs) imp_rules
+  = idCoreRules fn ++ concatMap imp_rules rule_base
   where
-    imp_rules = lookupNameEnv rule_base (idName fn) `orElse` []
+    imp_rules rb = filter (ruleIsVisible orphs) (lookupNameEnv rb (idName fn) `orElse` [])
 
 ruleIsVisible :: ModuleSet -> CoreRule -> Bool
 ruleIsVisible _ BuiltinRule{} = True
@@ -365,15 +365,12 @@ extendRuleBaseList :: RuleBase -> [CoreRule] -> RuleBase
 extendRuleBaseList rule_base new_guys
   = foldl' extendRuleBase rule_base new_guys
 
-unionRuleBase :: RuleBase -> RuleBase -> RuleBase
-unionRuleBase rb1 rb2 = plusNameEnv_C (++) rb1 rb2
-
 extendRuleBase :: RuleBase -> CoreRule -> RuleBase
 extendRuleBase rule_base rule
   = extendNameEnv_Acc (:) Utils.singleton rule_base (ruleIdName rule) rule
 
 extendRuleEnv :: RuleEnv -> RuleBase -> RuleEnv
-extendRuleEnv (RuleEnv rules orphs) rb = (RuleEnv (rules `unionRuleBase` rb) orphs)
+extendRuleEnv (RuleEnv rules orphs) rb = (RuleEnv (rb:rules) orphs)
 
 pprRuleBase :: RuleBase -> SDoc
 pprRuleBase rules = pprUFM rules $ \rss ->