CPR: Detect constructed products in `runRW#` apps (#19822)

In #19822, we realised that the Simplifier's new habit of floating cases into
`runRW#` continuations inhibits CPR analysis from giving key functions of `text`
the CPR property, such as `singleton`.

This patch fixes that by anticipating part of !5667 (Nested CPR) to give
`runRW#` the proper CPR transformer it now deserves: Namely, `runRW# (\s -> e)`
should have the CPR property iff `e` has it.

The details are in `Note [Simplification of runRW#]` in GHC.CoreToStg.Prep.

The output of T18086 changed a bit: `panic` (which calls `runRW#`) now has
`botCpr`. As outlined in Note [Bottom CPR iff Dead-Ending Divergence], that's
OK.

Fixes #19822.

Metric Decrease:
    T9872d

7 files changed, 75 insertions, 27 deletions

diff --git a/compiler/GHC/Core/Opt/CprAnal.hs b/compiler/GHC/Core/Opt/CprAnal.hs
index 0a35583acf..a697dd65d0 100644
--- a/compiler/GHC/Core/Opt/CprAnal.hs
+++ b/compiler/GHC/Core/Opt/CprAnal.hs
@@ -16,6 +16,7 @@ import GHC.Types.Cpr
 import GHC.Core
 import GHC.Core.Seq
 import GHC.Utils.Outputable
+import GHC.Builtin.Names ( runRWKey )
 import GHC.Types.Var.Env
 import GHC.Types.Basic
 import GHC.Types.Id
@@ -148,8 +149,6 @@ cprAnal' _ (Lit lit)     = (topCprType, Lit lit)
 cprAnal' _ (Type ty)     = (topCprType, Type ty)      -- Doesn't happen, in fact
 cprAnal' _ (Coercion co) = (topCprType, Coercion co)
 
-cprAnal' env (Var var)   = (cprTransform env var, Var var)
-
 cprAnal' env (Cast e co)
   = (cpr_ty, Cast e' co)
   where
@@ -160,19 +159,10 @@ cprAnal' env (Tick t e)
   where
     (cpr_ty, e') = cprAnal env e
 
-cprAnal' env (App fun (Type ty))
-  = (fun_ty, App fun' (Type ty))
-  where
-    (fun_ty, fun') = cprAnal env fun
-
-cprAnal' env (App fun arg)
-  = (res_ty, App fun' arg')
-  where
-    (fun_ty, fun') = cprAnal env fun
-    -- In contrast to DmdAnal, there is no useful (non-nested) CPR info to be
-    -- had by looking into the CprType of arg.
-    (_, arg')      = cprAnal env arg
-    res_ty         = applyCprTy fun_ty
+cprAnal' env e@(Var{})
+  = cprAnalApp env e [] []
+cprAnal' env e@(App{})
+  = cprAnalApp env e [] []
 
 cprAnal' env (Lam var body)
   | isTyVar var
@@ -234,26 +224,56 @@ cprAnalAlt env scrut_ty (Alt con bndrs rhs)
 -- * CPR transformer
 --
 
-cprTransform :: AnalEnv         -- ^ The analysis environment
-             -> Id              -- ^ The function
-             -> CprType         -- ^ The demand type of the function
-cprTransform env id
-  = -- pprTrace "cprTransform" (vcat [ppr id, ppr sig])
+cprAnalApp :: AnalEnv -> CoreExpr -> [CoreArg] -> [CprType] -> (CprType, CoreExpr)
+cprAnalApp env e args' arg_tys
+  -- Collect CprTypes for (value) args (inlined collectArgs):
+  | App fn arg <- e, isTypeArg arg -- Don't analyse Type args
+  = cprAnalApp env fn (arg:args') arg_tys
+  | App fn arg <- e
+  , (arg_ty, arg') <- cprAnal env arg
+  = cprAnalApp env fn (arg':args') (arg_ty:arg_tys)
+
+  | Var fn <- e
+  = (cprTransform env fn arg_tys, mkApps e args')
+
+  | otherwise -- e is not an App and not a Var
+  , (e_ty, e') <- cprAnal env e
+  = (applyCprTy e_ty (length arg_tys), mkApps e' args')
+
+cprTransform :: AnalEnv   -- ^ The analysis environment
+             -> Id        -- ^ The function
+             -> [CprType] -- ^ info about incoming /value/ arguments
+             -> CprType   -- ^ The demand type of the application
+cprTransform env id args
+  = -- pprTrace "cprTransform" (vcat [ppr id, ppr args, ppr sig])
     sig
   where
     sig
-      -- Top-level binding, local let-binding or case binder
+      -- Top-level binding, local let-binding, lambda arg or case binder
       | Just sig <- lookupSigEnv env id
-      = getCprSig sig
+      = applyCprTy (getCprSig sig) (length args)
+      -- CPR transformers for special Ids
+      | Just cpr_ty <- cprTransformSpecial id args
+      = cpr_ty
       -- See Note [CPR for data structures]
       | Just rhs <- cprDataStructureUnfolding_maybe id
       = fst $ cprAnal env rhs
       -- Imported function or data con worker
       | isGlobalId id
-      = getCprSig (idCprSig id)
+      = applyCprTy (getCprSig (idCprSig id)) (length args)
       | otherwise
       = topCprType
 
+-- | CPR transformers for special Ids
+cprTransformSpecial :: Id -> [CprType] -> Maybe CprType
+cprTransformSpecial id args
+  -- See Note [Simplification of runRW#] in GHC.CoreToStg.Prep
+  | idUnique id == runRWKey -- `runRW (\s -> e)`
+  , [arg] <- args           -- `\s -> e` has CPR type `arg` (e.g. `. -> 2`)
+  = Just $ applyCprTy arg 1 -- `e` has CPR type `2`
+  | otherwise
+  = Nothing
+
 --
 -- * Bindings
 --
diff --git a/compiler/GHC/CoreToStg/Prep.hs b/compiler/GHC/CoreToStg/Prep.hs
index fa20e39e70..4fff314839 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -1310,6 +1310,14 @@ in straight-line code. Consequently, GHC.Core.Opt.SetLevels.lvlApp has special
 treatment for runRW# applications, ensure the arguments are not floated as
 MFEs.
 
+Now that we float evaluation context into runRW#, we also have to give runRW# a
+special higher-order CPR transformer lest we risk #19822. E.g.,
+
+  case runRW# (\s -> doThings) of x -> Data.Text.Text x something something'
+      ~>
+  runRW# (\s -> case doThings s of x -> Data.Text.Text x something something')
+
+The former had the CPR property, and so should the latter.
 
 Other considered designs
 ------------------------
diff --git a/compiler/GHC/Types/Cpr.hs b/compiler/GHC/Types/Cpr.hs
index 0e12765314..c07b614e58 100644
--- a/compiler/GHC/Types/Cpr.hs
+++ b/compiler/GHC/Types/Cpr.hs
@@ -124,9 +124,9 @@ lubCprType ty1@(CprType n1 cpr1) ty2@(CprType n2 cpr2)
   | n1 == n2                   = CprType n1 (lubCpr cpr1 cpr2)
   | otherwise                  = topCprType
 
-applyCprTy :: CprType -> CprType
-applyCprTy (CprType n res)
-  | n > 0         = CprType (n-1) res
+applyCprTy :: CprType -> Arity -> CprType
+applyCprTy (CprType n res) k
+  | n >= k        = CprType (n-k) res
   | res == botCpr = botCprType
   | otherwise     = topCprType
 
diff --git a/testsuite/tests/cpranal/sigs/T19822.hs b/testsuite/tests/cpranal/sigs/T19822.hs
new file mode 100644
index 0000000000..ca83bafcc3
--- /dev/null
+++ b/testsuite/tests/cpranal/sigs/T19822.hs
@@ -0,0 +1,14 @@
+{-# OPTIONS_GHC -O2 -fforce-recomp #-}
+{-# LANGUAGE MagicHash, UnboxedTuples #-}
+
+module T19822 where
+
+import GHC.Exts
+import Data.Char
+
+data Text = MkText !Int Char
+
+-- | Should have the CPR property
+singleton :: Char -> Text
+singleton c = MkText (runRW# (\_ -> 42 + ord c)) c
+{-# NOINLINE singleton #-} -- to force WW
diff --git a/testsuite/tests/cpranal/sigs/T19822.stderr b/testsuite/tests/cpranal/sigs/T19822.stderr
new file mode 100644
index 0000000000..8e4636d322
--- /dev/null
+++ b/testsuite/tests/cpranal/sigs/T19822.stderr
@@ -0,0 +1,5 @@
+
+==================== Cpr signatures ====================
+T19822.singleton: 1
+
+
diff --git a/testsuite/tests/cpranal/sigs/all.T b/testsuite/tests/cpranal/sigs/all.T
index 0647c8a611..99cdebe716 100644
--- a/testsuite/tests/cpranal/sigs/all.T
+++ b/testsuite/tests/cpranal/sigs/all.T
@@ -8,3 +8,4 @@ setTestOpts(extra_hc_opts('-dno-typeable-binds -ddump-cpr-signatures'))
 test('CaseBinderCPR', normal, compile, [''])
 test('T19232', normal, compile, [''])
 test('T19398', normal, compile, [''])
+test('T19822', normal, compile, [''])
diff --git a/testsuite/tests/stranal/sigs/T18086.stderr b/testsuite/tests/stranal/sigs/T18086.stderr
index 43266ad0da..1748a0c145 100644
--- a/testsuite/tests/stranal/sigs/T18086.stderr
+++ b/testsuite/tests/stranal/sigs/T18086.stderr
@@ -9,7 +9,7 @@ T18086.panic: <L>x
 ==================== Cpr signatures ====================
 T18086.$trModule:
 T18086.m: b
-T18086.panic:
+T18086.panic: b