Arity: Refactor fixed-point iteration in GHC.Core.Opt.Arity

Arity analysis used to propagate optimistic arity types during
fixed-point interation through the `ArityEnv`'s `ae_cheap_fun` field,
which is like `GHC.Core.Utils.exprIsCheap`, but also considers the
current iteration's optimistic arity, for the binder in question only.

In #18793, we have seen that this is a problematic design, because it
doesn't allow us to look through PAP bindings of that binder.

Hence this patch refactors to a more traditional form with an explicit
signature environment, in which we record the optimistic `ArityType` of
the binder in question (and at the moment is the *only* binder that is
recorded in the arity environment).

2 files changed, 99 insertions, 70 deletions

diff --git a/compiler/GHC/Core/Opt/Arity.hs b/compiler/GHC/Core/Opt/Arity.hs
index d223a79870..cd2dd5c648 100644
--- a/compiler/GHC/Core/Opt/Arity.hs
+++ b/compiler/GHC/Core/Opt/Arity.hs
@@ -175,13 +175,10 @@ exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
 -- and gives them a suitable strictness signatures.  It's used during
 -- float-out
 exprBotStrictness_maybe e
-  = case getBotArity (arityType env e) of
+  = case getBotArity (arityType botStrictnessArityEnv e) of
         Nothing -> Nothing
         Just ar -> Just (ar, sig ar)
   where
