[Stack Switching] Supporting sending/receiving values from continuations (#7775)

kripken · web-flow · commit 62dae747f248 · 2025-08-04T15:33:35.000-07:00
With this, the next part of the spec test can be added.
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
@@ -191,6 +191,11 @@ struct ContData {
   // will emit a single Literals for itself, or possibly a few bundles.
   std::vector<Literals> resumeInfo;
 
+  // The arguments sent when resuming (on first execution these appear as
+  // parameters to the function; on later resumes, they are returned from the
+  // suspend).
+  Literals resumeArguments;
+
   // Whether we executed. Continuations are one-shot, so they may not be
   // executed a second time.
   bool executed = false;
@@ -403,6 +408,10 @@ class ExpressionRunner : public OverriddenVisitor<SubType, Flow> {
         auto iter = restoredValuesMap.find(curr);
         if (iter != restoredValuesMap.end()) {
           ret = iter->second;
+#if WASM_INTERPRETER_DEBUG
+          std::cout << indent() << "consume restored value: " << ret.values
+                    << '\n';
+#endif
           restoredValuesMap.erase(iter);
           hasValue = true;
         }
@@ -437,6 +446,10 @@ class ExpressionRunner : public OverriddenVisitor<SubType, Flow> {
               --num;
               auto value = popResumeEntry("child value");
               restoredValuesMap[child] = value;
+#if WASM_INTERPRETER_DEBUG
+              std::cout << indent() << "prepare restored value: " << value
+                        << '\n';
+#endif
             }
           }
 
@@ -4636,6 +4649,15 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
   }
   Flow visitContBind(ContBind* curr) { return Flow(NONCONSTANT_FLOW); }
   Flow visitSuspend(Suspend* curr) {
+    // Process the arguments, whether or not we are resuming. If we are resuming
+    // then we don't need these values (we sent them as part of the suspension),
+    // but must still handle them, so we finish re-winding the stack.
+    Literals arguments;
+    Flow flow = self()->generateArguments(curr->operands, arguments);
+    if (flow.breaking()) {
+      return flow;
+    }
+
     if (self()->resuming) {
       // This is a resume, so we have found our way back to where we
       // suspended.
@@ -4646,15 +4668,10 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
       // restoredValues map.
       assert(self()->currContinuation->resumeInfo.empty());
       assert(self()->restoredValuesMap.empty());
-      return Flow();
+      return self()->currContinuation->resumeArguments;
     }
 
     // We were not resuming, so this is a new suspend that we must execute.
-    Literals arguments;
-    Flow flow = self()->generateArguments(curr->operands, arguments);
-    if (flow.breaking()) {
-      return flow;
-    }
 
     // Copy the continuation (the old one cannot be resumed again). Note that no
     // old one may exist, in which case we still emit a continuation, but it is
@@ -4688,6 +4705,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
       trap("continuation already executed");
     }
     contData->executed = true;
+    contData->resumeArguments = arguments;
     Name func = contData->func;
     self()->currContinuation = contData;
     if (contData->resumeExpr) {
@@ -4699,12 +4717,24 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
 #if WASM_INTERPRETER_DEBUG
     std::cout << self()->indent() << "resuming func " << func << '\n';
 #endif
-    Flow ret = callFunction(func, arguments);
+    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 = callFunction(func, arguments);
+    }
 #if WASM_INTERPRETER_DEBUG
     std::cout << self()->indent() << "finished resuming, with " << ret << '\n';
 #endif
-    if (ret.suspendTag) {
-      // See if a suspension arrived that we support.
+    if (!ret.suspendTag) {
+      // No suspension: the coroutine finished normally. Mark it as no longer
+      // active.
+      self()->currContinuation.reset();
+    } else {
+      // We are suspending. See if a suspension arrived that we support.
       for (size_t i = 0; i < curr->handlerTags.size(); i++) {
         auto handlerTag = curr->handlerTags[i];
         if (handlerTag == ret.suspendTag) {
@@ -4719,7 +4749,6 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
         }
       }
       // No handler worked out, keep propagating.
-      return ret;
     }
     // No suspension; all done.
     return ret;
diff --git a/test/lit/exec/cont_simple.wast b/test/lit/exec/cont_simple.wast
@@ -2,13 +2,20 @@
 
 ;; RUN: foreach %s %t wasm-opt -all --fuzz-exec-before -q -o /dev/null 2>&1 | filecheck %s
 
-(module $state
+(module
   (type $f (func))
   (type $k (cont $f))
 
+  (type $f-i32 (func (result i32)))
+  (type $k-i32 (cont $f-i32))
+
+  (type $f-get-i32 (func (param i32)))
+  (type $k-get-i32 (cont $f-get-i32))
+
   (import "fuzzing-support" "log" (func $log (param i32)))
 
   (tag $more)
+  (tag $more-i32 (result i32))
 
   (func $run (param $k (ref $k))
     ;; Run a coroutine, continuing to resume it until it is complete.
@@ -552,4 +559,116 @@
       (cont.new $k (ref.func $trinary))
     )
   )
+
+  (func $run-i32 (param $k-i32 (ref $k-i32)) (result i32)
+    ;; As $run, but the coroutine returns an i32.
+    (call $log (i32.const 100)) ;; start
+    (loop $loop
+      (block $on (result (ref $k-i32))
+        (resume $k-i32 (on $more $on)
+          (local.get $k-i32)
+        )
+        (call $log (i32.const 300)) ;; stop
+        (return)
+      )
+      (call $log (i32.const 200)) ;; continue
+      (local.set $k-i32)
+      (br $loop)
+    )
+    (unreachable)
+  )
+
+  (func $ret-i32 (result i32)
+    ;; Just immediately return.
+    (i32.const 42)
+  )
+
+  ;; CHECK:      [fuzz-exec] calling run-ret-i32
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 100]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 300]
+  ;; CHECK-NEXT: [fuzz-exec] note result: run-ret-i32 => 42
+  (func $run-ret-i32 (export "run-ret-i32") (result i32)
+    (call $run-i32
+      (cont.new $k-i32 (ref.func $ret-i32))
+    )
+  )
+
+  (func $pause-i32 (result i32)
+    (local $x i32)
+    ;; Pause before returning.
+    (local.set $x
+      (i32.const 1336)
+    )
+    (suspend $more)
+    (i32.add
+      (local.get $x)
+      (i32.const 1)
+    )
+  )
+
+  ;; CHECK:      [fuzz-exec] calling run-pause-i32
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 100]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 200]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 300]
+  ;; CHECK-NEXT: [fuzz-exec] note result: run-pause-i32 => 1337
+  (func $run-pause-i32 (export "run-pause-i32") (result i32)
+    (call $run-i32
+      (cont.new $k-i32 (ref.func $pause-i32))
+    )
+  )
+
+  (func $run-get-i32 (param $x i32) (param $k-get-i32 (ref $k-get-i32))
+    ;; As $run, but the coroutine receives an i32.
+    (call $log (i32.const 100)) ;; start
+    (loop $loop
+      (block $on (result (ref $k-get-i32))
+        (resume $k-get-i32 (on $more-i32 $on)
+          (local.get $x)
+          (local.get $k-get-i32)
+        )
+        (call $log (i32.const 300)) ;; stop
+        (return)
+      )
+      (call $log (i32.const 200)) ;; continue
+      ;; Modify $x, so we can see differences in the loggings.
+      (local.set $x
+        (i32.sub
+          (local.get $x)
+          (i32.const 1)
+        )
+      )
+      (local.set $k-get-i32)
+      (br $loop)
+    )
+    (unreachable)
+  )
+
+  (func $param (param $x i32)
+    (call $log (local.get $x))
+    (local.set $x
+      (i32.add
+        (local.get $x)
+        (i32.const 1295)
+      )
+    )
+    (call $log (suspend $more-i32))
+    (call $log (local.get $x))
+    (call $log (suspend $more-i32))
+  )
+
+  ;; CHECK:      [fuzz-exec] calling run-param
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 100]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 42]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 200]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 41]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 1337]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 200]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 40]
+  ;; CHECK-NEXT: [LoggingExternalInterface logging 300]
+  (func $run-param (export "run-param")
+    (call $run-get-i32
+      (i32.const 42)
+      (cont.new $k-get-i32 (ref.func $param))
+    )
+  )
 )
