Account for local rules in specImports

As #23024 showed, in GHC.Core.Opt.Specialise.specImports, we were
generating specialisations (a locally-define function) for imported
functions; and then generating specialisations for those
locally-defined functions.  The RULE for the latter should be
attached to the local Id, not put in the rules-for-imported-ids
set.

Fix is easy; similar to what happens in GHC.HsToCore.addExportFlagsAndRules

3 files changed, 65 insertions, 46 deletions

diff --git a/compiler/GHC/Core/Opt/Specialise.hs b/compiler/GHC/Core/Opt/Specialise.hs
index 13bff9f170..f028bd428e 100644
--- a/compiler/GHC/Core/Opt/Specialise.hs
+++ b/compiler/GHC/Core/Opt/Specialise.hs
@@ -64,6 +64,7 @@ import GHC.Unit.Module( Module )
 import GHC.Unit.Module.ModGuts
 import GHC.Core.Unfold
 
+import Data.List( partition )
 import Data.List.NonEmpty ( NonEmpty (..) )
 
 {-
@@ -726,6 +727,33 @@ specialisation (see canSpecImport):
     Specialise even INLINE things; it hasn't inlined yet, so perhaps
     it never will.  Moreover it may have calls inside it that we want
     to specialise
+
+Wrinkle (W1): If we specialise an imported Id M.foo, we make a /local/
+binding $sfoo.  But specImports may further specialise $sfoo. So we end up
+with RULES for both M.foo (imported) and $sfoo (local).  Rules for local
+Ids should be attached to the Ids themselves (see GHC.HsToCore
+Note [Attach rules to local ids]); so we must partition the rules and
+attach the local rules.  That is done in specImports, via addRulesToId.
+
+Note [Glom the bindings if imported functions are specialised]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Suppose we have an imported, *recursive*, INLINABLE function
+   f :: Eq a => a -> a
+   f = /\a \d x. ...(f a d)...
+In the module being compiled we have
+   g x = f (x::Int)
+Now we'll make a specialised function
+   f_spec :: Int -> Int
+   f_spec = \x -> ...(f Int dInt)...
+   {-# RULE  f Int _ = f_spec #-}
+   g = \x. f Int dInt x
+Note that f_spec doesn't look recursive
+After rewriting with the RULE, we get
+   f_spec = \x -> ...(f_spec)...
+BUT since f_spec was non-recursive before it'll *stay* non-recursive.
+The occurrence analyser never turns a NonRec into a Rec.  So we must
+make sure that f_spec is recursive.  Easiest thing is to make all
+the specialisations for imported bindings recursive.
 -}
 
 specImports :: SpecEnv
@@ -740,16 +768,24 @@ specImports top_env (MkUD { ud_binds = dict_binds, ud_calls = calls })
   = do { let env_w_dict_bndrs = top_env `bringFloatedDictsIntoScope` dict_binds
        ; (_env, spec_rules, spec_binds) <- spec_imports env_w_dict_bndrs [] dict_binds calls
 
-             -- Don't forget to wrap the specialized bindings with
-             -- bindings for the needed dictionaries.
-             -- See Note [Wrap bindings returned by specImports]
-             -- and Note [Glom the bindings if imported functions are specialised]
-       ; let final_binds
+             -- Make a Rec: see Note [Glom the bindings if imported functions are specialised]
+             --
+             -- wrapDictBinds: don't forget to wrap the specialized bindings with
+             --   bindings for the needed dictionaries.
+             --   See Note [Wrap bindings returned by specImports]
+             --
+             -- addRulesToId: see Wrinkle (W1) in Note [Specialising imported functions]
+             --               c.f. GHC.HsToCore.addExportFlagsAndRules
+       ; let (rules_for_locals, rules_for_imps) = partition isLocalRule spec_rules
+             local_rule_base = extendRuleBaseList emptyRuleBase rules_for_locals
+             final_binds
                | null spec_binds = wrapDictBinds dict_binds []
-               | otherwise       = [Rec $ flattenBinds $
-                                    wrapDictBinds dict_binds spec_binds]
+               | otherwise       = [Rec $ mapFst (addRulesToId local_rule_base) $
+                                          flattenBinds                          $
+                                          wrapDictBinds dict_binds              $
+                                          spec_binds]
 
-       ; return (spec_rules, final_binds)
+       ; return (rules_for_imps, final_binds)
     }
 
 -- | Specialise a set of calls to imported bindings
@@ -1111,27 +1147,6 @@ And if the call is to the same type, one specialisation is enough.
 Avoiding this recursive specialisation loop is one reason for the
 'callers' stack passed to specImports and specImport.
 
-Note [Glom the bindings if imported functions are specialised]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Suppose we have an imported, *recursive*, INLINABLE function
-   f :: Eq a => a -> a
-   f = /\a \d x. ...(f a d)...
-In the module being compiled we have
-   g x = f (x::Int)
-Now we'll make a specialised function
-   f_spec :: Int -> Int
-   f_spec = \x -> ...(f Int dInt)...
-   {-# RULE  f Int _ = f_spec #-}
-   g = \x. f Int dInt x
-Note that f_spec doesn't look recursive
-After rewriting with the RULE, we get
-   f_spec = \x -> ...(f_spec)...
-BUT since f_spec was non-recursive before it'll *stay* non-recursive.
-The occurrence analyser never turns a NonRec into a Rec.  So we must
-make sure that f_spec is recursive.  Easiest thing is to make all
-the specialisations for imported bindings recursive.
-
-
 
 ************************************************************************
 *                                                                      *
diff --git a/compiler/GHC/Core/Rules.hs b/compiler/GHC/Core/Rules.hs
index df763835cf..67f47e9d9c 100644
--- a/compiler/GHC/Core/Rules.hs
+++ b/compiler/GHC/Core/Rules.hs
@@ -22,7 +22,7 @@ module GHC.Core.Rules (
 
         -- ** Manipulating 'RuleInfo' rules
         extendRuleInfo, addRuleInfo,
-        addIdSpecialisations,
+        addIdSpecialisations, addRulesToId,
 
         -- ** RuleBase and RuleEnv
 
@@ -349,6 +349,14 @@ addIdSpecialisations id rules
   = setIdSpecialisation id $
     extendRuleInfo (idSpecialisation id) rules
 
+addRulesToId :: RuleBase -> Id -> Id
+-- Add rules in the RuleBase to the rules in the Id
+addRulesToId rule_base bndr
+  | Just rules <- lookupNameEnv rule_base (idName bndr)
+  = bndr `addIdSpecialisations` rules
+  | otherwise
+  = bndr
+
 -- | Gather all the rules for locally bound identifiers from the supplied bindings
 rulesOfBinds :: [CoreBind] -> [CoreRule]
 rulesOfBinds binds = concatMap (concatMap idCoreRules . bindersOf) binds
diff --git a/compiler/GHC/HsToCore.hs b/compiler/GHC/HsToCore.hs
index 755fe3e198..57f1b13391 100644
--- a/compiler/GHC/HsToCore.hs
+++ b/compiler/GHC/HsToCore.hs
@@ -362,32 +362,28 @@ deSugarExpr hsc_env tc_expr = do
 addExportFlagsAndRules
     :: Backend -> NameSet -> NameSet -> [CoreRule]
     -> [(Id, t)] -> [(Id, t)]
-addExportFlagsAndRules bcknd exports keep_alive rules = mapFst add_one
+addExportFlagsAndRules bcknd exports keep_alive rules
+  = mapFst (addRulesToId rule_base . add_export_flag)
+        -- addRulesToId: see Note [Attach rules to local ids]
+        -- NB: the binder might have some existing rules,
+        -- arising from specialisation pragmas
+
   where
-    add_one bndr = add_rules name (add_export name bndr)
-       where
-         name = idName bndr
 
     ---------- Rules --------
-        -- See Note [Attach rules to local ids]
-        -- NB: the binder might have some existing rules,
-        -- arising from specialisation pragmas
-    add_rules name bndr
-        | Just rules <- lookupNameEnv rule_base name
-        = bndr `addIdSpecialisations` rules
-        | otherwise
-        = bndr
     rule_base = extendRuleBaseList emptyRuleBase rules
 
     ---------- Export flag --------
     -- See Note [Adding export flags]
-    add_export name bndr
-        | dont_discard name = setIdExported bndr
+    add_export_flag bndr
+        | dont_discard bndr = setIdExported bndr
         | otherwise         = bndr
 
-    dont_discard :: Name -> Bool
-    dont_discard name = is_exported name
+    dont_discard :: Id -> Bool
+    dont_discard bndr = is_exported name
                      || name `elemNameSet` keep_alive
+       where
+         name = idName bndr
 
         -- In interactive mode, we don't want to discard any top-level
         -- entities at all (eg. do not inline them away during