[Stack Switching] Implement resume_throw in the interpreter (#7815)

kripken · web-flow · commit 5dae9f4569a8 · 2025-08-11T17:10:17.000-07:00
Pretty straightforward: Add an exception tag to ContData, and throw if
it is set.

With this, the next big spec test passes.

Throwing from the first resumption is not in the spec test, so add that
manually (for us, but I guess not the spec interpreter, this checks
another code path).

This uncovered an existing bug with StackValueNoter. We added one of
these during Resume, to ensure a stack value noting scope always
existed. But we track those by ModuleRunner, so it doesn't help when we
do a cross-module suspend. Just remove that extra logic, and handle the
lack of a scope directly, which makes this all a lot simpler.
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
@@ -226,6 +226,9 @@ struct ContData {
   // suspend).
   Literals resumeArguments;
 
+  // If set, this is the exception to be thrown at the resume point.
+  Tag* exceptionTag = nullptr;
+
   // Whether we executed. Continuations are one-shot, so they may not be
   // executed a second time.
   bool executed = false;
@@ -528,9 +531,8 @@ class ExpressionRunner : public OverriddenVisitor<SubType, Flow> {
       }
       if (!hasValue) {
         // We must execute this instruction. Set up the logic to note the values
-        // of children (we mainly need this for non-control flow structures,
-        // but even control flow ones must add a scope on the value stack, to
-        // not confuse the others).
+        // of children. TODO: as an optimization, we could avoid this for
+        // control flow structures, at the cost of more complexity
         StackValueNoter noter(this);
 
         if (Properties::isControlFlowStructure(curr)) {
@@ -592,12 +594,16 @@ class ExpressionRunner : public OverriddenVisitor<SubType, Flow> {
       // values on the stack, not values sent on a break/suspend; suspending is
       // handled above).
       if (!ret.breaking() && ret.getType().isConcrete()) {
-        assert(!valueStack.empty());
-        auto& values = valueStack.back();
-        values.push_back(ret.values);
+        // The value stack may be empty, if we lack a parent that needs our
+        // value. That is the case when we are the toplevel expression, etc.
+        if (!valueStack.empty()) {
+          auto& values = valueStack.back();
+          values.push_back(ret.values);
 #if WASM_INTERPRETER_DEBUG
-        std::cout << indent() << "added to valueStack: " << ret.values << '\n';
+          std::cout << indent() << "added to valueStack: " << ret.values
+                    << '\n';
 #endif
+        }
       }
     }
 
@@ -4863,6 +4869,12 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
       // restoredValues map.
       assert(currContinuation->resumeInfo.empty());
       assert(self()->restoredValuesMap.empty());
+      // Throw, if we were resumed by resume_throw;
+      if (auto* tag = currContinuation->exceptionTag) {
+        // XXX tag->name lacks cross-module support
+        throwException(WasmException{
+          self()->makeExnData(tag->name, currContinuation->resumeArguments)});
+      }
       return currContinuation->resumeArguments;
     }
 
@@ -4895,7 +4907,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     new_->resumeExpr = curr;
     return Flow(SUSPEND_FLOW, tag, std::move(arguments));
   }
-  Flow visitResume(Resume* curr) {
+  template<typename T> Flow doResume(T* curr, Tag* exceptionTag = nullptr) {
     Literals arguments;
     Flow flow = self()->generateArguments(curr->operands, arguments);
     if (flow.breaking()) {
@@ -4923,6 +4935,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
         // the immediate ones here. TODO
         contData->resumeArguments = arguments;
       }
+      contData->exceptionTag = exceptionTag;
       self()->pushCurrContinuation(contData);
       self()->continuationStore->resuming = true;
 #if WASM_INTERPRETER_DEBUG
@@ -4931,15 +4944,8 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
 #endif
     }
 
-    Flow ret;
-    {
-      // Create a stack value scope. This ensures that we always have a scope,
-      // and so the code that pushes/pops doesn't need to check if a scope
-      // exists. (We do not need the values in this scope, of course, as no
-      // expression is above them, so we cannot suspend and need these values).
-      typename ExpressionRunner<SubType>::StackValueNoter noter(this);
-      ret = func.getFuncData()->doCall(arguments);
-    }
+    Flow ret = func.getFuncData()->doCall(arguments);
+
 #if WASM_INTERPRETER_DEBUG
     if (!self()->isResuming()) {
       std::cout << self()->indent() << "finished resuming, with " << ret
@@ -4995,7 +5001,11 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     // No suspension; all done.
     return ret;
   }
-  Flow visitResumeThrow(ResumeThrow* curr) { return Flow(NONCONSTANT_FLOW); }
+  Flow visitResume(Resume* curr) { return doResume(curr); }
+  Flow visitResumeThrow(ResumeThrow* curr) {
+    // TODO: should the Resume and ResumeThrow classes be merged?
+    return doResume(curr, self()->getModule()->getTag(curr->tag));
+  }
   Flow visitStackSwitch(StackSwitch* curr) { return Flow(NONCONSTANT_FLOW); }
 
   void trap(const char* why) override {
@@ -5058,14 +5068,20 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
 
     if (self()->isResuming()) {
       // The arguments are in the continuation data.
-      arguments = self()->getCurrContinuation()->resumeArguments;
+      auto currContinuation = self()->getCurrContinuation();
+      arguments = currContinuation->resumeArguments;
 
-      if (!self()->getCurrContinuation()->resumeExpr) {
+      if (!currContinuation->resumeExpr) {
         // This is the first time we resume, that is, there is no suspend which
         // is the resume expression that we need to execute up to. All we need
         // to do is just start calling this function (with the arguments we've
-        // set), so resuming is done
+        // set), so resuming is done. (And throw, if resume_throw.)
         self()->continuationStore->resuming = false;
+        if (auto* tag = currContinuation->exceptionTag) {
+          // XXX tag->name lacks cross-module support
+          throwException(WasmException{
+            self()->makeExnData(tag->name, currContinuation->resumeArguments)});
+        }
       }
     }
 
diff --git a/test/lit/exec/cont_simple.wast b/test/lit/exec/cont_simple.wast
@@ -844,4 +844,24 @@
       )
     )
   )
+
+  (func $never
+    ;; This will be resume_throw'd on the first execution, so this code is never
+    ;; reached.
+    (call $log (i32.const 1337))
+  )
+
+  ;; CHECK:      [fuzz-exec] calling resume_throw-never
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 42]
+  (func $resume_throw-never (export "resume_throw-never")
+    (block $more
+      (try_table (catch $more $more)
+        (resume_throw $k $more
+          (cont.new $k (ref.func $never))
+        )
+      )
+    )
+    ;; This will be reached.
+    (call $log (i32.const 42))
+  )
 )
