[Stack Switching] Update continuation types based on resume tag & target block (#7786)

The resume tag defines how the continuation should be called,
perhaps with different parameters than before. To handle that,
update its type and also properly handle function parameters
(if we are resuming with a different type, we should not error -
the function params are already saved as local state, and we
just need to reload them). The actual type also depends on the
target block.

This allows the next part of the spec test to run.

Author: kripken

src/wasm-interpreter.h

@@ -3458,14 +3458,20 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
       parent.scope = this;
       parent.callDepth++;
       parent.functionStack.push_back(function->name);
+      locals.resize(function->getNumLocals());
+
+      if (parent.resuming) {
+        // Nothing more to do here: we are resuming execution, so there is old
+        // locals state that will be restored.
+        return;
+      }
 
       if (function->getParams().size() != arguments.size()) {
         std::cerr << "Function `" << function->name << "` expects "
                   << function->getParams().size() << " parameters, got "
                   << arguments.size() << " arguments." << std::endl;
         WASM_UNREACHABLE("invalid param count");
       }
-      locals.resize(function->getNumLocals());
       Type params = function->getParams();
       for (size_t i = 0; i < function->getNumLocals(); i++) {
         if (i < arguments.size()) {
@@ -4697,6 +4703,10 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     assert(!old || old->executed);
     auto new_ = std::make_shared<ContData>(old ? old->func : Name(),
                                            old ? old->type : HeapType::none);
+    // Note we cannot update the type yet, so it will be wrong in debug
+    // logging. To update it, we must find the block that receives this value,
+    // which means we cannot do it here (we don't even know what that block is).
+
     // Switch to the new continuation, so that as we unwind, we will save the
     // information we need to resume it later in the proper place.
     self()->currContinuation = new_;
@@ -4758,8 +4768,30 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
           // Switch the flow from suspending to branching.
           ret.suspendTag = Name();
           ret.breakTo = curr->handlerBlocks[i];
+          // We can now update the continuation type, which was wrong until now
+          // (see comment in visitSuspend). The type is taken from the block we
+          // branch to (which we find in a quite inefficient manner).
+          struct BlockFinder : public PostWalker<BlockFinder> {
+            Name target;
+            Type type = Type::none;
+            void visitBlock(Block* curr) {
+              if (curr->name == target) {
+                type = curr->type;
+              }
+            }
+          } finder;
+          finder.target = ret.breakTo;
+          // We must be in a function scope.
+          assert(self()->scope->function);
+          finder.walk(self()->scope->function->body);
+          // We must have found the type, and it must be valid.
+          assert(finder.type.isConcrete());
+          assert(finder.type.size() >= 1);
+          // The continuation is the final value/type there.
+          auto cont = self()->currContinuation;
+          cont->type = finder.type[finder.type.size() - 1].getHeapType();
           // Add the continuation as the final value being sent.
-          ret.values.push_back(Literal(self()->currContinuation));
+          ret.values.push_back(Literal(cont));
           // We are not longer processing that continuation.
           self()->currContinuation.reset();
           return ret;

test/spec/cont.wast

@@ -279,3 +279,58 @@
 
 (assert_return (invoke "run") (i32.const 19))
 
+;; Simple generator example
+
+(module $generator
+  (type $gen (func (param i64)))
+  (type $geny (func (param i32)))
+  (type $cont0 (cont $gen))
+  (type $cont (cont $geny))
+
+  (tag $yield (param i64) (result i32))
+
+  ;; Hook for logging purposes
+  (global $hook (export "hook") (mut (ref $gen)) (ref.func $dummy))
+  (func $dummy (param i64))
+
+  (func $gen (export "start") (param $i i64)
+    (loop $l
+      (br_if 1 (suspend $yield (local.get $i)))
+      (call_ref $gen (local.get $i) (global.get $hook))
+      (local.set $i (i64.add (local.get $i) (i64.const 1)))
+      (br $l)
+    )
+  )
+
+  (elem declare func $gen)
+
+  (func (export "sum") (param $i i64) (param $j i64) (result i64)
+    (local $sum i64)
+    (local $n i64)
+    (local $k (ref null $cont))
+    (local.get $i)
+    (cont.new $cont0 (ref.func $gen))
+    (block $on_first_yield (param i64 (ref $cont0)) (result i64 (ref $cont))
+      (resume $cont0 (on $yield $on_first_yield))
+      (unreachable)
+    )
+    (loop $on_yield (param i64) (param (ref $cont))
+      (local.set $k)
+      (local.set $n)
+      (local.set $sum (i64.add (local.get $sum) (local.get $n)))
+      (i64.eq (local.get $n) (local.get $j))
+      (local.get $k)
+      (resume $cont (on $yield $on_yield))
+    )
+    (return (local.get $sum))
+  )
+)
+
+(register "generator")
+
+(assert_return (invoke "sum" (i64.const 0) (i64.const 0)) (i64.const 0))
+(assert_return (invoke "sum" (i64.const 2) (i64.const 2)) (i64.const 2))
+(assert_return (invoke "sum" (i64.const 0) (i64.const 3)) (i64.const 6))
+(assert_return (invoke "sum" (i64.const 1) (i64.const 10)) (i64.const 55))
+(assert_return (invoke "sum" (i64.const 100) (i64.const 2000)) (i64.const 1_996_050))
+