-    env    = AE { ae_ped_bot = True
-                , ae_cheap_fn = \ _ _ -> False
-                , ae_joins = emptyVarSet }
     sig ar = mkClosedStrictSig (replicate ar topDmd) botDiv
 
 {-
@@ -552,34 +549,18 @@ maxWithArity at@(ATop oss) ar
 vanillaArityType :: ArityType
 vanillaArityType = ATop []      -- Totally uninformative
 
--- ^ The Arity returned is the number of value args the
+-- | The Arity returned is the number of value args the
 -- expression can be applied to without doing much work
 exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
 -- exprEtaExpandArity is used when eta expanding
 --      e  ==>  \xy -> e x y
-exprEtaExpandArity dflags e
-  = arityType env e
-  where
-    env = AE { ae_cheap_fn = mk_cheap_fn dflags isCheapApp
-             , ae_ped_bot  = gopt Opt_PedanticBottoms dflags
-             , ae_joins    = emptyVarSet }
+exprEtaExpandArity dflags e = arityType (initArityEnv dflags) e
 
 getBotArity :: ArityType -> Maybe Arity
 -- Arity of a divergent function
 getBotArity (ABot n) = Just n
 getBotArity _        = Nothing
 
-mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun
-mk_cheap_fn dflags cheap_app
-  | not (gopt Opt_DictsCheap dflags)
-  = \e _     -> exprIsCheapX cheap_app e
-  | otherwise
-  = \e mb_ty -> exprIsCheapX cheap_app e
-             || case mb_ty of
-                  Nothing -> False
-                  Just ty -> isDictTy ty
-
-
 ----------------------
 findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> ArityType
 -- This implements the fixpoint loop for arity analysis
@@ -589,20 +570,16 @@ findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> ArityType
 --      so it is safe to expand e  ==>  (\x1..xn. e x1 .. xn)
 --  (b) if is_bot=True, then e applied to n args is guaranteed bottom
 findRhsArity dflags bndr rhs old_arity
-  = go (get_arity init_cheap_app)
-       -- We always call exprEtaExpandArity once, but usually
-       -- that produces a result equal to old_arity, and then
-       -- we stop right away (since arities should not decrease)
-       -- Result: the common case is that there is just one iteration
+  = go (step botArityType)
+      -- We always do one step, but usually that produces a result equal to
+      -- old_arity, and then we stop right away (since arities should not
+      -- decrease)
+      -- Result: the common case is that there is just one iteration
   where
-    init_cheap_app :: CheapAppFun
-    init_cheap_app fn n_val_args
-      | fn == bndr = True   -- On the first pass, this binder gets infinite arity
-      | otherwise  = isCheapApp fn n_val_args
-
     go :: ArityType -> ArityType
+    go cur_atype@(ATop oss)
+      | length oss <= old_arity = cur_atype
     go cur_atype
-      | cur_arity <= old_arity = cur_atype
       | new_atype == cur_atype = cur_atype
       | otherwise =
 #if defined(DEBUG)
@@ -612,20 +589,12 @@ findRhsArity dflags bndr rhs old_arity
 #endif
                     go new_atype
       where
-        new_atype = get_arity cheap_app
-
-        cur_arity = arityTypeArity cur_atype
-        cheap_app :: CheapAppFun
-        cheap_app fn n_val_args
-          | fn == bndr = n_val_args < cur_arity
-          | otherwise  = isCheapApp fn n_val_args
+        new_atype = step cur_atype
 
-    get_arity :: CheapAppFun -> ArityType
-    get_arity cheap_app = arityType env rhs
+    step :: ArityType -> ArityType
+    step at = arityType env rhs
       where
-         env = AE { ae_cheap_fn = mk_cheap_fn dflags cheap_app
-                  , ae_ped_bot  = gopt Opt_PedanticBottoms dflags
-                  , ae_joins    = emptyVarSet }
+        env = extendSigEnv (initArityEnv dflags) bndr at
 
 {-
 Note [Arity analysis]
@@ -757,22 +726,80 @@ encountered a cast, but that is far too conservative: see #5475
 -}
 
 ---------------------------
-type CheapFun = CoreExpr -> Maybe Type -> Bool
-        -- How to decide if an expression is cheap
-        -- If the Maybe is Just, the type is the type
-        -- of the expression; Nothing means "don't know"
+
+data AnalysisMode
+  = BotStrictness
+  -- ^ Used during 'exprBotStrictness_maybe'.
+  | ArityAnalysis { aa_ped_bot :: !Bool
+                  , aa_dicts_cheap :: !Bool
+                  , aa_sigs :: !(IdEnv ArityType) }
+  -- ^ Used for regular arity analysis ('exprEtaExpandArity', 'findRhsArity').
 
 data ArityEnv
-  = AE { ae_cheap_fn :: CheapFun
-       , ae_ped_bot  :: Bool       -- True <=> be pedantic about bottoms
-       , ae_joins    :: IdSet      -- In-scope join points
-                                   -- See Note [Eta-expansion and join points]
+  = AE
+  { ae_mode   :: !AnalysisMode
+  -- ^ The analysis mode. Called during 'exprBotStrictness_maybe' or not?
+  , ae_joins  :: !IdSet
+  -- ^ In-scope join points. See Note [Eta-expansion and join points]
   }
 
+-- | A regular, initial @ArityEnv@ used in arity analysis.
+initArityEnv :: DynFlags -> ArityEnv
+initArityEnv dflags
+  = AE { ae_mode  = ArityAnalysis { aa_ped_bot = gopt Opt_PedanticBottoms dflags
+                                  , aa_dicts_cheap = gopt Opt_DictsCheap dflags
+                                  , aa_sigs = emptyVarEnv }
+       , ae_joins = emptyVarSet }
+
+-- | The @ArityEnv@ used by 'exprBotStrictness_maybe'. Pedantic about bottoms
+-- and no application is ever considered cheap.
+botStrictnessArityEnv :: ArityEnv
+botStrictnessArityEnv = AE { ae_mode = BotStrictness, ae_joins = emptyVarSet }
+
 extendJoinEnv :: ArityEnv -> [JoinId] -> ArityEnv
 extendJoinEnv env@(AE { ae_joins = joins }) join_ids
   = env { ae_joins = joins `extendVarSetList` join_ids }
 
+extendSigEnv :: ArityEnv -> Id -> ArityType -> ArityEnv
+extendSigEnv env id ar_ty = env { ae_mode = go (ae_mode env) }
+  where
+    go BotStrictness = BotStrictness
+    go aa            = aa { aa_sigs = extendVarEnv (aa_sigs aa) id ar_ty }
+
+lookupSigEnv :: ArityEnv -> Id -> Maybe ArityType
+lookupSigEnv AE{ ae_mode = mode } id = case mode of
+  BotStrictness                   -> Nothing
+  ArityAnalysis{ aa_sigs = sigs } -> lookupVarEnv sigs id
+
+-- | Whether the analysis should be pedantic about bottoms.
+-- 'exprBotStrictness_maybe' always is.
+pedanticBottoms :: ArityEnv -> Bool
+pedanticBottoms AE{ ae_mode = mode } = case mode of
+  BotStrictness                         -> True
+  ArityAnalysis{ aa_ped_bot = ped_bot } -> ped_bot
+
+-- | A version of 'exprIsCheap' that considers results from arity analysis
+-- and optionally the expression's type.
+-- Under 'exprBotStrictness_maybe', no expressions are cheap.
+myExprIsCheap :: ArityEnv -> CoreExpr -> Maybe Type -> Bool
+myExprIsCheap AE{ae_mode = mode} e mb_ty = case mode of
+  BotStrictness                                               -> False
+  ArityAnalysis{aa_dicts_cheap = dicts_cheap, aa_sigs = sigs} ->
+    cheap_dict || exprIsCheapX (myIsCheapApp sigs) e
+    where
+      cheap_dict = dicts_cheap && fmap isDictTy mb_ty == Just True
+
+-- | A version of 'isCheapApp' that considers results from arity analysis.
+myIsCheapApp :: IdEnv ArityType -> CheapAppFun
+myIsCheapApp sigs fn n_val_args = case lookupVarEnv sigs fn of
+  -- Nothing means not a local function, fall back to regular
+  -- 'GHC.Core.Utils.isCheapApp'
+  Nothing         -> isCheapApp fn n_val_args
+  -- @Just at@ means local function with @at@ as current ArityType.
+  -- Roughly approximate what 'isCheapApp' is doing.
+  Just (ABot _)   -> True -- See Note [isCheapApp: bottoming functions] in GHC.Core.Utils
+  Just (ATop oss) -> n_val_args < length oss -- Essentially isWorkFreeApp
+
 ----------------
 arityType :: ArityEnv -> CoreExpr -> ArityType
 
@@ -793,6 +820,8 @@ arityType env (Cast e co)
 arityType env (Var v)
   | v `elemVarSet` ae_joins env
   = botArityType  -- See Note [Eta-expansion and join points]
+  | Just at <- lookupSigEnv env v -- Local binding
+  = at
   | otherwise
   = idArityType v
 
@@ -805,7 +834,7 @@ arityType env (Lam x e)
 arityType env (App fun (Type _))
    = arityType env fun
 arityType env (App fun arg )
-   = arityApp (arityType env fun) (ae_cheap_fn env arg Nothing)
+   = arityApp (arityType env fun) (myExprIsCheap env arg Nothing)
 
         -- Case/Let; keep arity if either the expression is cheap
         -- or it's a 1-shot lambda
@@ -825,10 +854,10 @@ arityType env (Case scrut _ _ alts)
              | otherwise -> botArityType  -- if RHS is bottomming
                                           -- See Note [Dealing with bottom (2)]
 
-     ATop as | not (ae_ped_bot env)    -- See Note [Dealing with bottom (3)]
-             , ae_cheap_fn env scrut Nothing -> ATop as
-             | exprOkForSpeculation scrut    -> ATop as
-             | otherwise                     -> ATop (takeWhile isOneShotInfo as)
+     ATop as | not (pedanticBottoms env)  -- See Note [Dealing with bottom (3)]
+             , myExprIsCheap env scrut Nothing -> ATop as
+             | exprOkForSpeculation scrut      -> ATop as
+             | otherwise                       -> ATop (takeWhile isOneShotInfo as)
   where
     alts_type = foldr1 andArityType [arityType env rhs | (_,_,rhs) <- alts]
 
@@ -855,11 +884,12 @@ arityType env (Let (Rec pairs) body)
       = pprPanic "arityType:joinrec" (ppr pairs)
 
 arityType env (Let b e)
-  = floatIn (cheap_bind b) (arityType env e)
+  = floatIn cheap_bind (arityType env e)
   where
-    cheap_bind (NonRec b e) = is_cheap (b,e)
-    cheap_bind (Rec prs)    = all is_cheap prs
-    is_cheap (b,e) = ae_cheap_fn env e (Just (idType b))
+    cheap_bind = case b of
+      NonRec b e -> is_cheap (b,e)
+      Rec prs    -> all is_cheap prs
+    is_cheap (b,e) = myExprIsCheap env e (Just (idType b))
 
 arityType env (Tick t e)
   | not (tickishIsCode t)     = arityType env e
@@ -1743,4 +1773,3 @@ freshEtaId n subst ty
                   -- "OrCoVar" since this can be used to eta-expand
                   -- coercion abstractions
         subst'  = extendTCvInScope subst eta_id'
-
diff --git a/testsuite/tests/simplCore/should_compile/T18231.stderr b/testsuite/tests/simplCore/should_compile/T18231.stderr
index 445192538b..ee5f474423 100644
--- a/testsuite/tests/simplCore/should_compile/T18231.stderr
+++ b/testsuite/tests/simplCore/should_compile/T18231.stderr
@@ -1,6 +1,6 @@
 
 ==================== Tidy Core ====================
-Result size of Tidy Core = {terms: 30, types: 22, coercions: 5, joins: 0/0}
+Result size of Tidy Core = {terms: 24, types: 20, coercions: 5, joins: 0/0}
 
 -- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
 T18231.$trModule4 :: GHC.Prim.Addr#
@@ -23,14 +23,14 @@ T18231.$trModule :: GHC.Types.Module
 T18231.$trModule = GHC.Types.Module T18231.$trModule3 T18231.$trModule1
 
 Rec {
--- RHS size: {terms: 6, types: 1, coercions: 0, joins: 0/0}
-lvl :: GHC.Prim.Int# -> Data.Functor.Identity.Identity ((), Int)
-lvl = \ (x :: GHC.Prim.Int#) -> T18231.m1 (GHC.Types.I# (GHC.Prim.+# x 1#))
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+lvl :: Data.Functor.Identity.Identity ((), Int)
+lvl = lvl
+end Rec }
 
--- RHS size: {terms: 6, types: 3, coercions: 0, joins: 0/0}
+-- RHS size: {terms: 5, types: 3, coercions: 0, joins: 0/0}
 T18231.m1 :: Int -> Data.Functor.Identity.Identity ((), Int)
-T18231.m1 = \ (s1 :: Int) -> case s1 of { GHC.Types.I# x -> lvl x }
-end Rec }
+T18231.m1 = \ (eta2 :: Int) -> case eta2 of { GHC.Types.I# x -> lvl }
 
 -- RHS size: {terms: 1, types: 0, coercions: 5, joins: 0/0}
 m :: State Int ()