Decouple dynflags in GHC.Core.Opt.Arity (related to #17957)

Metric Decrease:
    T16875

2 files changed, 31 insertions, 24 deletions

diff --git a/compiler/GHC/Core/Opt/Arity.hs b/compiler/GHC/Core/Opt/Arity.hs
index 7125397637..b318c75f59 100644
--- a/compiler/GHC/Core/Opt/Arity.hs
+++ b/compiler/GHC/Core/Opt/Arity.hs
@@ -15,6 +15,7 @@ module GHC.Core.Opt.Arity
    , exprEtaExpandArity, findRhsArity
    , etaExpand, etaExpandAT
    , exprBotStrictness_maybe
+   , ArityOpts(..)
 
    -- ** ArityType
    , ArityType(..), mkBotArityType, mkTopArityType, expandableArityType
@@ -31,8 +32,6 @@ where
 
 import GHC.Prelude
 
-import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt )
-
 import GHC.Core
 import GHC.Core.FVs
 import GHC.Core.Utils
@@ -622,12 +621,17 @@ takeWhileOneShot (AT oss div)
   | isDeadEndDiv div = AT (takeWhile isOneShotInfo oss) topDiv
   | otherwise        = AT (takeWhile isOneShotInfo oss) div
 
+data ArityOpts = ArityOpts
+  { ao_ped_bot :: !Bool -- See Note [Dealing with bottom]
+  , ao_dicts_cheap :: !Bool -- See Note [Eta expanding through dictionaries]
+  }
+
 -- | The Arity returned is the number of value args the
 -- expression can be applied to without doing much work
-exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
+exprEtaExpandArity :: ArityOpts -> CoreExpr -> ArityType
 -- exprEtaExpandArity is used when eta expanding
 --      e  ==>  \xy -> e x y
-exprEtaExpandArity dflags e = arityType (etaExpandArityEnv dflags) e
+exprEtaExpandArity opts e = arityType (etaExpandArityEnv opts) e
 
 getBotArity :: ArityType -> Maybe Arity
 -- Arity of a divergent function
@@ -636,14 +640,14 @@ getBotArity (AT oss div)
   | otherwise        = Nothing
 
 ----------------------
-findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> ArityType
+findRhsArity :: ArityOpts -> Id -> CoreExpr -> Arity -> ArityType
 -- This implements the fixpoint loop for arity analysis
 -- See Note [Arity analysis]
 -- If findRhsArity e = (n, is_bot) then
 --  (a) any application of e to <n arguments will not do much work,
 --      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
+findRhsArity opts bndr rhs old_arity
   = go 0 botArityType
       -- We always do one step, but usually that produces a result equal to
       -- old_arity, and then we stop right away, because old_arity is assumed
@@ -668,7 +672,7 @@ findRhsArity dflags bndr rhs old_arity
     step at = -- pprTrace "step" (ppr bndr <+> ppr at <+> ppr (arityType env rhs)) $
               arityType env rhs
       where
