Make keepAlive# out-of-linewip/ghc-9.2-T21708

This is a naive approach to fixing the unsoundness noticed in #21708.
Specifically, we remove the lowering of `keepAlive#` via CorePrep and
instead turn it into an out-of-line primop.
This is simple, inefficient (since the continuation must now be heap
allocated), but good enough for 9.4.1. We will revisit this
(particiularly via #16098) in a future release.

Metric Increase:
    T4978
    T7257
    T9203

(cherry picked from commit d75c540d439510491b45f64c1113762dcb251ae1)

7 files changed, 31 insertions, 36 deletions

diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index 0810d39b16..611869c1e0 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -3257,6 +3257,7 @@ primop KeepAliveOp "keepAlive#" GenPrimOp
    { \tt{keepAlive# x s k} keeps the value \tt{x} alive during the execution
      of the computation \tt{k}. }
    with
+   out_of_line = True
    strictness = { \ _arity -> mkClosedStrictSig [topDmd, topDmd, strictOnceApply1Dmd] topDiv }
 
 
diff --git a/compiler/GHC/CoreToStg/Prep.hs b/compiler/GHC/CoreToStg/Prep.hs
index bd99bf976d..3f6a93d27d 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -32,10 +32,8 @@ import GHC.Tc.Utils.Env
 import GHC.Unit
 
 import GHC.Builtin.Names
-import GHC.Builtin.PrimOps
 import GHC.Builtin.Types
-import GHC.Builtin.Types.Prim ( realWorldStatePrimTy )
-import GHC.Types.Id.Make ( realWorldPrimId, mkPrimOpId )
+import GHC.Types.Id.Make ( realWorldPrimId )
 
 import GHC.Core.Utils
 import GHC.Core.Opt.Arity
@@ -1012,36 +1010,6 @@ cpeApp top_env expr
         = let (terminal, args', depth') = collect_args arg
           in cpe_app env terminal (args' ++ args) (depth + depth' - 1)
 
-    -- See Note [keepAlive# magic].
-    cpe_app env
-            (Var f)
-            args
-            n
-        | Just KeepAliveOp <- isPrimOpId_maybe f
-        , CpeApp (Type arg_rep)
-          : CpeApp (Type arg_ty)
-          : CpeApp (Type _result_rep)
-          : CpeApp (Type result_ty)
-          : CpeApp arg
-          : CpeApp s0
-          : CpeApp k
-          : rest <- args
-        = do { y  <- newVar (cpSubstTy env result_ty)
-             ; s2 <- newVar realWorldStatePrimTy
-             ; -- beta reduce if possible
-             ; (floats, k') <- case k of
-                  Lam s body -> cpe_app (extendCorePrepEnvExpr env s s0) body rest (n-2)
-                  _          -> cpe_app env k (CpeApp s0 : rest) (n-1)
-             ; let touchId = mkPrimOpId TouchOp
-                   expr = Case k' y result_ty [Alt DEFAULT [] rhs]
-                   rhs = let scrut = mkApps (Var touchId) [Type arg_rep, Type arg_ty, arg, Var realWorldPrimId]
-                         in Case scrut s2 result_ty [Alt DEFAULT [] (Var y)]
-             ; (floats', expr') <- cpeBody env expr
-             ; return (floats `appendFloats` floats', expr')
-             }
-        | Just KeepAliveOp <- isPrimOpId_maybe f
-        = panic "invalid keepAlive# application"
-
     cpe_app env (Var f) (CpeApp _runtimeRep@Type{} : CpeApp _type@Type{} : CpeApp arg : rest) n
         | f `hasKey` runRWKey
         -- N.B. While it may appear that n == 1 in the case of runRW#
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index 41fbaf28ba..619fdd37a5 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -1673,9 +1673,7 @@ emitPrimOp dflags primop = case primop of
   TraceEventBinaryOp -> alwaysExternal
   TraceMarkerOp -> alwaysExternal
   SetThreadAllocationCounter -> alwaysExternal
-
-  -- See Note [keepAlive# magic] in GHC.CoreToStg.Prep.
-  KeepAliveOp -> panic "keepAlive# should have been eliminated in CorePrep"
+  KeepAliveOp -> alwaysExternal
 
  where
   profile = targetProfile dflags
diff --git a/includes/rts/storage/Closures.h b/includes/rts/storage/Closures.h
index ebb836bca2..8dbc739541 100644
--- a/includes/rts/storage/Closures.h
+++ b/includes/rts/storage/Closures.h
@@ -186,6 +186,11 @@ typedef struct _StgUpdateFrame {
 
 typedef struct {
     StgHeader  header;
+    StgClosure *c;
+} StgKeepAliveFrame;
+
+typedef struct {
+    StgHeader  header;
     StgWord    exceptions_blocked;
     StgClosure *handler;
 } StgCatchFrame;
diff --git a/includes/stg/MiscClosures.h b/includes/stg/MiscClosures.h
index d8aefd8035..a988d55b05 100644
--- a/includes/stg/MiscClosures.h
+++ b/includes/stg/MiscClosures.h
@@ -61,6 +61,7 @@ RTS_RET(stg_unmaskAsyncExceptionszh_ret);
 RTS_RET(stg_maskUninterruptiblezh_ret);
 RTS_RET(stg_maskAsyncExceptionszh_ret);
 RTS_RET(stg_stack_underflow_frame);
+RTS_RET(stg_keepAlive_frame);
 RTS_RET(stg_restore_cccs);
 RTS_RET(stg_restore_cccs_eval);
 
@@ -495,6 +496,7 @@ RTS_FUN_DECL(stg_raiseUnderflowzh);
 RTS_FUN_DECL(stg_raiseOverflowzh);
 RTS_FUN_DECL(stg_raiseIOzh);
 RTS_FUN_DECL(stg_paniczh);
+RTS_FUN_DECL(stg_keepAlivezh);
 RTS_FUN_DECL(stg_absentErrorzh);
 
 RTS_FUN_DECL(stg_makeStableNamezh);
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index e950f0de9a..26568f961d 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -2896,3 +2896,23 @@ stg_setThreadAllocationCounterzh ( I64 counter )
     StgTSO_alloc_limit(CurrentTSO) = counter + TO_I64(offset);
     return ();
 }
+
+
+#define KEEP_ALIVE_FRAME_FIELDS(w_,p_,info_ptr,p1,p2,c)   \
+  w_ info_ptr,                                            \
+  PROF_HDR_FIELDS(w_,p1,p2)                               \
+  p_ c
+
+stg_keepAlivezh ( P_ c,      /* :: v */
+                  P_ io      /* :: IO p */ )
+{
+    STK_CHK_GEN();
+    jump stg_ap_v_fast
+        (KEEP_ALIVE_FRAME_FIELDS(,,stg_keepAlive_frame_info, CCCS, 0, c))(io);
+}
+
+INFO_TABLE_RET(stg_keepAlive_frame, RET_SMALL, KEEP_ALIVE_FRAME_FIELDS(W_,P_, info_ptr, p1, p2, c))
+    return (P_ ret)
+{
+    return (ret);
+}
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 987755651d..eabdb16d27 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -780,6 +780,7 @@ extern char **environ;
       SymI_HasProto(stg_raiseUnderflowzh)                               \
       SymI_HasProto(stg_raiseOverflowzh)                                \
       SymI_HasProto(stg_raiseIOzh)                                      \
+      SymI_HasProto(stg_keepAlivezh)                                    \
       SymI_HasProto(stg_paniczh)                                        \
       SymI_HasProto(stg_absentErrorzh)                                  \
       SymI_HasProto(stg_readTVarzh)                                     \