diff --git a/test/spec/cont.wast b/test/spec/cont.wast
@@ -124,8 +124,8 @@
 (assert_return (invoke "handled"))
 
 (assert_exception (invoke "uncaught-1"))
-;; TODO: resume_throw (assert_exception (invoke "uncaught-2"))
-;; TODO: resume_throw (assert_exception (invoke "uncaught-3"))
+(assert_exception (invoke "uncaught-2"))
+(assert_exception (invoke "uncaught-3"))
 
 (assert_trap (invoke "non-linear-1") "continuation already consumed")
 (assert_trap (invoke "non-linear-2") "continuation already consumed")
@@ -641,3 +641,88 @@
     (unreachable))
 )
 
+(module $co2
+  (type $task (func (result i32))) ;; type alias task = [] -> []
+  (type $ct   (cont $task)) ;; type alias   ct = $task
+  (tag $pause (export "pause"))   ;; pause : [] -> []
+  (tag $cancel (export "cancel"))   ;; cancel : [] -> []
+  ;; run : [(ref $task) (ref $task)] -> []
+  ;; implements a 'seesaw' (c.f. Ganz et al. (ICFP@99))
+  (func $run (export "seesaw") (param $up (ref $ct)) (param $down (ref $ct)) (result i32)
+    (local $result i32)
+    ;; run $up
+    (loop $run_next (result i32)
+      (block $on_pause (result (ref $ct))
+        (resume $ct (on $pause $on_pause)
+                    (local.get $up))
+        ;; $up finished, store its result
+        (local.set $result)
+        ;; next cancel $down
+        (block $on_cancel
+          (try_table (catch $cancel $on_cancel)
+            ;; inject the cancel exception into $down
+            (resume_throw $ct $cancel (local.get $down))
+            (drop) ;; drop the return value if it handled $cancel
+                   ;; itself and returned normally...
+          )
+        ) ;; ... otherwise catch $cancel and return $up's result.
+        (return (local.get $result))
+      ) ;; on_pause clause, stack type: [(cont $ct)]
+      (local.set $up)
+      ;; swap $up and $down
+      (local.get $down)
+      (local.set $down (local.get $up))
+      (local.set $up)
+      (br $run_next)
+    )
+  )
+)
+(register "co2")
+
+(module $client
+  (type $task-0 (func (param i32) (result i32)))
+  (type $ct-0 (cont $task-0))
+  (type $task (func (result i32)))
+  (type $ct (cont $task))
+
+  (func $seesaw (import "co2" "seesaw") (param (ref $ct)) (param (ref $ct)) (result i32))
+  (func $print-i32 (import "spectest" "print_i32") (param i32))
+  (tag $pause (import "co2" "pause"))
+
+  (func $even (param $niter i32) (result i32)
+     (local $next i32) ;; zero initialised.
+     (local $i i32)
+     (loop $print-next
+       (call $print-i32 (local.get $next))
+       (suspend $pause)
+       (local.set $next (i32.add (local.get $next) (i32.const 2)))
+       (local.set $i (i32.add (local.get $i) (i32.const 1)))
+       (br_if $print-next (i32.lt_u (local.get $i) (local.get $niter)))
+     )
+     (local.get $next)
+  )
+  (func $odd (param $niter i32) (result i32)
+     (local $next i32) ;; zero initialised.
+     (local $i i32)
+     (local.set $next (i32.const 1))
+     (loop $print-next
+       (call $print-i32 (local.get $next))
+       (suspend $pause)
+       (local.set $next (i32.add (local.get $next) (i32.const 2)))
+       (local.set $i (i32.add (local.get $i) (i32.const 1)))
+       (br_if $print-next (i32.lt_u (local.get $i) (local.get $niter)))
+     )
+     (local.get $next)
+  )
+
+  (func (export "main") (result i32)
+    (call $seesaw
+       (cont.bind $ct-0 $ct
+         (i32.const 5) (cont.new $ct-0 (ref.func $even)))
+       (cont.bind $ct-0 $ct
+         (i32.const 5) (cont.new $ct-0 (ref.func $odd)))))
+
+  (elem declare func $even $odd)
+)
+(assert_return (invoke "main") (i32.const 10))
+