-        env = extendSigEnv (findRhsArityEnv dflags) bndr at
+        env = extendSigEnv (findRhsArityEnv opts) bndr at
 
 
 {-
@@ -872,12 +876,10 @@ We do this regardless of -fdicts-cheap; it's not really a dictionary.
 data AnalysisMode
   = BotStrictness
   -- ^ Used during 'exprBotStrictness_maybe'.
-  | EtaExpandArity { am_ped_bot :: !Bool
-                   , am_dicts_cheap :: !Bool }
+  | EtaExpandArity { am_opts :: !ArityOpts }
   -- ^ Used for finding an expression's eta-expanding arity quickly, without
   -- fixed-point iteration ('exprEtaExpandArity').
-  | FindRhsArity { am_ped_bot :: !Bool
-                 , am_dicts_cheap :: !Bool
+  | FindRhsArity { am_opts :: !ArityOpts
                  , am_sigs :: !(IdEnv ArityType) }
   -- ^ Used for regular, fixed-point arity analysis ('findRhsArity').
   --   See Note [Arity analysis] for details about fixed-point iteration.
@@ -898,17 +900,15 @@ botStrictnessArityEnv :: ArityEnv
 botStrictnessArityEnv = AE { ae_mode = BotStrictness, ae_joins = emptyVarSet }
 
 -- | The @ArityEnv@ used by 'exprEtaExpandArity'.
-etaExpandArityEnv :: DynFlags -> ArityEnv
-etaExpandArityEnv dflags
-  = AE { ae_mode  = EtaExpandArity { am_ped_bot = gopt Opt_PedanticBottoms dflags
-                                   , am_dicts_cheap = gopt Opt_DictsCheap dflags }
+etaExpandArityEnv :: ArityOpts -> ArityEnv
+etaExpandArityEnv opts
+  = AE { ae_mode  = EtaExpandArity { am_opts = opts }
        , ae_joins = emptyVarSet }
 
 -- | The @ArityEnv@ used by 'findRhsArity'.
-findRhsArityEnv :: DynFlags -> ArityEnv
-findRhsArityEnv dflags
-  = AE { ae_mode  = FindRhsArity { am_ped_bot = gopt Opt_PedanticBottoms dflags
-                                 , am_dicts_cheap = gopt Opt_DictsCheap dflags
+findRhsArityEnv :: ArityOpts -> ArityEnv
+findRhsArityEnv opts
+  = AE { ae_mode  = FindRhsArity { am_opts = opts
                                  , am_sigs = emptyVarEnv }
        , ae_joins = emptyVarSet }
 
@@ -967,8 +967,8 @@ lookupSigEnv AE{ ae_mode = mode } id = case mode of
 pedanticBottoms :: ArityEnv -> Bool
 pedanticBottoms AE{ ae_mode = mode } = case mode of
   BotStrictness                          -> True
-  EtaExpandArity{ am_ped_bot = ped_bot } -> ped_bot
-  FindRhsArity{ am_ped_bot = ped_bot }   -> ped_bot
+  EtaExpandArity{ am_opts = ArityOpts{ ao_ped_bot = ped_bot } } -> ped_bot
+  FindRhsArity{ am_opts = ArityOpts{ ao_ped_bot = ped_bot } }   -> ped_bot
 
 -- | A version of 'exprIsCheap' that considers results from arity analysis
 -- and optionally the expression's type.
@@ -980,7 +980,7 @@ myExprIsCheap AE{ae_mode = mode} e mb_ty = case mode of
     where
       cheap_dict = case mb_ty of
                      Nothing -> False
-                     Just ty -> (am_dicts_cheap mode && isDictTy ty)
+                     Just ty -> (ao_dicts_cheap (am_opts mode) && isDictTy ty)
                                 || isCallStackPredTy ty
         -- See Note [Eta expanding through dictionaries]
         -- See Note [Eta expanding through CallStacks]
diff --git a/compiler/GHC/Core/Opt/Simplify/Utils.hs b/compiler/GHC/Core/Opt/Simplify/Utils.hs
index ce69e35aea..92a2a318d7 100644
--- a/compiler/GHC/Core/Opt/Simplify/Utils.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Utils.hs
@@ -1690,7 +1690,7 @@ mkLam env bndrs body cont
       | not (contIsRhs cont)   -- See Note [Eta expanding lambdas]
       , sm_eta_expand mode
       , any isRuntimeVar bndrs
-      , let body_arity = {-# SCC "eta" #-} exprEtaExpandArity dflags body
+      , let body_arity = {-# SCC "eta" #-} exprEtaExpandArity (initArityOpts dflags) body
       , expandableArityType body_arity
       = do { tick (EtaExpansion (head bndrs))
            ; let res = {-# SCC "eta3" #-}
@@ -1803,9 +1803,10 @@ tryEtaExpandRhs env bndr rhs
     mode      = getMode env
     in_scope  = getInScope env
     dflags    = sm_dflags mode
+    arityOpts = initArityOpts dflags
     old_arity = exprArity rhs
 
-    arity_type = findRhsArity dflags bndr rhs old_arity
+    arity_type = findRhsArity arityOpts bndr rhs old_arity
                  `maxWithArity` idCallArity bndr
     new_arity = arityTypeArity arity_type
 
@@ -1824,6 +1825,12 @@ tryEtaExpandRhs env bndr rhs
                    ABot {}     -> True
 -}
 
+initArityOpts :: DynFlags -> ArityOpts
+initArityOpts dflags = ArityOpts
+  { ao_ped_bot = gopt Opt_PedanticBottoms dflags
+  , ao_dicts_cheap = gopt Opt_DictsCheap dflags
+  }
+
 {-
 Note [Eta-expanding at let bindings]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~