[Stack Switching] Fix resuming of return calls (#7860)

A return call is a form of control flow, so we need to handle it. Specifically,
the original function is unwound before calling the next one, so when we
resume, we should just call the next one. To do that, save the call target
along with the locals during unwinding.

Author: kripken

src/wasm-interpreter.h

@@ -4664,6 +4664,30 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
 
     // We may have to call multiple functions in the event of return calls.
     while (true) {
+      if (self()->isResuming()) {
+        // See which function to call. Re-winding the stack, we are calling the
+        // function that the parent called, but the target that was called may
+        // have return-called. In that case, the original target function should
+        // not be called, as it was returned from, and we noted the proper
+        // target during that return.
+        auto entry = self()->popResumeEntry("function-target");
+        assert(entry.size() == 1);
+        auto func = entry[0];
+        auto data = func.getFuncData();
+        // We must be in the right module to do the call using that name.
+        if (data->self != self()) {
+          // Restore the entry to the resume stack, as the other module's
+          // callFunction() will read it. Then call into the other module. This
+          // sets this up as if we called into the proper module in the first
+          // place.
+          self()->pushResumeEntry(entry, "function-target");
+          return data->doCall(arguments);
+        }
+
+        // We are in the right place, and can just call the given function.
+        name = data->name;
+      }
+
       Function* function = wasm.getFunction(name);
       assert(function);
 
@@ -4684,7 +4708,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
         // Restore the local state (see below for the ordering, we push/pop).
         for (Index i = 0; i < scope.locals.size(); i++) {
           auto l = scope.locals.size() - 1 - i;
-          scope.locals[l] = self()->popResumeEntry("function");
+          scope.locals[l] = self()->popResumeEntry("function-local");
 #ifndef NDEBUG
           // Must have restored valid data. The type must match the local's
           // type, except for the case of a non-nullable local that has not yet
@@ -4718,8 +4742,15 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
       if (flow.suspendTag) {
         // Save the local state.
         for (auto& local : scope.locals) {
-          self()->pushResumeEntry(local, "function");
+          self()->pushResumeEntry(local, "function-local");
         }
+
+        // Save the function we called (in the case of a return call, this is
+        // not the original function that was called, and the original has been
+        // returned from already; we should call the last return_called
+        // function).
+        auto target = self()->makeFuncData(name, function->type);
+        self()->pushResumeEntry({target}, "function-target");
       }
 
       if (flow.breakTo != RETURN_CALL_FLOW) {

test/lit/exec/cont_return_call.wast

@@ -0,0 +1,75 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --output=fuzz-exec and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt -all --fuzz-exec-before -q -o /dev/null 2>&1 | filecheck %s
+
+;; Test suspend/resume through a return call. $cont return_calls $sub, which
+;; then suspends. The return_call removes $cont from the stack before that
+;; suspend, so we should jump into $sub during resuming, and not hit any errors.
+
+(module
+ (type $i32 (func (result i32)))
+ (type $cont (cont $i32))
+ (type $none (func))
+
+ (import "fuzzing-support" "log" (func $log (param i32)))
+
+ (tag $tag (type $none))
+
+ ;; CHECK:      [fuzz-exec] calling main
+ ;; CHECK-NEXT: [LoggingExternalInterface logging 50]
+ ;; CHECK-NEXT: [LoggingExternalInterface logging -1]
+ ;; CHECK-NEXT: [LoggingExternalInterface logging 100]
+ ;; CHECK-NEXT: [LoggingExternalInterface logging 200]
+ ;; CHECK-NEXT: [LoggingExternalInterface logging -1]
+ ;; CHECK-NEXT: [LoggingExternalInterface logging 300]
+ ;; CHECK-NEXT: [LoggingExternalInterface logging 400]
+ ;; CHECK-NEXT: [LoggingExternalInterface logging -1]
+ ;; CHECK-NEXT: [LoggingExternalInterface logging 500]
+ ;; CHECK-NEXT: [trap unreachable]
+ (func $main (export "main")
+  (local $c (ref $cont))
+  ;; Create a continuation and keep resuming it while it suspends.
+  (local.set $c
+   (cont.new $cont
+    (ref.func $cont)
+   )
+  )
+  (loop $label
+   (local.set $c
+    (block $block (result (ref $cont))
+     (drop
+      (resume $cont (on $tag $block)
+       (local.get $c)
+      )
+     )
+     (unreachable)
+    )
+   )
+   (call $log (i32.const -1)) ;; logged each time we suspend
+   (br $label)
+  )
+ )
+
+ (func $cont (type $i32) (result i32)
+  (call $log (i32.const 50)) ;; only reached once
+  (suspend $tag)
+  (call $log (i32.const 100)) ;; only reached once
+  (return_call $sub)
+  (call $log (i32.const -100)) ;; never reached
+ )
+
+ (func $sub (result i32)
+  (call $log (i32.const 200))
+  (suspend $tag)
+  (call $log (i32.const 300))
+  (return_call $subsub)
+ )
+
+ (func $subsub (result i32)
+  (call $log (i32.const 400))
+  (suspend $tag)
+  (call $log (i32.const 500))
+  (i32.const 0)
+ )
+)
+