DmdAnal: Keep alive RULE vars in LetUp (#18971)

I also took the liberty to refactor the logic around `ruleFVs`.

3 files changed, 67 insertions, 85 deletions

diff --git a/compiler/GHC/Core/FVs.hs b/compiler/GHC/Core/FVs.hs
index 7fc5c6994c..92386a6d2f 100644
--- a/compiler/GHC/Core/FVs.hs
+++ b/compiler/GHC/Core/FVs.hs
@@ -28,12 +28,13 @@ module GHC.Core.FVs (
         varTypeTyCoFVs,
         idUnfoldingVars, idFreeVars, dIdFreeVars,
         bndrRuleAndUnfoldingVarsDSet,
+        bndrRuleAndUnfoldingIds,
         idFVs,
-        idRuleVars, idRuleRhsVars, stableUnfoldingVars,
+        idRuleVars, stableUnfoldingVars,
         ruleFreeVars, rulesFreeVars,
         rulesFreeVarsDSet, mkRuleInfo,
         ruleLhsFreeIds, ruleLhsFreeIdsList,
-        ruleRhsFreeVars, ruleRhsFreeIds,
+        ruleRhsFreeVars, rulesRhsFreeIds,
 
         expr_fvs,
 
@@ -62,8 +63,6 @@ import GHC.Core
 import GHC.Types.Id
 import GHC.Types.Id.Info
 import GHC.Types.Name.Set
-import GHC.Types.Unique.Set
-import GHC.Types.Unique (Uniquable (..))
 import GHC.Types.Name
 import GHC.Types.Var.Set
 import GHC.Types.Var
@@ -76,7 +75,6 @@ import GHC.Core.FamInstEnv
 import GHC.Builtin.Types( unrestrictedFunTyConName )
 import GHC.Builtin.Types.Prim( funTyConName )
 import GHC.Data.Maybe( orElse )
-import GHC.Types.Basic( Activation )
 
 import GHC.Utils.FV as FV
 import GHC.Utils.Misc
@@ -450,87 +448,71 @@ orph_names_of_fun_ty_con _ = emptyNameSet
 ************************************************************************
 -}
 
+data RuleFVsFrom
+  = LhsOnly
+  | RhsOnly
+  | BothSides
+
+-- | Those locally-defined variables free in the left and/or right hand sides
+-- of the rule, depending on the first argument. Returns an 'FV' computation.
+ruleFVs :: RuleFVsFrom -> CoreRule -> FV
+ruleFVs !_   (BuiltinRule {}) = emptyFV
+ruleFVs from (Rule { ru_fn = _do_not_include
+                     -- See Note [Rule free var hack]
+                   , ru_bndrs = bndrs
+                   , ru_rhs = rhs, ru_args = args })
+  = filterFV isLocalVar $ addBndrs bndrs (exprs_fvs exprs)
+  where
+    exprs = case from of
+      LhsOnly   -> args
+      RhsOnly   -> [rhs]
+      BothSides -> rhs:args
+
+-- | Those locally-defined variables free in the left and/or right hand sides
+-- from several rules, depending on the first argument.
+-- Returns an 'FV' computation.
+rulesFVs :: RuleFVsFrom -> [CoreRule] -> FV
+rulesFVs from = mapUnionFV (ruleFVs from)
+
 -- | Those variables free in the right hand side of a rule returned as a
 -- non-deterministic set
 ruleRhsFreeVars :: CoreRule -> VarSet
-ruleRhsFreeVars (BuiltinRule {}) = noFVs
-ruleRhsFreeVars (Rule { ru_fn = _, ru_bndrs = bndrs, ru_rhs = rhs })
-  = fvVarSet $ filterFV isLocalVar $ addBndrs bndrs (expr_fvs rhs)
-      -- See Note [Rule free var hack]
+ruleRhsFreeVars = fvVarSet . ruleFVs RhsOnly
 
--- | Those variables free in the both the left right hand sides of a rule
+-- | Those locally-defined free 'Id's in the right hand side of several rules
 -- returned as a non-deterministic set
-ruleFreeVars :: CoreRule -> VarSet
-ruleFreeVars = fvVarSet . ruleFVs
+rulesRhsFreeIds :: [CoreRule] -> VarSet
+rulesRhsFreeIds = fvVarSet . filterFV isLocalId . rulesFVs RhsOnly
 
--- | Those variables free in the both the left right hand sides of a rule
--- returned as FV computation
-ruleFVs :: CoreRule -> FV
-ruleFVs (BuiltinRule {}) = emptyFV
-ruleFVs (Rule { ru_fn = _do_not_include
-                  -- See Note [Rule free var hack]
-              , ru_bndrs = bndrs
-              , ru_rhs = rhs, ru_args = args })
-  = filterFV isLocalVar $ addBndrs bndrs (exprs_fvs (rhs:args))
+ruleLhsFreeIds :: CoreRule -> VarSet
+-- ^ This finds all locally-defined free Ids on the left hand side of a rule
+-- and returns them as a non-deterministic set
+ruleLhsFreeIds = fvVarSet . filterFV isLocalId . ruleFVs LhsOnly
 
--- | Those variables free in the both the left right hand sides of rules
--- returned as FV computation
-rulesFVs :: [CoreRule] -> FV
-rulesFVs = mapUnionFV ruleFVs
+ruleLhsFreeIdsList :: CoreRule -> [Var]
+-- ^ This finds all locally-defined free Ids on the left hand side of a rule
+-- and returns them as a deterministically ordered list
+ruleLhsFreeIdsList = fvVarList . filterFV isLocalId . ruleFVs LhsOnly
+
+-- | Those variables free in the both the left right hand sides of a rule
+-- returned as a non-deterministic set
+ruleFreeVars :: CoreRule -> VarSet
+ruleFreeVars = fvVarSet . ruleFVs BothSides
 
 -- | Those variables free in the both the left right hand sides of rules
 -- returned as a deterministic set
 rulesFreeVarsDSet :: [CoreRule] -> DVarSet
-rulesFreeVarsDSet rules = fvDVarSet $ rulesFVs rules
+rulesFreeVarsDSet rules = fvDVarSet $ rulesFVs BothSides rules
+
+-- | Those variables free in both the left right hand sides of several rules
+rulesFreeVars :: [CoreRule] -> VarSet
+rulesFreeVars rules = fvVarSet $ rulesFVs BothSides rules
 
 -- | Make a 'RuleInfo' containing a number of 'CoreRule's, suitable
 -- for putting into an 'IdInfo'
 mkRuleInfo :: [CoreRule] -> RuleInfo
 mkRuleInfo rules = RuleInfo rules (rulesFreeVarsDSet rules)
 
-idRuleRhsVars :: (Activation -> Bool) -> Id -> VarSet
--- Just the variables free on the *rhs* of a rule
-idRuleRhsVars is_active id
-  = mapUnionVarSet get_fvs (idCoreRules id)
-  where
-    get_fvs (Rule { ru_fn = fn, ru_bndrs = bndrs
-                  , ru_rhs = rhs, ru_act = act })
-      | is_active act
-            -- See Note [Finding rule RHS free vars] in "GHC.Core.Opt.OccurAnal"
-      = delOneFromUniqSet_Directly fvs (getUnique fn)
-            -- Note [Rule free var hack]
-      where
-        fvs = fvVarSet $ filterFV isLocalVar $ addBndrs bndrs (expr_fvs rhs)
-    get_fvs _ = noFVs
-
--- | Those variables free in the right hand side of several rules
-rulesFreeVars :: [CoreRule] -> VarSet
-rulesFreeVars rules = mapUnionVarSet ruleFreeVars rules
-
-ruleLhsFreeIds :: CoreRule -> VarSet
--- ^ This finds all locally-defined free Ids on the left hand side of a rule
--- and returns them as a non-deterministic set
-ruleLhsFreeIds = fvVarSet . ruleLhsFVIds
-
-ruleLhsFreeIdsList :: CoreRule -> [Var]
--- ^ This finds all locally-defined free Ids on the left hand side of a rule
--- and returns them as a deterministically ordered list
-ruleLhsFreeIdsList = fvVarList . ruleLhsFVIds
-
-ruleLhsFVIds :: CoreRule -> FV
--- ^ This finds all locally-defined free Ids on the left hand side of a rule
--- and returns an FV computation
-ruleLhsFVIds (BuiltinRule {}) = emptyFV
-ruleLhsFVIds (Rule { ru_bndrs = bndrs, ru_args = args })
-  = filterFV isLocalId $ addBndrs bndrs (exprs_fvs args)
-
-ruleRhsFreeIds :: CoreRule -> VarSet
--- ^ This finds all locally-defined free Ids on the right hand side of a rule
--- and returns them as a non-deterministic set
-ruleRhsFreeIds (BuiltinRule {}) = emptyVarSet
-ruleRhsFreeIds (Rule { ru_bndrs = bndrs, ru_rhs = rhs })
-  = fvVarSet $ filterFV isLocalId $ addBndrs bndrs $ expr_fvs rhs
-
 {-
 Note [Rule free var hack]  (Not a hack any more)
 ~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -586,9 +568,6 @@ freeVarsOf (fvs, _) = fvs
 freeVarsOfAnn :: FVAnn -> DIdSet
 freeVarsOfAnn fvs = fvs
 
-noFVs :: VarSet
-noFVs = emptyVarSet
-
 aFreeVar :: Var -> DVarSet
 aFreeVar = unitDVarSet
 
@@ -660,6 +639,9 @@ idFVs id = ASSERT( isId id)
 bndrRuleAndUnfoldingVarsDSet :: Id -> DVarSet
 bndrRuleAndUnfoldingVarsDSet id = fvDVarSet $ bndrRuleAndUnfoldingFVs id
 
+bndrRuleAndUnfoldingIds :: Id -> IdSet
+bndrRuleAndUnfoldingIds id = fvVarSet $ filterFV isId $ bndrRuleAndUnfoldingFVs id
+
 bndrRuleAndUnfoldingFVs :: Id -> FV
 bndrRuleAndUnfoldingFVs id
   | isId id   = idRuleFVs id `unionFV` idUnfoldingFVs id
diff --git a/compiler/GHC/Core/Opt/DmdAnal.hs b/compiler/GHC/Core/Opt/DmdAnal.hs
index 854d8c586e..de4b435c5f 100644
--- a/compiler/GHC/Core/Opt/DmdAnal.hs
+++ b/compiler/GHC/Core/Opt/DmdAnal.hs
@@ -34,7 +34,7 @@ import GHC.Types.Id
 import GHC.Core.Utils
 import GHC.Core.TyCon
 import GHC.Core.Type
-import GHC.Core.FVs      ( exprFreeIds, ruleRhsFreeIds )
+import GHC.Core.FVs      ( rulesRhsFreeIds, bndrRuleAndUnfoldingIds )
 import GHC.Core.Coercion ( Coercion, coVarsOfCo )
 import GHC.Core.FamInstEnv
 import GHC.Core.Opt.Arity ( typeArity )
@@ -96,7 +96,7 @@ dmdAnalProgram opts fam_envs rules binds
       = dmd_ty
 
     rule_fvs :: IdSet
-    rule_fvs = foldr (unionVarSet . ruleRhsFreeIds) emptyVarSet rules
+    rule_fvs = rulesRhsFreeIds rules
 
 -- | We attach useful (e.g. not 'topDmd') 'idDemandInfo' to top-level bindings
 -- that satisfy this function.
@@ -265,7 +265,10 @@ dmdAnalBindLetUp top_lvl env id rhs anal_body = (final_ty, NonRec id' rhs', body
     (body_ty', id_dmd) = findBndrDmd env notArgOfDfun body_ty id
     id'                = setBindIdDemandInfo top_lvl id id_dmd
     (rhs_ty, rhs')     = dmdAnalStar env (dmdTransformThunkDmd rhs id_dmd) rhs
-    final_ty           = body_ty' `plusDmdType` rhs_ty
+
+    -- See Note [Absence analysis for stable unfoldings and RULES]
+    rule_fvs           = bndrRuleAndUnfoldingIds id
+    final_ty           = body_ty' `plusDmdType` rhs_ty `keepAliveDmdType` rule_fvs
 
 -- | Let bindings can be processed in two ways:
 -- Down (RHS before body) or Up (body before RHS).
@@ -809,21 +812,12 @@ dmdAnalRhsSig top_lvl rec_flag env let_dmd id rhs
                 Recursive    -> reuseEnv rhs_fv
                 NonRecursive -> rhs_fv
 
-    rhs_fv2 = rhs_fv1 `keepAliveDmdEnv` extra_fvs
-    -- Find the RHS free vars of the unfoldings and RULES
     -- See Note [Absence analysis for stable unfoldings and RULES]
-    extra_fvs = foldr (unionVarSet . ruleRhsFreeIds) unf_fvs $
-                idCoreRules id
+    rhs_fv2 = rhs_fv1 `keepAliveDmdEnv` bndrRuleAndUnfoldingIds id
 
     -- See Note [Lazy and unleashable free variables]
     (lazy_fv, sig_fv) = partitionVarEnv isWeakDmd rhs_fv2
 
-    unf = realIdUnfolding id
-    unf_fvs | isStableUnfolding unf
-            , Just unf_body <- maybeUnfoldingTemplate unf
-            = exprFreeIds unf_body
-            | otherwise = emptyVarSet
-
 -- | If given the (local, non-recursive) let-bound 'Id', 'useLetUp' determines
 -- whether we should process the binding up (body before rhs) or down (rhs
 -- before body).
diff --git a/compiler/GHC/Types/Demand.hs b/compiler/GHC/Types/Demand.hs
index 3587618d4d..ba5e5266c9 100644
--- a/compiler/GHC/Types/Demand.hs
+++ b/compiler/GHC/Types/Demand.hs
@@ -55,6 +55,7 @@ module GHC.Types.Demand (
     PlusDmdArg, mkPlusDmdArg, toPlusDmdArg,
     -- ** Other operations
     peelFV, findIdDemand, addDemand, splitDmdTy, deferAfterPreciseException,
+    keepAliveDmdType,
 
     -- * Demand signatures
     StrictSig(..), mkStrictSigForArity, mkClosedStrictSig,
@@ -1196,6 +1197,11 @@ findIdDemand (DmdType fv _ res) id
 deferAfterPreciseException :: DmdType -> DmdType
 deferAfterPreciseException = lubDmdType exnDmdType
 
+-- | See 'keepAliveDmdEnv'.
+keepAliveDmdType :: DmdType -> VarSet -> DmdType
+keepAliveDmdType (DmdType fvs ds res) vars =
+  DmdType (fvs `keepAliveDmdEnv` vars) ds res
+
 {-
 Note [Demand type Divergence]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~