From 1b2db3048b87ffcc5e7f8998bd17afefc3cb3bef Mon Sep 17 00:00:00 2001
From: MaartenS11 <maarten.steevens@gmail.com>
Date: Tue, 17 Feb 2026 14:51:23 +0100
Subject: [PATCH] Add tracing checkpoint policy

This option is configurable, this way you can choose when it sends data, for example: only after input primitives. Additionally, the data it sends can also be configured, for example: only send the program counter.
---
 src/Debug/debugger.cpp | 56 ++++++++++++++++++++++++++++++++++++++----
 src/Debug/debugger.h   |  3 +++
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/src/Debug/debugger.cpp b/src/Debug/debugger.cpp
index 2631af78..4384ef74 100644
--- a/src/Debug/debugger.cpp
+++ b/src/Debug/debugger.cpp
@@ -25,6 +25,9 @@ Debugger::Debugger(Channel *duplex) {
     this->checkpointInterval = 10;
     this->instructions_executed = 0;
     this->fidx_called = {};
+    this->min_return_values = 10;
+    this->checkpoint_state = nullptr;
+    this->checkpoint_state_size = 0;
     this->remaining_instructions = -1;
 }
 
@@ -965,11 +968,33 @@ void Debugger::inspect(Module *m, const uint16_t sizeStateArray,
 }
 
 void Debugger::setSnapshotPolicy(Module *m, uint8_t *interruptData) {
-    snapshotPolicy = SnapshotPolicy{*interruptData};
+    uint8_t **data_ptr = &interruptData;
+    if (*interruptData <= 2) {
+        snapshotPolicy = SnapshotPolicy{*interruptData};
+        min_return_values = 0;
+        if (checkpoint_state) {
+            free(checkpoint_state);
+        }
+        checkpoint_state = nullptr;
+        checkpoint_state_size = 0;
+    } else {
+        snapshotPolicy = SnapshotPolicy::checkpointing;
+        *data_ptr += 1;
+        min_return_values = read_LEB_32(data_ptr);
+        if (checkpoint_state) {
+            free(checkpoint_state);
+        }
+        checkpoint_state_size = read_LEB_32(data_ptr);
+        checkpoint_state = new uint8_t[checkpoint_state_size];
+        for (int i = 0; i < checkpoint_state_size; i++) {
+            checkpoint_state[i] = **data_ptr;
+            *data_ptr += 1;
+        }
+    }
 
     // Make a checkpoint when you first enable checkpointing
     if (snapshotPolicy == SnapshotPolicy::checkpointing) {
-        uint8_t *ptr = interruptData + 1;
+        uint8_t *ptr = *data_ptr + 1;
         checkpointInterval = read_B32(&ptr);
         checkpoint(m, true);
     }
@@ -999,7 +1024,14 @@ void Debugger::handleSnapshotPolicy(Module *m) {
         this->channel->write("\n");
     } else if (snapshotPolicy == SnapshotPolicy::checkpointing) {
         if (instructions_executed >= checkpointInterval || fidx_called) {
-            checkpoint(m);
+            if (min_return_values == 0) {
+                checkpoint(m);
+            } else if (fidx_called) {
+                const Type *type = m->functions[*fidx_called].type;
+                if (type->result_count >= min_return_values) {
+                    checkpoint(m);
+                }
+            }
         }
         instructions_executed++;
 
@@ -1032,9 +1064,23 @@ void Debugger::checkpoint(Module *m, bool force) {
             comma = true;
         }
         this->channel->write("], ");
+
+        // Return values:
+        this->channel->write(R"("returns": [)");
+        comma = false;
+        for (uint32_t i = 0; i < func_block.type->result_count; i++) {
+            channel->write("%s%d", comma ? ", " : "",
+                           m->stack[m->sp - i].value.uint32);
+            comma = true;
+        }
+        this->channel->write("], ");
+    }
+    this->channel->write(R"("snapshot": )");
+    if (!checkpoint_state) {
+        snapshot(m);
+    } else {
+        inspect(m, checkpoint_state_size, checkpoint_state);
     }
-    this->channel->write(R"("snapshot": )", instructions_executed);
-    snapshot(m);
     this->channel->write("}\n");
     instructions_executed = 0;
 }
diff --git a/src/Debug/debugger.h b/src/Debug/debugger.h
index 98f28e48..cec77b23 100644
--- a/src/Debug/debugger.h
+++ b/src/Debug/debugger.h
@@ -138,6 +138,9 @@ class Debugger {
     uint32_t instructions_executed;       // #instructions since last checkpoint
     std::optional<uint32_t> fidx_called;  // The primitive that was executed
     uint32_t prim_args[8];                // The arguments of the executed prim
+    uint32_t min_return_values;
+    uint32_t checkpoint_state_size;
+    uint8_t *checkpoint_state;
 
     // Continue for
     int32_t remaining_instructions;