Improve trimming of auto-rules

I hadn't got the new function trimAutoRules quite right, so we had
a left-over rule which mentioned a local variable whose binding had
been discarded.  (Result: crash when compiling Haddock.)

This patch merges trimAutoRules into an expanded version of
findExternalRules, gets it right, and adds lots of comments.

See Note [Finding external rules].

And indeed in one regression test we get to trim off more rules
(and hence code) than before.

compiler/main/TidyPgm.lhs

@@ -56,13 +56,12 @@ import ErrUtils (Severity(..))
 import Outputable
 import FastBool hiding ( fastOr )
 import SrcLoc
-import Util
 import FastString
 import qualified ErrUtils as Err
 
 import Control.Monad
 import Data.Function
-import Data.List        ( sortBy, partition )
+import Data.List        ( sortBy )
 import Data.IORef       ( atomicModifyIORef )
 \end{code}
 
@@ -330,12 +329,8 @@ tidyProgram hsc_env  (ModGuts { mg_module    = mod
         ; (unfold_env, tidy_occ_env)
               <- chooseExternalIds hsc_env mod omit_prags expose_all
                                    binds implicit_binds imp_rules (vectInfoVar vect_info)
-        ; let { ext_rules = findExternalRules omit_prags binds imp_rules unfold_env }
-                -- Glom together imp_rules and rules currently attached to binders
-                -- Then pick just the ones we need to expose
-                -- See Note [Which rules to expose]
-
-        ; let { (trimmed_binds, trimmed_rules) = trimAutoRules binds ext_rules }
+        ; let { (trimmed_binds, trimmed_rules) 
+                    = findExternalRules omit_prags binds imp_rules unfold_env }
 
         ; (tidy_env, tidy_binds)
                  <- tidyTopBinds hsc_env mod unfold_env tidy_occ_env trimmed_binds
@@ -425,65 +420,8 @@ lookup_aux_id type_env id
         _other          -> pprPanic "lookup_axu_id" (ppr id)
 \end{code}
 
-Note [Trimming auto rules]
-~~~~~~~~~~~~~~~~~~~~~~~~~~
-With auto-specialisation we may specialise local or imported dfuns or
-INLINE functions, and then later inline them.  That may leave behind
-something like
-   RULE "foo" forall d. f @ Int d = f_spec
-where there is no remaining reference to f_spec except from the RULE.
-
-Now that RULE *might* be useful to an importing module, but that is
-purely speculative, and meanwhile the code is taking up space and
-codegen time.  So is seeems better to drop the bidign for f_spec if
-the auto-generated rule is the only reason that it is being kept
-alive.
-
-Notice, though, that the RULE still might have been useful; that is,
-it was the right thing to have generated it in the first place.  See
-Note [Inline specialisations] in Specialise.  But now it has served
-its purpose, and can be discarded.
-
-So trimAutoRules does this:
-  * Remove all bindings that are kept alive *only* by isAutoRule rules
-  * Remove all auto rules that mention bindings that have been removed
-So if a binding is kept alive for some other reason (e.g. f_spec is
-called in the final code), we keep the rule too.
-
-I found that binary sizes jumped by 6-10% when I started to specialise
-INLINE functions (again, Note [Inline specialisations] in Specialise).
-Adding trimAutoRules removed all this bloat.
-
 
 \begin{code}
-trimAutoRules :: [CoreBind] -> [CoreRule] -> ([CoreBind], [CoreRule])
--- See Note [Trimming auto rules]
-trimAutoRules binds rules
-  | True {- null auto_rules -}  -- Temporrary fix
-  = (binds, rules)
-  | otherwise
-  = (binds', filter keep_rule auto_rules ++ user_rules)
-  where
-    (auto_rules, user_rules) = partition isAutoRule rules
-    rule_fvs = foldr (unionVarSet . ruleRhsFreeVars) emptyVarSet user_rules
-
-    (all_fvs, binds') = trim_binds binds
-
-    trim_binds :: [CoreBind] -> (VarSet, [CoreBind])
-    trim_binds []
-       = (rule_fvs, [])
-    trim_binds (bind:binds)
-       | keep_bind = (fvs `unionVarSet` bind_fvs, bind:binds')
-       | otherwise = (fvs, binds')
-       where
-         needed bndr   = isExportedId bndr || bndr `elemVarSet` fvs
-         keep_bind     = any needed (bindersOf bind)
-         (fvs, binds') = trim_binds binds
-         bind_fvs      = bindFreeVars bind
-
-    keep_rule rule = ruleFreeVars rule `subVarSet` all_fvs
-
-----------------------
 tidyTypeEnv :: Bool       -- Compiling without -O, so omit prags
             -> TypeEnv -> TypeEnv
 
@@ -618,7 +556,7 @@ Oh: two other reasons for injecting them late:
 
   - If implicit Ids are already in the bindings when we start TidyPgm,
     we'd have to be careful not to treat them as external Ids (in
-    the sense of findExternalIds); else the Ids mentioned in *their*
+    the sense of chooseExternalIds); else the Ids mentioned in *their*
     RHSs will be treated as external and you get an interface file
     saying      a18 = <blah>
     but nothing refererring to a18 (because the implicit Id is the
@@ -697,7 +635,7 @@ chooseExternalIds hsc_env mod omit_prags expose_all binds implicit_binds imp_id_
   -- (c) it is the vectorised version of an imported Id
   -- See Note [Which rules to expose]
   is_external id = isExportedId id || id `elemVarSet` rule_rhs_vars || id `elemVarSet` vect_var_vs
-  rule_rhs_vars  = foldr (unionVarSet . ruleRhsFreeVars) emptyVarSet imp_id_rules
+  rule_rhs_vars  = listFVs ruleRhsFreeVars imp_id_rules emptyVarSet
   vect_var_vs    = mkVarSet [var_v | (var, var_v) <- nameEnvElts vect_vars, isGlobalId var]
 
   binders          = bindersOfBinds binds
@@ -913,6 +851,152 @@ dffvLetBndr vanilla_unfold id
 \end{code}
 
 
+%************************************************************************
+%*                                                                      *
+               findExternalRules
+%*                                                                      *
+%************************************************************************
+
+Note [Finding external rules]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The complete rules are gotten by combining
+   a) local rules for imported Ids
+   b) rules embedded in the top-level Ids
+
+There are two complications:
+  * Note [Which rules to expose]
+  * Note [Trimming auto-rules]
+
+Note [Which rules to expose]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The function 'expose_rule' filters out rules that mention, on the LHS,
+Ids that aren't externally visible; these rules can't fire in a client
+module.
+
+The externally-visible binders are computed (by chooseExternalIds)
+assuming that all orphan rules are externalised (see init_ext_ids in
+function 'search'). So in fact it's a bit conservative and we may
+export more than we need.  (It's a sort of mutual recursion.)
+
+Note [Trimming auto-rules]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Second, with auto-specialisation we may specialise local or imported
+dfuns or INLINE functions, and then later inline them.  That may leave
+behind something like
+   RULE "foo" forall d. f @ Int d = f_spec
+where f is either local or imported, and there is no remaining
+reference to f_spec except from the RULE.
+
+Now that RULE *might* be useful to an importing module, but that is
+purely speculative, and meanwhile the code is taking up space and
+codegen time.  So is seeems better to drop the binding for f_spec if
+the auto-generated rule is the *only* reason that it is being kept
+alive.
+
+(The RULE still might have been useful in the past; that is, it was
+the right thing to have generated it in the first place.  See Note
+[Inline specialisations] in Specialise.  But now it has served its
+purpose, and can be discarded.)
+
+So findExternalRules does this:
+  * Remove all bindings that are kept alive *only* by isAutoRule rules
+      (this is done in trim_binds)
+  * Remove all auto rules that mention bindings that have been removed
+      (this is done by filtering by keep_rule)
+
+So if a binding is kept alive for some *other* reason (e.g. f_spec is
+called in the final code), we keep the rule too.
+
+I found that binary sizes jumped by 6-10% when I started to specialise
+INLINE functions (again, Note [Inline specialisations] in Specialise).
+Adding trimAutoRules removed all this bloat.
+
+
+\begin{code}
+findExternalRules :: Bool       -- Omit pragmas
+                  -> [CoreBind]
+                  -> [CoreRule] -- Local rules for imported fns
+                  -> UnfoldEnv  -- Ids that are exported, so we need their rules
+                  -> ([CoreBind], [CoreRule])
+-- See Note [Finding external rules]
+findExternalRules omit_prags binds imp_id_rules unfold_env
+  = (trimmed_binds, filter keep_rule all_rules)
+  where
+    imp_rules         = filter expose_rule imp_id_rules
+    imp_user_rule_fvs = listFVs user_rule_rhs_fvs imp_rules emptyVarSet
+
+    user_rule_rhs_fvs rule | isAutoRule rule = emptyVarSet
+                           | otherwise       = ruleRhsFreeVars rule
+
+    (trimmed_binds, local_bndrs, _, all_rules) = trim_binds binds
+
+    keep_rule rule = ruleFreeVars rule `subVarSet` local_bndrs
+        -- Remove rules that make no sense, because they mention a
+        -- local binder (on LHS or RHS) that we have now discarded.
+        -- (NB: ruleFreeVars only includes LocalIds)
+        --
+        -- LHS: we have alrady filtered out rules that mention internal Ids
+        --     on LHS but that isn't enough because we might have by now
+        --     discarded a binding with an external Id. (How?
+        --     chooseExternalIds is a bit conservative.)
+        --
+        -- RHS: the auto rules that might mention a binder that has
+        --      been discarded; see Note [Trimming auto-rules]
+
+    expose_rule rule
+        | omit_prags = False
+        | otherwise  = all is_external_id (varSetElems (ruleLhsFreeIds rule))
+                -- Don't expose a rule whose LHS mentions a locally-defined
+                -- Id that is completely internal (i.e. not visible to an
+                -- importing module).  NB: ruleLhsFreeIds only returns LocalIds.
+                -- See Note [Which rules to expose]
+
+    is_external_id id = case lookupVarEnv unfold_env id of
+                          Just (name, _) -> isExternalName name
+                          Nothing        -> False
+
+    trim_binds :: [CoreBind]
+               -> ( [CoreBind]   -- Trimmed bindings
+                  , VarSet       -- Binders of those bindings
+                  , VarSet       -- Free vars of those bindings + rhs of user rules
+                                 -- (we don't bother to delete the binders)
+                  , [CoreRule])  -- All rules, imported + from the bindings
+    -- This function removes unnecessary bindings, and gathers up rules from
+    -- the bindings we keep.  See Note [Trimming auto-rules]
+    trim_binds []  -- Base case, start with imp_user_rule_fvs
+       = ([], emptyVarSet, imp_user_rule_fvs, imp_rules)
+
+    trim_binds (bind:binds)
+       | any needed bndrs    -- Keep binding
+       = ( bind : binds', bndr_set', needed_fvs', local_rules ++ rules )
+       | otherwise           -- Discard binding altogether
+       = stuff
+       where
+         stuff@(binds', bndr_set, needed_fvs, rules)
+                       = trim_binds binds
+         needed bndr   = isExportedId bndr || bndr `elemVarSet` needed_fvs
+
+         bndrs         = bindersOf  bind
+         rhss          = rhssOfBind bind
+         bndr_set'     = bndr_set `extendVarSetList` bndrs
+
+         needed_fvs'   = listFVs idUnfoldingVars   bndrs       $
+                              -- Ignore type variables in the type of bndrs
+                         listFVs exprFreeVars      rhss        $
+                         listFVs user_rule_rhs_fvs local_rules $
+                         needed_fvs
+            -- In needed_fvs', we don't bother to delete binders from the fv set
+
+         local_rules  = [ rule
+                        | id <- bndrs
+                        , is_external_id id   -- Only collect rules for external Ids
+                        , rule <- idCoreRules id
+                        , expose_rule rule ]  -- and ones that can fire in a client
+
+listFVs :: (a -> VarSet) -> [a] -> VarSet -> VarSet
+listFVs fv_fn xs fvs = foldr (unionVarSet . fv_fn) fvs xs
+\end{code}
+
 %************************************************************************
 %*                                                                      *
                tidyTopName
@@ -994,44 +1078,6 @@ tidyTopName mod nc_var maybe_ref occ_env id
         -- use the same name for externally-visible things as we did before.
 \end{code}
 
-\begin{code}
-findExternalRules :: Bool       -- Omit pragmas
-                  -> [CoreBind]
-                  -> [CoreRule] -- Local rules for imported fns
-                  -> UnfoldEnv  -- Ids that are exported, so we need their rules
-                  -> [CoreRule]
-  -- The complete rules are gotten by combining
-  --    a) local rules for imported Ids
-  --    b) rules embedded in the top-level Ids
-findExternalRules omit_prags binds imp_id_rules unfold_env
-  | omit_prags = []
-  | otherwise  = filterOut internal_rule (imp_id_rules ++ local_rules)
-  where
-    local_rules  = [ rule
-                   | id <- bindersOfBinds binds,
-                     external_id id,
-                     rule <- idCoreRules id
-                   ]
-
-    internal_rule rule
-        =  any (not . external_id) (varSetElems (ruleLhsFreeIds rule))
-                -- Don't export a rule whose LHS mentions a locally-defined
-                --  Id that is completely internal (i.e. not visible to an
-                -- importing module)
-
-    external_id id
-      | Just (name,_) <- lookupVarEnv unfold_env id = isExternalName name
-      | otherwise = False
-\end{code}
-
-Note [Which rules to expose]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-findExternalRules filters imp_rules to avoid binders that
-aren't externally visible; but the externally-visible binders
-are computed (by findExternalIds) assuming that all orphan
-rules are externalised (see init_ext_ids in function
-'search'). So in fact we may export more than we need.
-(It's a sort of mutual recursion.)
 
 %************************************************************************
 %*                                                                      *

testsuite/tests/simplCore/should_run/T2486.stderr

 
 ==================== Tidy Core rules ====================
-"SPEC Main.fib @ GHC.Types.Double" [ALWAYS]
-    forall ($dNum :: Num Double) ($dOrd :: Ord Double).
-      fib @ Double $dNum $dOrd
-      = fib_$sfib1
-"SPEC Main.fib @ GHC.Types.Int" [ALWAYS]
-    forall ($dNum :: Num Int) ($dOrd :: Ord Int).
-      fib @ Int $dNum $dOrd
-      = fib_$sfib
-"SPEC Main.tak @ GHC.Types.Double" [ALWAYS]
-    forall ($dNum :: Num Double) ($dOrd :: Ord Double).
-      tak @ Double $dNum $dOrd
-      = tak_$stak1
-"SPEC Main.tak @ GHC.Types.Int" [ALWAYS]
-    forall ($dNum :: Num Int) ($dOrd :: Ord Int).
-      tak @ Int $dNum $dOrd
-      = tak_$stak