Very hackily set the stable unfolding if there are static args and keep

on re-setting it

4 files changed, 45 insertions, 46 deletions

diff --git a/compiler/GHC/Core/Opt/OccurAnal.hs b/compiler/GHC/Core/Opt/OccurAnal.hs
index 56eabfc684..cd31f53804 100644
--- a/compiler/GHC/Core/Opt/OccurAnal.hs
+++ b/compiler/GHC/Core/Opt/OccurAnal.hs
@@ -1472,7 +1472,7 @@ nodeScore env new_bndr lb_deps
   = (0, 0, True)                   -- a NOINLINE pragma) makes a great loop breaker
 
   | is_self_rec          -- Self-recursive things are great loop breakers
-  , not has_static_args  -- See Note [Self-recursion and loop breakers]
+--  , not has_static_args  -- See Note [Self-recursion and loop breakers]
   = (0, 0, True)
 
   | exprIsTrivial rhs
diff --git a/compiler/GHC/Core/Opt/Simplify.hs b/compiler/GHC/Core/Opt/Simplify.hs
index e4146a86bc..c485bd1cfe 100644
--- a/compiler/GHC/Core/Opt/Simplify.hs
+++ b/compiler/GHC/Core/Opt/Simplify.hs
@@ -3807,8 +3807,6 @@ simplLetUnfolding env top_lvl cont_mb id new_rhs rhs_ty arity unf
   = simplStableUnfolding env top_lvl cont_mb id rhs_ty arity unf
   | isExitJoinId id
   = return noUnfolding -- See Note [Do not inline exit join points] in GHC.Core.Opt.Exitify
-  | Just mk_unf <- mkSatUnfolding (seUnfoldingOpts env) top_lvl InlineRhs id new_rhs
-  = mk_unf
   | otherwise
   = mkLetUnfolding (seUnfoldingOpts env) top_lvl InlineRhs id new_rhs
 
@@ -3820,36 +3818,15 @@ mkLetUnfolding uf_opts top_lvl src id new_rhs
     return (mkUnfolding uf_opts src is_top_lvl is_bottoming new_rhs)
             -- We make an  unfolding *even for loop-breakers*.
             -- Reason: (a) It might be useful to know that they are WHNF
