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

(cherry picked from commit b0df2f012ede2f9017577198b5285ca6015b582b)

7 files changed, 81 insertions, 26 deletions

diff --git a/compiler/GHC/Core/Opt/CprAnal.hs b/compiler/GHC/Core/Opt/CprAnal.hs
index 4fcda8c4a8..61617ef3e9 100644
--- a/compiler/GHC/Core/Opt/CprAnal.hs
+++ b/compiler/GHC/Core/Opt/CprAnal.hs
@@ -19,6 +19,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.Core.DataCon
@@ -145,8 +146,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
@@ -157,19 +156,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
@@ -225,26 +215,56 @@ cprAnalAlt env _ _ (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
       -- See Note [CPR for expandable unfoldings]
       | Just rhs <- cprExpandUnfolding_maybe id
       = fst $ cprAnal env rhs
+      -- CPR transformers for special Ids
+      | Just cpr_ty <- cprTransformSpecial id args
+      = cpr_ty
       -- Imported function or data con worker
       | isGlobalId id
-      = getCprSig (idCprInfo id)
+      = applyCprTy (getCprSig (idCprInfo id)) (length args)
       -- Local let-bound
       | Just sig <- lookupSigEnv env id
-      = getCprSig sig
+      = applyCprTy (getCprSig sig) (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 1322984243..40f0317d72 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -1138,6 +1138,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 a884091cef..1907a6da3a 100644
--- a/compiler/GHC/Types/Cpr.hs
+++ b/compiler/GHC/Types/Cpr.hs
@@ -93,9 +93,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..5f1522fd00
--- /dev/null
+++ b/testsuite/tests/cpranal/sigs/T19822.stderr
@@ -0,0 +1,5 @@
+
+==================== Cpr signatures ====================
+T19822.singleton: m1
+
+
diff --git a/testsuite/tests/cpranal/sigs/all.T b/testsuite/tests/cpranal/sigs/all.T
new file mode 100644
index 0000000000..749a5f739a
--- /dev/null
+++ b/testsuite/tests/cpranal/sigs/all.T
@@ -0,0 +1,8 @@
+# We are testing the result of an optimization, so no use
+# running them in various runtimes
+setTestOpts(only_ways(['optasm']))
+# This directory contains tests where we annotate functions with expected
+# CPR signatures, and verify that these are actually those found by the compiler
+setTestOpts(extra_hc_opts('-dno-typeable-binds -ddump-cpr-signatures'))
+
+test('T19822', normal, compile, [''])
diff --git a/testsuite/tests/stranal/sigs/T18086.stderr b/testsuite/tests/stranal/sigs/T18086.stderr
index 6941e233f8..8488b24557 100644
--- a/testsuite/tests/stranal/sigs/T18086.stderr
+++ b/testsuite/tests/stranal/sigs/T18086.stderr
@@ -9,7 +9,7 @@ T18086.panic: <L,U>x
 ==================== Cpr signatures ====================
 T18086.$trModule:
 T18086.m: b
-T18086.panic:
+T18086.panic: b