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

5 files changed, 30 insertions, 2 deletions

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 c66da4ae4e..3cec74cfa9 100644
--- a/rts/Trace.h
+++ b/rts/Trace.h
@@ -67,6 +67,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] */