WorkWrap: Nuke CPR signatures of join points (#18824)

In #18824 we saw that the Simplifier didn't nuke a CPR signature of a join point
when it pushed a continuation into it when it better should have.

But join points are local, mostly non-exported bindings. We don't use their
CPR signature anyway and would discard it at the end of the Core pipeline.
Their main purpose is to propagate CPR info during CPR analysis and by the time
worker/wrapper runs the signature will have served its purpose. So we zap it!

Fixes #18824.

4 files changed, 79 insertions, 24 deletions

diff --git a/compiler/GHC/Core/Opt/WorkWrap.hs b/compiler/GHC/Core/Opt/WorkWrap.hs
index 8d0dc2367f..2c204795c9 100644
--- a/compiler/GHC/Core/Opt/WorkWrap.hs
+++ b/compiler/GHC/Core/Opt/WorkWrap.hs
@@ -390,9 +390,21 @@ tuple.
           in case z of A -> j 1 2
                        B -> j 2 3
 
-Note that we still want to give @j@ the CPR property, so that @f@ has it. So
+Note that we still want to give `j` the CPR property, so that `f` has it. So
 CPR *analyse* join points as regular functions, but don't *transform* them.
 
+We could retain the CPR /signature/ on the worker after W/W, but it would
+become outright wrong if the Simplifier pushes a non-trivial continuation
+into it. For example:
+    case (let $j x = (x,x) in ...) of alts
+    ==>
+    let $j x = case (x,x) of alts in case ... of alts
+Before pushing the case in, `$j` has the CPR property, but not afterwards.
+
+So we simply zap the CPR signature for join pints as part of the W/W pass.
+The signature served its purpose during CPR analysis in propagating the
+CPR property of `$j`.
+
 Doing W/W for returned products on a join point would be tricky anyway, as the
 worker could not be a join point because it would not be tail-called. However,
 doing the *argument* part of W/W still works for join points, since the wrapper
@@ -529,7 +541,7 @@ tryWW ww_opts is_rec fn_id rhs
   = return [ (new_fn_id, rhs ) ]
 
   | is_fun && is_eta_exp
-  = splitFun ww_opts new_fn_id fn_info rhs
+  = splitFun ww_opts new_fn_id rhs
 
   -- See Note [Thunk splitting]
   | isNonRec is_rec, is_thunk
@@ -541,11 +553,17 @@ tryWW ww_opts is_rec fn_id rhs
   where
     fn_info        = idInfo fn_id
     (wrap_dmds, _) = splitDmdSig (dmdSigInfo fn_info)
+    new_fn_id      = zap_join_cpr $ zap_usage fn_id
 
-    new_fn_id = zapIdUsedOnceInfo (zapIdUsageEnvInfo fn_id)
+    zap_usage = zapIdUsedOnceInfo . zapIdUsageEnvInfo
         -- See Note [Zapping DmdEnv after Demand Analyzer] and
         -- See Note [Zapping Used Once info in WorkWrap]
 
+    zap_join_cpr id
+      | isJoinId id = id `setIdCprSig` topCprSig
+      | otherwise   = id
+        -- See Note [Don't w/w join points for CPR]
+
     -- is_eta_exp: see Note [Don't eta expand in w/w]
     is_eta_exp = length wrap_dmds == manifestArity rhs
     is_fun     = notNull wrap_dmds || isJoinId fn_id
@@ -690,11 +708,11 @@ by LitRubbish (see Note [Drop absent bindings]) so there is no great harm.
 
 
 ---------------------
-splitFun :: WwOpts -> Id -> IdInfo -> CoreExpr -> UniqSM [(Id, CoreExpr)]
-splitFun ww_opts fn_id fn_info rhs
+splitFun :: WwOpts -> Id -> CoreExpr -> UniqSM [(Id, CoreExpr)]
+splitFun ww_opts fn_id rhs
   = warnPprTrace (not (wrap_dmds `lengthIs` (arityInfo fn_info)))
                  (ppr fn_id <+> (ppr wrap_dmds $$ ppr cpr)) $
-    do { mb_stuff <- mkWwBodies ww_opts fn_id arg_vars (exprType body) wrap_dmds use_cpr_info
+    do { mb_stuff <- mkWwBodies ww_opts fn_id arg_vars (exprType body) wrap_dmds cpr
        ; case mb_stuff of
             Nothing -> -- No useful wrapper; leave the binding alone
                        return [(fn_id, rhs)]
@@ -710,9 +728,10 @@ splitFun ww_opts fn_id fn_info rhs
               | otherwise
               -> do { work_uniq <- getUniqueM
                     ; return (mkWWBindPair ww_opts fn_id fn_info arg_vars body
-                                           work_uniq div cpr stuff) } }
+                                           work_uniq div stuff) } }
   where
     uf_opts = so_uf_opts (wo_simple_opts ww_opts)
+    fn_info = idInfo fn_id
     (arg_vars, body) = collectBinders rhs
             -- collectBinders was not enough for GHC.Event.IntTable.insertWith
             -- last time I checked, where manifest lambdas were wrapped in casts
@@ -727,17 +746,11 @@ splitFun ww_opts fn_id fn_info rhs
                       <+> text "arityInfo:" <+> ppr (arityInfo fn_info)) $
           ct_cpr cpr_ty
 
