rts: Fix `prompt#` when profiling is enabled

This commit also adds a new -Dk RTS option to the debug RTS to assist
debugging continuation captures. Currently, the printed information is
quite minimal, but more can be added in the future if it proves to be
useful when debugging future issues.

fixes #23001

(cherry picked from commit e5794ede9e2af208669438a7f72958aeecbec111)

7 files changed, 33 insertions, 2 deletions

diff --git a/docs/users_guide/runtime_control.rst b/docs/users_guide/runtime_control.rst
index a51fb7f2a2..64d38367a7 100644
--- a/docs/users_guide/runtime_control.rst
+++ b/docs/users_guide/runtime_control.rst
@@ -1367,6 +1367,7 @@ recommended for everyday use!
 .. rts-flag::  -Dc  DEBUG: program coverage
 .. rts-flag::  -Dr  DEBUG: sparks
 .. rts-flag::  -DC  DEBUG: compact
+.. rts-flag::  -Dk  DEBUG: continuation
 
     Debug messages will be sent to the binary event log file instead of
     stdout if the :rts-flag:`-l ⟨flags⟩` option is added. This might be useful
diff --git a/rts/Continuation.c b/rts/Continuation.c
index fbc279574f..f2ec62fe60 100644
--- a/rts/Continuation.c
+++ b/rts/Continuation.c
@@ -12,6 +12,7 @@
 #include "sm/Storage.h"
 #include "sm/Sanity.h"
 #include "Continuation.h"
+#include "Printer.h"
 #include "Threads.h"
 
 #include <string.h>