-            --         (b) They might have static arguments, in which case we
-            --             provide a non-rec unfolding that specialises for those
-            --         (c) And even without static arguments, in GHC.Iface.Tidy we
-            --             currently assume that, if we want to expose the unfolding
-            --             then indeed we *have* an unfolding to expose. (We could
-            --             instead use the RHS, but currently we don't.) The simple
-            --             thing is always to have one.
+            --         (c) In GHC.Iface.Tidy we currently assume that, if we
+            --             want to expose the unfolding then indeed we *have*
+            --             an unfolding to expose. (We could instead use the
+            --             RHS, but currently we don't.) The simple thing is
+            --             always to have one.
   where
     is_top_lvl   = isTopLevel top_lvl
     is_bottoming = isDeadEndId id
 
-mkSatUnfolding :: UnfoldingOpts -> TopLevelFlag -> UnfoldingSource
-               -> InId -> OutExpr -> Maybe (SimplM Unfolding)
-mkSatUnfolding opts top_lvl src id new_rhs
-  | (lam_bndrs, lam_body) <- collectBinders new_rhs
-  , static_args <- takeStaticArgs (length lam_bndrs) $ idStaticArgs id
-  , static_args /= noStaticArgs
-  = Just $ do  { unf_rhs <- return $ saTransform id (getStaticArgs static_args) lam_bndrs lam_body
-               ; return (mkCoreUnfolding src is_top_lvl unf_rhs guidance) }
-  | otherwise
-  = Nothing
-  where
-    is_top_lvl = isTopLevel top_lvl
-    !is_top_bottoming = is_top_lvl && isDeadEndId id
-    guidance          = calcUnfoldingGuidance opts is_top_bottoming new_rhs
-        -- NB: we calculate the guidance from the untransformed new_rhs!
-        -- The SAT'd unfolding will introduce new let binds and lambdas that
-        -- all go away, but will be counted by 'calcUnfoldingGuidance', which
-        -- penalises SAT'd unfoldings too much.
-
 -------------------
 simplStableUnfolding :: SimplEnv -> TopLevelFlag
                      -> MaybeJoinCont  -- Just k => a join point with continuation k
diff --git a/compiler/GHC/Core/Opt/StaticArgs.hs b/compiler/GHC/Core/Opt/StaticArgs.hs
index 0dd6ed8601..1ee130811c 100644
--- a/compiler/GHC/Core/Opt/StaticArgs.hs
+++ b/compiler/GHC/Core/Opt/StaticArgs.hs
@@ -62,6 +62,8 @@ import GHC.Core
 import GHC.Core.Utils
 import GHC.Core.Type
 import GHC.Core.Coercion
+import GHC.Core.Unfold
+import GHC.Core.Unfold.Make
 import GHC.Types.Id
 import GHC.Types.Name
 import GHC.Types.Var.Env
@@ -81,7 +83,7 @@ import Data.Bifunctor (second)
 -- import GHC.Driver.Ppr
 
 satAnalProgram :: CoreProgram -> CoreProgram
-satAnalProgram bs = map (snd . satAnalBind initSatEnv) bs
+satAnalProgram bs = map (snd . satAnalBind initSatEnv TopLevel) bs
 
 -- | Lambda binders ('TyVar's, 'CoVar's and 'Id's) of a let-bound RHS, thus
 -- parameters to a function.
@@ -151,11 +153,11 @@ peelSatOccs :: SatOccs -> Id -> (StaticArgs, SatOccs)
 peelSatOccs (SO env) fn = case delLookupVarEnv env fn of
   (mb_sa, env') -> (mb_sa `orElse` noStaticArgs, SO env')
 
-satAnalBind :: SatEnv -> CoreBind -> (SatOccs, CoreBind)
-satAnalBind env (NonRec id rhs) = (occs, NonRec id rhs')
+satAnalBind :: SatEnv -> TopLevelFlag -> CoreBind -> (SatOccs, CoreBind)
+satAnalBind env top_lvl (NonRec id rhs) = (occs, NonRec id rhs')
   where
     (occs, rhs') = satAnalExpr (env `addInScopeVar` id) rhs
-satAnalBind env (Rec pairs) = (combineSatOccsList occss, Rec pairs')
+satAnalBind env top_lvl (Rec pairs) = (combineSatOccsList occss, Rec pairs')
   where
     (occss, pairs') = mapAndUnzip anal_one pairs
     anal_one (fn, rhs) = (occs', (fn', rhs'))
@@ -167,9 +169,16 @@ satAnalBind env (Rec pairs) = (combineSatOccsList occss, Rec pairs')
         (occs, rhs_body')    = satAnalExpr env2 rhs_body
         rhs'                 = mkLams bndrs rhs_body'
         (static_args, occs') = peelSatOccs occs fn
-        !fn' | allStaticAndUnliftedBody (length bndrs) static_args rhs_body = fn -- otherwise we end up with an unlifted worker body
+        !fn' | allStaticAndUnliftedBody (length bndrs) static_args rhs_body
+             || (isStableUnfolding (realIdUnfolding fn) && idStaticArgs fn /= noStaticArgs)
+             = fn -- otherwise we end up with an unlifted worker body
              | otherwise     = -- pprTrace "satAnalBind:set" (ppr fn $$ ppr static_args) $
-                               setIdStaticArgs fn static_args
+                               fn `setIdStaticArgs` static_args
+                                  `setIdUnfolding` sat_unf
+        sat_unf
+          | Just unf <- mkSatUnfolding top_lvl fn rhs = unf
+          | otherwise = realIdUnfolding fn
+
 
 allStaticAndUnliftedBody :: Arity -> StaticArgs -> CoreExpr -> Bool
 allStaticAndUnliftedBody arty sa body =
@@ -193,7 +202,7 @@ satAnalExpr env e@Lam{}        = (occs, mkLams bndrs body')
     (occs, body') = satAnalExpr (env `addInScopeVars` bndrs) body
 satAnalExpr env (Let bnd body) = (occs, Let bnd' body')
   where
-    (occs_bind, bnd')  = satAnalBind env bnd
+    (occs_bind, bnd')  = satAnalBind env NotTopLevel bnd
     (occs_body, body') = satAnalExpr (env `addInScopeVars` bindersOf bnd) body
     !occs              = combineSatOccs occs_body occs_bind
 satAnalExpr env (Case scrut bndr ty alts) = (occs, Case scrut' bndr ty alts')
@@ -601,3 +610,24 @@ saTransform binder arg_staticness rhs_binders rhs_body
 isStaticValue :: Staticness App -> Bool
 isStaticValue (Static (VarApp _)) = True
 isStaticValue _                   = False
+
+mkSatUnfolding :: TopLevelFlag -> Id -> CoreExpr -> Maybe Unfolding
+mkSatUnfolding top_lvl id new_rhs
+  | (lam_bndrs, lam_body) <- collectBinders new_rhs
+  , static_args <- takeStaticArgs (length lam_bndrs) $ idStaticArgs id
+  , static_args /= noStaticArgs
+  , let unf_rhs = saTransform id (getStaticArgs static_args) lam_bndrs lam_body
+  , let unf = mkCoreUnfolding InlineStable is_top_lvl unf_rhs guidance
+  = Just unf
+  | otherwise
+  = Nothing
+  where
+    is_top_lvl = isTopLevel top_lvl
+    !is_top_bottoming = is_top_lvl && isDeadEndId id
+    opts              = defaultUnfoldingOpts
+    guidance          = calcUnfoldingGuidance opts is_top_bottoming new_rhs
+        -- NB: we calculate the guidance from the untransformed new_rhs!
+        -- The SAT'd unfolding will introduce new let binds and lambdas that
+        -- all go away, but will be counted by 'calcUnfoldingGuidance', which
+        -- penalises SAT'd unfoldings too much.
+
diff --git a/compiler/GHC/Core/Unfold.hs b/compiler/GHC/Core/Unfold.hs
index cc6459c92c..244de58c5d 100644
--- a/compiler/GHC/Core/Unfold.hs
+++ b/compiler/GHC/Core/Unfold.hs
@@ -1076,16 +1076,11 @@ instance Outputable CallCtxt where
   ppr RuleArgCtxt = text "RuleArgCtxt"
 
 callSiteInline dflags id active_unfolding lone_variable arg_infos cont_info
-  = case unfoldingInfo (idInfo id) of
-      -- Don't call idUnfolding, because we want unfolding for loop-breakers
-      -- if they have static arguments
-      -- Things with an INLINE pragma may have an unfolding *and*
-      -- be a loop breaker  (maybe the knot is not yet untied)
+  = case idUnfolding id of
         CoreUnfolding { uf_tmpl = unf_template
                       , uf_is_work_free = is_wf
                       , uf_guidance = guidance, uf_expandable = is_exp }
           | active_unfolding
-          , isStrongLoopBreaker (idOccInfo id) ==> has_static_args id
           -> tryUnfolding dflags id lone_variable
                                     arg_infos cont_info unf_template
                                     is_wf is_exp guidance
@@ -1094,9 +1089,6 @@ callSiteInline dflags id active_unfolding lone_variable arg_infos cont_info
         BootUnfolding    -> Nothing
         OtherCon {}      -> Nothing
         DFunUnfolding {} -> Nothing     -- Never unfold a DFun
-  where
-    b ==> t = not b || t
-    has_static_args id = idStaticArgs id /= noStaticArgs
 
 -- | Report the inlining of an identifier's RHS to the user, if requested.
 traceInline :: DynFlags -> Id -> String -> SDoc -> a -> a
@@ -1136,7 +1128,7 @@ tryUnfolding dflags id lone_variable
           enough_args = (n_val_args >= uf_arity) || (unsat_ok && n_val_args > 0)
 
      UnfIfGoodArgs { ug_args = arg_discounts, ug_res = res_discount, ug_size = size }
-        | unfoldingVeryAggressive uf_opts || idStaticArgs id /= noStaticArgs
+        | unfoldingVeryAggressive uf_opts
         -> traceInline dflags id str (mk_doc some_benefit extra_doc True) (Just unf_template)
         | is_wf && some_benefit && small_enough
         -> traceInline dflags id str (mk_doc some_benefit extra_doc True) (Just unf_template)