-    -- use_cpr_info is the CPR we w/w for. Note that we kill it for join points,
-    -- see Note [Don't w/w join points for CPR].
-    use_cpr_info  | isJoinId fn_id = topCpr
-                  | otherwise      = cpr
-
-
 mkWWBindPair :: WwOpts -> Id -> IdInfo
-             -> [Var] -> CoreExpr -> Unique -> Divergence -> Cpr
+             -> [Var] -> CoreExpr -> Unique -> Divergence
              -> ([Demand], JoinArity, Id -> CoreExpr, Expr CoreBndr -> CoreExpr)
              -> [(Id, CoreExpr)]
-mkWWBindPair ww_opts fn_id fn_info fn_args fn_body work_uniq div cpr
+mkWWBindPair ww_opts fn_id fn_info fn_args fn_body work_uniq div
              (work_demands, join_arity, wrap_fn, work_fn)
   = [(work_id, work_rhs), (wrap_id, wrap_rhs)]
      -- Worker first, because wrapper mentions it
@@ -783,7 +796,7 @@ mkWWBindPair ww_opts fn_id fn_info fn_args fn_body work_uniq div cpr
                         -- Even though we may not be at top level,
                         -- it's ok to give it an empty DmdEnv
 
-                `setIdCprSig` mkCprSig work_arity work_cpr_info
+                `setIdCprSig` topCprSig
 
                 `setIdDemandInfo` worker_demand
 
@@ -814,13 +827,6 @@ mkWWBindPair ww_opts fn_id fn_info fn_args fn_body work_uniq div cpr
     fn_inline_spec  = inl_inline fn_inl_prag
     fn_unfolding    = realUnfoldingInfo fn_info
 
-    -- Even if we don't w/w join points for CPR, we might still do so for
-    -- strictness. In which case a join point worker keeps its original CPR
-    -- property; see Note [Don't w/w join points for CPR]. Otherwise, the worker
-    -- doesn't have the CPR property anymore.
-    work_cpr_info | isJoinId fn_id = cpr
-                  | otherwise      = topCpr
-
 mkStrWrapperInlinePrag :: InlinePragma -> InlinePragma
 mkStrWrapperInlinePrag (InlinePragma { inl_act = act, inl_rule = rule_info })
   = InlinePragma { inl_src    = SourceText "{-# INLINE"
diff --git a/testsuite/tests/cpranal/should_compile/T18824.hs b/testsuite/tests/cpranal/should_compile/T18824.hs
new file mode 100644
index 0000000000..79f6c49aaa
--- /dev/null
+++ b/testsuite/tests/cpranal/should_compile/T18824.hs
@@ -0,0 +1,39 @@
+{-# OPTIONS_GHC -O2 -fforce-recomp #-}
+{-# LANGUAGE DeriveFunctor #-}
+
+module T18824 where
+
+import Control.Applicative
+
+data Box a = Box a
+  deriving Functor
+
+instance Applicative Box where
+  pure = Box
+  Box f <*> Box a = Box (f a)
+
+data X = X
+  (Maybe String)
+  (Maybe String)
+  (Maybe String)
+  (Maybe String)
+  (Maybe String)
+
+mb :: (String -> Box a) -> String -> Box (Maybe a)
+mb _ ""  = Box Nothing
+mb _ "-" = Box Nothing
+mb p xs  = Just <$> p xs
+
+run :: [String] -> Box X
+run
+  [ x1
+  , x2
+  , x3
+  , x4
+  , x5
+  ] = X
+  <$> mb pure x1
+  <*> mb pure x2
+  <*> mb pure x3
+  <*> mb pure x4
+  <*> mb pure x5
diff --git a/testsuite/tests/cpranal/should_compile/all.T b/testsuite/tests/cpranal/should_compile/all.T
index 570b78228f..03d2b10d68 100644
--- a/testsuite/tests/cpranal/should_compile/all.T
+++ b/testsuite/tests/cpranal/should_compile/all.T
@@ -8,7 +8,17 @@ test('Cpr001', [], multimod_compile, ['Cpr001', '-v0'])
 # The following tests grep for type signatures of worker functions.
 test('T18109', [ grep_errmsg(r'^T18109\.\$w\S+ ::') ], compile, ['-ddump-simpl -dsuppress-idinfo -dppr-cols=9999'])
 test('T18174', [ grep_errmsg(r'^T18174\.\$w\S+ ::') ], compile, ['-ddump-simpl -dsuppress-idinfo -dppr-cols=9999'])
+
 # T18401 probably needs -flate-dmd-anal so that it runs after SpecConstr.
 # It is currently broken, but not marked expect_broken. We can't know the exact
 # name of the function before it is fixed, so expect_broken doesn't make sense.
 test('T18401', [ grep_errmsg(r'^T18401\.\S+ ::') ], compile, ['-ddump-simpl -dsuppress-idinfo -dppr-cols=9999 -flate-dmd-anal'])
+
+# Exitify is (one of) the first passes after worker/wrapper.
+# The CPR sig on the join point that sucked in a `case of` should have been nuked here.
+# The Regex works by grepping for the IdDetails line of the join point, which used to look like
+#     [LclId[JoinId(1)], Arity=1, Str=<L>, Cpr=1, Unf=...]
+# It won't match if the Cpr=1 is missing, which is what we're trying to assess.
+test('T18824', [ grep_errmsg(r'JoinId[^\n]*Cpr') ], compile, ['-ddump-exitify -dppr-cols=1000 -dsuppress-uniques'])
+
+
diff --git a/testsuite/tests/simplCore/should_compile/T3772.stdout b/testsuite/tests/simplCore/should_compile/T3772.stdout
index a43c4fc349..96830336bc 100644
--- a/testsuite/tests/simplCore/should_compile/T3772.stdout
+++ b/testsuite/tests/simplCore/should_compile/T3772.stdout
@@ -42,7 +42,7 @@ T3772.$trModule
 Rec {
 -- RHS size: {terms: 10, types: 2, coercions: 0, joins: 0/0}
 $wxs :: GHC.Prim.Int# -> (# #)
-[GblId, Arity=1, Str=<1L>, Cpr=1, Unf=OtherCon []]
+[GblId, Arity=1, Str=<1L>, Unf=OtherCon []]
 $wxs
   = \ (ww :: GHC.Prim.Int#) ->
       case ww of ds1 {