diff --git a/test/spec/cont.wast b/test/spec/cont.wast
@@ -227,3 +227,55 @@
   (type $c2 (cont $f2))
 )
 
+;; Simple state example
+
+(module $state
+  (tag $get (result i32))
+  (tag $set (param i32) (result i32))
+
+  (type $f (func (param i32) (result i32)))
+  (type $k (cont $f))
+
+  (func $runner (param $s i32) (param $k (ref $k)) (result i32)
+    (loop $loop
+      (block $on_get (result (ref $k))
+        (block $on_set (result i32 (ref $k))
+          (resume $k (on $get $on_get) (on $set $on_set)
+            (local.get $s) (local.get $k)
+          )
+          (return)
+        )
+        ;; on set
+        (local.set $k)
+        (local.set $s)
+        (br $loop)
+      )
+      ;; on get
+      (local.set $k)
+      (br $loop)
+    )
+    (unreachable)
+  )
+
+  (func $f (param i32) (result i32)
+    (drop (suspend $set (i32.const 7)))
+    (i32.add
+      (suspend $get)
+      (i32.mul
+        (i32.const 2)
+        (i32.add
+          (suspend $set (i32.const 3))
+          (suspend $get)
+        )
+      )
+    )
+  )
+
+  (elem declare func $f)
+  (func (export "run") (result i32)
+    (call $runner (i32.const 0) (cont.new $k (ref.func $f)))
+  )
+)
+
+(assert_return (invoke "run") (i32.const 19))
+
