Boxity: Don't update Boxity unless worker/wrapper follows (#21754)wip/T21754

A small refactoring in our Core Opt pipeline and some new functions for
transfering argument boxities from one signature to another to facilitate
`Note [Don't change boxity without worker/wrapper]`.

Fixes #21754.

9 files changed, 133 insertions, 17 deletions

diff --git a/compiler/GHC/Core/Opt/DmdAnal.hs b/compiler/GHC/Core/Opt/DmdAnal.hs
index 36c512d656..1263792d05 100644
--- a/compiler/GHC/Core/Opt/DmdAnal.hs
+++ b/compiler/GHC/Core/Opt/DmdAnal.hs
@@ -59,9 +59,18 @@ _ = pprTrace -- Tired of commenting out the import all the time
 
 -- | Options for the demand analysis
 data DmdAnalOpts = DmdAnalOpts
-   { dmd_strict_dicts    :: !Bool -- ^ Use strict dictionaries
-   , dmd_unbox_width     :: !Int  -- ^ Use strict dictionaries
+   { dmd_strict_dicts    :: !Bool
+   -- ^ Value of `-fdicts-strict` (on by default).
+   -- When set, all functons are implicitly strict in dictionary args.
+   , dmd_do_boxity       :: !Bool
+   -- ^ Governs whether the analysis should update boxity signatures.
+   -- See Note [Don't change boxity without worker/wrapper].
+   , dmd_unbox_width     :: !Int
+   -- ^ Value of `-fdmd-unbox-width`.
+   -- See Note [Unboxed demand on function bodies returning small products]
    , dmd_max_worker_args :: !Int
+   -- ^ Value of `-fmax-worker-args`.
+   -- Don't unbox anything if we end up with more than this many args.
    }
 
 -- This is a strict alternative to (,)
@@ -146,6 +155,40 @@ unforced thunks in demand or strictness information; and it is the
 most memory-intensive part of the compilation process, so this added
 seqBinds makes a big difference in peak memory usage.
 
+Note [Don't change boxity without worker/wrapper]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider (T21754)
+  f n = n+1
+  {-# NOINLINE f #-}
+With `-fno-worker-wrapper`, we should not give `f` a boxity signature that says
+that it unboxes its argument! Client modules would never be able to cancel away
+the box for n. Likewise we shouldn't give `f` the CPR property.
+
+Similarly, in the last run of DmdAnal before codegen (which does not have a
+worker/wrapper phase) we should not change boxity in any way. Remember: an
+earlier result of the demand analyser, complete with worker/wrapper, has aleady
+given a demand signature (with boxity info) to the function.
+(The "last run" is mainly there to attach demanded-once info to let-bindings.)
+
+In general, we should not run Note [Boxity analysis] unless worker/wrapper
+follows to exploit the boxity and make sure that calling modules can observe the
+reported boxity.
+
+Hence DmdAnal is configured by a flag `dmd_do_boxity` that is True only
+if worker/wrapper follows after DmdAnal. If it is not set, and the signature
+is not subject to Note [Boxity for bottoming functions], DmdAnal tries
+to transfer over the previous boxity to the new demand signature, in
+`setIdDmdAndBoxSig`.
+
+Why isn't CprAnal configured with a similar flag? Because if we aren't going to
+do worker/wrapper we don't run CPR analysis at all. (see GHC.Core.Opt.Pipeline)
+
+It might be surprising that we only try to preserve *arg* boxity, not boxity on
+FVs. But FV demands won't make it into interface files anyway, so it's a waste
+of energy.
+Besides, W/W zaps the `DmdEnv` portion of a signature, so we don't know the old
+boxity to begin with; see Note [Zapping DmdEnv after Demand Analyzer].
+
 Note [Analysing top-level bindings]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Consider a CoreProgram like
@@ -257,6 +300,16 @@ setBindIdDemandInfo top_lvl id dmd = setIdDemandInfo id $ case top_lvl of
   TopLevel | not (isInterestingTopLevelFn id) -> topDmd
   _                                           -> dmd
 
+-- | Update the demand signature, but be careful not to change boxity info if
+-- `dmd_do_boxity` is True or if the signature is bottom.
+-- See Note [Don't change boxity without worker/wrapper]
+-- and Note [Boxity for bottoming functions].
+setIdDmdAndBoxSig :: DmdAnalOpts -> Id -> DmdSig -> Id
+setIdDmdAndBoxSig opts id sig = setIdDmdSig id $
+  if dmd_do_boxity opts || isBottomingSig sig
+    then sig
+    else transferArgBoxityDmdSig (idDmdSig id) sig
+
 -- | Let bindings can be processed in two ways:
 -- Down (RHS before body) or Up (body before RHS).
 -- This function handles the up variant.
@@ -1018,7 +1071,8 @@ dmdAnalRhsSig top_lvl rec_flag env let_dmd id rhs
 
     sig = mkDmdSigForArity threshold_arity (DmdType sig_fv final_rhs_dmds rhs_div)
 
-    final_id   = id `setIdDmdSig` sig
+    opts       = ae_opts env
+    final_id   = setIdDmdAndBoxSig opts id sig
     !final_env = extendAnalEnv top_lvl env final_id sig
 
     -- See Note [Aggregated demand for cardinality]
@@ -1858,8 +1912,9 @@ dmdFix :: TopLevelFlag
 dmdFix top_lvl env let_dmd orig_pairs
   = loop 1 initial_pairs
   where
+    opts = ae_opts env
     -- See Note [Initialising strictness]
-    initial_pairs | ae_virgin env = [(setIdDmdSig id botSig, rhs) | (id, rhs) <- orig_pairs ]
+    initial_pairs | ae_virgin env = [(setIdDmdAndBoxSig opts id botSig, rhs) | (id, rhs) <- orig_pairs ]
                   | otherwise     = orig_pairs
 
     -- If fixed-point iteration does not yield a result we use this instead
diff --git a/compiler/GHC/Core/Opt/Pipeline.hs b/compiler/GHC/Core/Opt/Pipeline.hs
index 28871d9fb7..6ed1adf84a 100644
--- a/compiler/GHC/Core/Opt/Pipeline.hs
+++ b/compiler/GHC/Core/Opt/Pipeline.hs
@@ -150,7 +150,7 @@ getCoreToDo dflags rule_base extra_vars
     maybe_rule_check phase = runMaybe rule_check (CoreDoRuleCheck phase)
 
     maybe_strictness_before (Phase phase)
-      | phase `elem` strictnessBefore dflags = CoreDoDemand
+      | phase `elem` strictnessBefore dflags = CoreDoDemand False
     maybe_strictness_before _
       = CoreDoNothing
 
@@ -171,8 +171,8 @@ getCoreToDo dflags rule_base extra_vars
     simpl_gently = CoreDoSimplify $ initSimplifyOpts dflags extra_vars max_iter
                                     (initGentleSimplMode dflags) rule_base
 
-    dmd_cpr_ww = if ww_on then [CoreDoDemand,CoreDoCpr,CoreDoWorkerWrapper]
-                          else [CoreDoDemand,CoreDoCpr]
+    dmd_cpr_ww = if ww_on then [CoreDoDemand True,CoreDoCpr,CoreDoWorkerWrapper]
+                          else [CoreDoDemand False] -- NB: No CPR! See Note [Don't change boxity without worker/wrapper]
 
 
     demand_analyser = (CoreDoPasses (
@@ -340,7 +340,7 @@ getCoreToDo dflags rule_base extra_vars
         -- has run at all. See Note [Final Demand Analyser run] in GHC.Core.Opt.DmdAnal
         -- It is EXTREMELY IMPORTANT to run this pass, otherwise execution
         -- can become /exponentially/ more expensive. See #11731, #12996.
-        runWhen (strictness || late_dmd_anal) CoreDoDemand,
+        runWhen (strictness || late_dmd_anal) (CoreDoDemand False),
 
         maybe_rule_check FinalPhase,
 
@@ -491,8 +491,8 @@ doCorePass pass guts = do
     CoreDoExitify             -> {-# SCC "Exitify" #-}
                                  updateBinds exitifyProgram
 
-    CoreDoDemand              -> {-# SCC "DmdAnal" #-}
-                                 updateBindsM (liftIO . dmdAnal logger dflags fam_envs (mg_rules guts))
+    CoreDoDemand before_ww    -> {-# SCC "DmdAnal" #-}
+                                 updateBindsM (liftIO . dmdAnal logger before_ww dflags fam_envs (mg_rules guts))
 
     CoreDoCpr                 -> {-# SCC "CprAnal" #-}
                                  updateBindsM (liftIO . cprAnalProgram logger fam_envs)
@@ -557,10 +557,11 @@ ruleCheckPass current_phase pat guts = do
                         rule_fn (mg_binds guts))
         return guts
 
-dmdAnal :: Logger -> DynFlags -> FamInstEnvs -> [CoreRule] -> CoreProgram -> IO CoreProgram
-dmdAnal logger dflags fam_envs rules binds = do
+dmdAnal :: Logger -> Bool -> DynFlags -> FamInstEnvs -> [CoreRule] -> CoreProgram -> IO CoreProgram
+dmdAnal logger before_ww dflags fam_envs rules binds = do
   let !opts = DmdAnalOpts
                { dmd_strict_dicts    = gopt Opt_DictsStrict dflags
+               , dmd_do_boxity       = before_ww -- only run Boxity Analysis immediately preceding WW
                , dmd_unbox_width     = dmdUnboxWidth dflags
                , dmd_max_worker_args = maxWorkerArgs dflags
                }
diff --git a/compiler/GHC/Core/Opt/Pipeline/Types.hs b/compiler/GHC/Core/Opt/Pipeline/Types.hs
index ff871b08ff..1630506a7d 100644
--- a/compiler/GHC/Core/Opt/Pipeline/Types.hs
+++ b/compiler/GHC/Core/Opt/Pipeline/Types.hs
@@ -45,7 +45,8 @@ data CoreToDo           -- These are diff core-to-core passes,
   | CoreDoStaticArgs
   | CoreDoCallArity
   | CoreDoExitify
-  | CoreDoDemand
+  | CoreDoDemand Bool  -- Bool: Do worker/wrapper afterwards?
+                       -- See Note [Don't change boxity without worker/wrapper]
   | CoreDoCpr
   | CoreDoWorkerWrapper
   | CoreDoSpecialising
@@ -74,7 +75,8 @@ instance Outputable CoreToDo where
   ppr CoreDoStaticArgs         = text "Static argument"
   ppr CoreDoCallArity          = text "Called arity analysis"
   ppr CoreDoExitify            = text "Exitification transformation"
-  ppr CoreDoDemand             = text "Demand analysis"
+  ppr (CoreDoDemand True)      = text "Demand analysis (including Boxity)"
+  ppr (CoreDoDemand False)     = text "Demand analysis"
   ppr CoreDoCpr                = text "Constructed Product Result analysis"
   ppr CoreDoWorkerWrapper      = text "Worker Wrapper binds"
   ppr CoreDoSpecialising       = text "Specialise"
diff --git a/compiler/GHC/Driver/Config/Core/Lint.hs b/compiler/GHC/Driver/Config/Core/Lint.hs
index cde05fa8b7..8ea716a36c 100644
--- a/compiler/GHC/Driver/Config/Core/Lint.hs
+++ b/compiler/GHC/Driver/Config/Core/Lint.hs
@@ -83,7 +83,7 @@ coreDumpFlag CoreLiberateCase         = Just Opt_D_verbose_core2core
 coreDumpFlag CoreDoStaticArgs         = Just Opt_D_verbose_core2core
 coreDumpFlag CoreDoCallArity          = Just Opt_D_dump_call_arity
 coreDumpFlag CoreDoExitify            = Just Opt_D_dump_exitify
-coreDumpFlag CoreDoDemand             = Just Opt_D_dump_stranal
+coreDumpFlag (CoreDoDemand {})        = Just Opt_D_dump_stranal
 coreDumpFlag CoreDoCpr                = Just Opt_D_dump_cpranal
 coreDumpFlag CoreDoWorkerWrapper      = Just Opt_D_dump_worker_wrapper
 coreDumpFlag CoreDoSpecialising       = Just Opt_D_dump_spec
diff --git a/compiler/GHC/Types/Demand.hs b/compiler/GHC/Types/Demand.hs
index 5956340187..71693748e0 100644
--- a/compiler/GHC/Types/Demand.hs
+++ b/compiler/GHC/Types/Demand.hs
@@ -64,7 +64,8 @@ module GHC.Types.Demand (
     -- * Demand signatures
     DmdSig(..), mkDmdSigForArity, mkClosedDmdSig, mkVanillaDmdSig,
     splitDmdSig, dmdSigDmdEnv, hasDemandEnvSig,
-    nopSig, botSig, isNopSig, isDeadEndSig, isDeadEndAppSig, trimBoxityDmdSig,
+    nopSig, botSig, isNopSig, isBottomingSig, isDeadEndSig, isDeadEndAppSig,
+    trimBoxityDmdSig, transferArgBoxityDmdSig,
 
     -- ** Handling arity adjustments
     prependArgsDmdSig, etaConvertDmdSig,
@@ -2147,6 +2148,13 @@ isNopSig (DmdSig ty) = isNopDmdType ty
 isDeadEndSig :: DmdSig -> Bool
 isDeadEndSig (DmdSig (DmdType _ _ res)) = isDeadEndDiv res
 
+-- | True if the signature diverges or throws an imprecise exception in a saturated call.
+-- NB: In constrast to 'isDeadEndSig' this returns False for 'exnDiv'.
+-- See Note [Dead ends]
+-- and Note [Precise vs imprecise exceptions].
+isBottomingSig :: DmdSig -> Bool
+isBottomingSig (DmdSig (DmdType _ _ res)) = res == botDiv
+
 -- | True when the signature indicates all arguments are boxed
 onlyBoxedArguments :: DmdSig -> Bool
 onlyBoxedArguments (DmdSig (DmdType _ dmds _)) = all demandIsBoxed dmds
@@ -2179,6 +2187,38 @@ trimBoxityDmdType (DmdType fvs ds res) =
 trimBoxityDmdSig :: DmdSig -> DmdSig
 trimBoxityDmdSig = coerce trimBoxityDmdType
 
+-- | Transfers the boxity of the left arg to the demand structure of the right
+-- arg. This only makes sense if applied to new and old demands of the same
+-- value.
+transferBoxity :: Demand -> Demand -> Demand
+transferBoxity from to = go_dmd from to
+  where
+    go_dmd (from_n :* from_sd) to_dmd@(to_n :* to_sd)
+      | isAbs from_n || isAbs to_n = to_dmd
+      | otherwise = case (from_sd, to_sd) of
+          (Poly from_b _, Poly _ to_c) ->
+            to_n :* Poly from_b to_c
+          (_, Prod _ to_ds)
+            | Just (from_b, from_ds) <- viewProd (length to_ds) from_sd
+            -> to_n :* mkProd from_b (strictZipWith go_dmd from_ds to_ds)
+          (Prod from_b from_ds, _)
+            | Just (_, to_ds) <- viewProd (length from_ds) to_sd
+            -> to_n :* mkProd from_b (strictZipWith go_dmd from_ds to_ds)
+          _ -> trimBoxity to_dmd
+
+transferArgBoxityDmdType :: DmdType -> DmdType -> DmdType
+transferArgBoxityDmdType _from@(DmdType _ from_ds _) to@(DmdType to_fvs to_ds to_res)
+  | equalLength from_ds to_ds
+  = -- pprTraceWith "transfer" (\r -> ppr _from $$ ppr to $$ ppr r) $
+    DmdType to_fvs -- Only arg boxity! See Note [Don't change boxity without worker/wrapper]
+            (zipWith transferBoxity from_ds to_ds)
+            to_res
+  | otherwise
+  = trimBoxityDmdType to
+
+transferArgBoxityDmdSig :: DmdSig -> DmdSig -> DmdSig
+transferArgBoxityDmdSig = coerce transferArgBoxityDmdType
+
 prependArgsDmdSig :: Int -> DmdSig -> DmdSig
 -- ^ Add extra ('topDmd') arguments to a strictness signature.
 -- In contrast to 'etaConvertDmdSig', this /prepends/ additional argument
diff --git a/testsuite/tests/simplCore/should_compile/spec-inline.stderr b/testsuite/tests/simplCore/should_compile/spec-inline.stderr
index 8705eeacea..97bbeabcc1 100644
--- a/testsuite/tests/simplCore/should_compile/spec-inline.stderr
+++ b/testsuite/tests/simplCore/should_compile/spec-inline.stderr
@@ -143,7 +143,7 @@ Roman.foo1 = GHC.Maybe.Just @Int Roman.foo2
 foo :: Int -> Int
 [GblId,
  Arity=1,
- Str=<1!P(L)>,
+ Str=<1L>,
  Cpr=1,
  Unf=Unf{Src=StableSystem, TopLvl=True, Value=True, ConLike=True,
          WorkFree=True, Expandable=True,
diff --git a/testsuite/tests/stranal/sigs/T21754.hs b/testsuite/tests/stranal/sigs/T21754.hs
new file mode 100644
index 0000000000..37aa9449f8
--- /dev/null
+++ b/testsuite/tests/stranal/sigs/T21754.hs
@@ -0,0 +1,7 @@
+{-# OPTIONS_GHC -fno-worker-wrapper #-}
+
+module Test where
+
+f :: Int -> Int
+f n = n+1
+{-# NOINLINE f #-}
diff --git a/testsuite/tests/stranal/sigs/T21754.stderr b/testsuite/tests/stranal/sigs/T21754.stderr
new file mode 100644
index 0000000000..0370380d8a
--- /dev/null
+++ b/testsuite/tests/stranal/sigs/T21754.stderr
@@ -0,0 +1,10 @@
+
+==================== Strictness signatures ====================
+Test.f: <1L>
+
+
+
+==================== Strictness signatures ====================
+Test.f: <1L>
+
+
diff --git a/testsuite/tests/stranal/sigs/all.T b/testsuite/tests/stranal/sigs/all.T
index 50b8176ce8..73ecf7be57 100644
--- a/testsuite/tests/stranal/sigs/all.T
+++ b/testsuite/tests/stranal/sigs/all.T
@@ -34,5 +34,6 @@ test('T20746b', normal, compile, [''])
 test('T21081', normal, compile, [''])
 test('T21119', normal, compile, [''])
 test('T21717', normal, compile, [''])
+test('T21754', normal, compile, [''])
 test('T21888', normal, compile, [''])
 test('T21888a', normal, compile, [''])