@@ -392,7 +393,14 @@ StgClosure *captureContinuationAndAbort(Capability *cap, StgTSO *tso, StgPromptT
 
   /* --- Phase 1: Find the matching prompt frame ---------------------------- */
 
+  IF_DEBUG(continuation,
+    debugBelch("captureContinuationAndAbort: searching for prompt\n");
+    debugBelch("  prompt_tag = "); printClosure(prompt_tag));
+
   while (true) {
+    IF_DEBUG(continuation,
+      printStackChunk(frame, frame + stack_frame_sizeW((StgClosure *)frame)));
+
     const StgInfoTable *info_ptr = ((StgClosure *)frame)->header.info;
     const StgRetInfoTable *info = get_ret_itbl((StgClosure *)frame);
     StgWord chunk_words = frame - stack->sp;
@@ -429,6 +437,8 @@ StgClosure *captureContinuationAndAbort(Capability *cap, StgTSO *tso, StgPromptT
                   || info->i.type == ATOMICALLY_FRAME
                   || info->i.type == CATCH_RETRY_FRAME
                   || info->i.type == CATCH_STM_FRAME)) {
+      IF_DEBUG(continuation,
+        debugBelch("captureContinuationAndAbort: could not find prompt, bailing out\n"));
       return NULL; // Bail out
     }
 
@@ -452,6 +462,10 @@ StgClosure *captureContinuationAndAbort(Capability *cap, StgTSO *tso, StgPromptT
 
   /* --- Phase 2: Perform the capture --------------------------------------- */
 
+  IF_DEBUG(continuation,
+    debugBelch("captureContinuationAndAbort: found prompt, "
+               "capturing %" FMT_Word " words of stack\n", total_words));
+
   dirty_TSO(cap, tso);
   dirty_STACK(cap, stack);
 
diff --git a/rts/ContinuationOps.cmm b/rts/ContinuationOps.cmm
index efb485057f..8e36148cd6 100644
--- a/rts/ContinuationOps.cmm
+++ b/rts/ContinuationOps.cmm
@@ -49,7 +49,12 @@ stg_newPromptTagzh()
   return (tag);
 }
 
-INFO_TABLE_RET(stg_prompt_frame, RET_SMALL, W_ info_ptr, P_ tag /* :: PromptTag# a */)
+#define PROMPT_FRAME_FIELDS(w_,p_,info_ptr,p1,p2,tag) \
+  w_ info_ptr,                                        \
+  PROF_HDR_FIELDS(w_,p1,p2)                           \
+  p_ tag
+
+INFO_TABLE_RET(stg_prompt_frame, RET_SMALL, PROMPT_FRAME_FIELDS(W_,P_, info_ptr, p1, p2, tag /* :: PromptTag# a */))
   return (P_ ret /* :: a */)
 {
   return (ret);
@@ -61,7 +66,9 @@ stg_promptzh(P_ tag /* :: PromptTag# a */, P_ io /* :: IO a */)
   STK_CHK_GEN();
   TICK_UNKNOWN_CALL();
   TICK_SLOW_CALL_fast_v();
-  jump stg_ap_v_fast (stg_prompt_frame_info, tag) (io);
+  jump stg_ap_v_fast
+    (PROMPT_FRAME_FIELDS(,,stg_prompt_frame_info, CCCS, 0, tag))
+    (io);
 }
 
 /* --------------------------------------------------------------------------
diff --git a/rts/RtsFlags.c b/rts/RtsFlags.c
index a1823b90b5..5f1e8d6403 100644
--- a/rts/RtsFlags.c
+++ b/rts/RtsFlags.c
@@ -205,6 +205,7 @@ void initRtsFlagsDefaults(void)
     RtsFlags.DebugFlags.sparks          = false;
     RtsFlags.DebugFlags.numa            = false;
     RtsFlags.DebugFlags.compact         = false;
+    RtsFlags.DebugFlags.continuation    = false;
 
 #if defined(PROFILING)
     RtsFlags.CcFlags.doCostCentres      = COST_CENTRES_NONE;
@@ -476,6 +477,7 @@ usage_text[] = {
 "  -Dc  DEBUG: program coverage",
 "  -Dr  DEBUG: sparks",
 "  -DC  DEBUG: compact",
+"  -Dk  DEBUG: continuation",
 "",
 "     NOTE: DEBUG events are sent to stderr by default; add -l to create a",
 "     binary event log file instead.",
@@ -2190,6 +2192,9 @@ static void read_debug_flags(const char* arg)
         case 'C':
             RtsFlags.DebugFlags.compact = true;
             break;
+        case 'k':
+            RtsFlags.DebugFlags.continuation = true;
+            break;
         default:
             bad_option( arg );
         }
diff --git a/rts/Trace.h b/rts/Trace.h
index 6e08d0d2fd..9c9a557296 100644
--- a/rts/Trace.h
+++ b/rts/Trace.h
@@ -65,6 +65,7 @@ enum CapsetType { CapsetTypeCustom = CAPSET_TYPE_CUSTOM,
 #define DEBUG_hpc         RtsFlags.DebugFlags.hpc
 #define DEBUG_sparks      RtsFlags.DebugFlags.sparks
 #define DEBUG_compact     RtsFlags.DebugFlags.compact
+#define DEBUG_continuation RtsFlags.DebugFlags.continuation
 
 // Event-enabled flags
 // These semantically booleans but we use a dense packing to minimize their
diff --git a/rts/include/rts/Flags.h b/rts/include/rts/Flags.h
index e33d97b17c..57fe0eb797 100644
--- a/rts/include/rts/Flags.h
+++ b/rts/include/rts/Flags.h
@@ -113,6 +113,7 @@ typedef struct _DEBUG_FLAGS {
     bool sparks;         /* 'r' */
     bool numa;           /* '--debug-numa' */
     bool compact;        /* 'C' */
+    bool continuation;   /* 'k' */
 } DEBUG_FLAGS;
 
 /* See Note [Synchronization of flags and base APIs] */
diff --git a/testsuite/tests/rts/continuations/all.T b/testsuite/tests/rts/continuations/all.T
index e8852416eb..efa2589ee2 100644
--- a/testsuite/tests/rts/continuations/all.T
+++ b/testsuite/tests/rts/continuations/all.T
@@ -1,4 +1,6 @@
 setTestOpts(js_broken(22261))
+if have_profiling():
+  setTestOpts(extra_ways(['prof']))
 
 test('cont_simple_shift', [extra_files(['ContIO.hs'])], multimod_compile_and_run, ['cont_simple_shift', ''])
 test('cont_exn_masking', [extra_files(['ContIO.hs'])], multimod_compile_and_run, ['cont_exn_masking', ''])