[Stack Switching] Handle imported tags properly (#7801)

We were using Tag names, but the internal name in a module may
overlap with other modules etc. - we really need to look at the
actual imported Tag, if it is imported. To do that, store Tag* and
add a way to get the canonical Tag* - i.e., the imported one if
imported, so all references, including from other modules, all
refer to the same thing.

To be able to test this, make `callImport` return a Flow (so we can
suspend through import calls).

As a bonus, replacing `Name tag` with `Tag* tag` makes our data
structures smaller.

Author: kripken

scripts/test/fuzzing.py

@@ -105,6 +105,7 @@
     'vacuum-desc.wast',
     'br_on_cast_desc.wast',
     # TODO: fuzzer support for stack switching
+    'tag_linked.wast',
     'stack_switching.wast',
     'stack_switching_contnew.wast',
     'stack_switching_contbind.wast',

src/shell-interface.h

@@ -133,7 +133,7 @@ struct ShellExternalInterface : ModuleRunner::ExternalInterface {
     });
   }
 
-  Literals callImport(Function* import, const Literals& arguments) override {
+  Flow callImport(Function* import, const Literals& arguments) override {
     if (import->module == SPECTEST && import->base.startsWith(PRINT)) {
       for (auto argument : arguments) {
         std::cout << argument << " : " << argument.type << '\n';
@@ -144,9 +144,7 @@ struct ShellExternalInterface : ModuleRunner::ExternalInterface {
       std::cout << "exit()\n";
       throw ExitException();
     } else if (auto* inst = getImportInstance(import)) {
-      auto flow = inst->callExport(import->base, arguments);
-      assert(!flow.suspendTag); // TODO: support stack switching on calls
-      return flow.values;
+      return inst->callExport(import->base, arguments);
     }
     Fatal() << "callImport: unknown import: " << import->module.str << "."
             << import->name.str;

src/tools/execution-results.h

@@ -75,7 +75,7 @@ struct LoggingExternalInterface : public ShellExternalInterface {
     }
   }
 
-  Literals callImport(Function* import, const Literals& arguments) override {
+  Flow callImport(Function* import, const Literals& arguments) override {
     if (import->module == "fuzzing-support") {
       if (import->base.startsWith("log")) {
         // This is a logging function like log-i32 or log-f64

src/tools/wasm-ctor-eval.cpp

@@ -242,7 +242,7 @@ struct CtorEvalExternalInterface : EvallingModuleRunner::ExternalInterface {
     });
   }
 
-  Literals callImport(Function* import, const Literals& arguments) override {
+  Flow callImport(Function* import, const Literals& arguments) override {
     Name WASI("wasi_snapshot_preview1");
 
     if (ignoreExternalInput) {

src/wasm-interpreter.h

@@ -80,15 +80,15 @@ class Flow {
   Flow(Name breakTo, Literal value) : values{value}, breakTo(breakTo) {}
   Flow(Name breakTo, Literals&& values)
     : values(std::move(values)), breakTo(breakTo) {}
-  Flow(Name breakTo, Name suspendTag, Literals&& values)
+  Flow(Name breakTo, Tag* suspendTag, Literals&& values)
     : values(std::move(values)), breakTo(breakTo), suspendTag(suspendTag) {
     assert(breakTo == SUSPEND_FLOW);
   }
 
   Literals values;
   Name breakTo; // if non-null, a break is going on
-  Name suspendTag; // if non-null, breakTo must be SUSPEND_FLOW, and this is the
-                   // tag being suspended
+  Tag* suspendTag = nullptr; // if non-null, breakTo must be SUSPEND_FLOW, and
+                             // this is the tag being suspended
 
   // A helper function for the common case where there is only one value
   const Literal& getSingleValue() {
@@ -123,7 +123,7 @@ class Flow {
       o << flow.values[i];
     }
     if (flow.suspendTag) {
-      o << " [suspend:" << flow.suspendTag << ']';
+      o << " [suspend:" << flow.suspendTag->name << ']';
     }
     o << "})";
     return o;
@@ -3056,8 +3056,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     virtual ~ExternalInterface() = default;
     virtual void init(Module& wasm, SubType& instance) {}
     virtual void importGlobals(GlobalValueSet& globals, Module& wasm) = 0;
-    virtual Literals callImport(Function* import,
-                                const Literals& arguments) = 0;
+    virtual Flow callImport(Function* import, const Literals& arguments) = 0;
     virtual bool growMemory(Name name, Address oldSize, Address newSize) = 0;
     virtual bool growTable(Name name,
                            const Literal& value,
@@ -3564,6 +3563,21 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     return inst->globals[global->name];
   }
 
+  // Get a tag object while looking through imports, i.e., this uses the name as
+  // the name of the tag in the current module, and finds the actual canonical
+  // Tag* object for it: the Tag in this module, if not imported, and if
+  // imported, the Tag in the originating module.
+  Tag* getCanonicalTag(Name name) {
+    auto* inst = self();
+    auto* tag = inst->wasm.getTag(name);
+    while (tag->imported()) {
+      inst = inst->linkedInstances.at(tag->module).get();
+      auto* tagExport = inst->wasm.getExport(tag->base);
+      tag = inst->wasm.getTag(*tagExport->getInternalName());
+    }
+    return tag;
+  }
+
 public:
   Flow visitCall(Call* curr) {
     Name target = curr->target;
@@ -4775,7 +4789,8 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     // We will resume from this precise spot, when the new continuation is
     // resumed.
     new_->resumeExpr = curr;
-    return Flow(SUSPEND_FLOW, curr->tag, std::move(arguments));
+    return Flow(
+      SUSPEND_FLOW, self()->getCanonicalTag(curr->tag), std::move(arguments));
   }
   Flow visitResume(Resume* curr) {
     Literals arguments;
@@ -4825,10 +4840,10 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     } 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];
+        auto* handlerTag = self()->getCanonicalTag(curr->handlerTags[i]);
         if (handlerTag == ret.suspendTag) {
           // Switch the flow from suspending to branching.
-          ret.suspendTag = Name();
+          ret.suspendTag = nullptr;
           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
@@ -4978,7 +4993,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
 
 #if WASM_INTERPRETER_DEBUG
       std::cout << self()->indent() << "exiting " << function->name << " with "
-                << flow.values << '\n';
+                << flow << '\n';
 #endif
 
       if (flow.suspendTag) {

test/spec/tag_linked.wast

@@ -0,0 +1,45 @@
+;; Test that we can properly use a tag from an imported module. The export name
+;; and the internal names all differ, to make sure we properly identify the tag.
+
+(module
+  (tag $suspend)
+
+  (export "suspendey" (tag $suspend))
+  (export "suspender" (func $suspender))
+
+  (func $suspender
+    (suspend $suspend)
+  )
+)
+
+(register "first")
+
+(module
+  (type $func (func))
+  (type $cont (cont $func))
+
+  (tag $imported-suspend (import "first" "suspendey"))
+
+  (import "first" "suspender" (func $suspender))
+
+  (tag $suspend) ;; name overlap
+
+  (func (export "run") (result i32)
+    (drop
+      (block $internal (result (ref $cont))
+        (block $imported (result (ref $cont))
+          (resume $cont (on $imported-suspend $imported) (on $suspend $internal)
+            (cont.new $cont (ref.func $suspender))
+          )
+          (unreachable)
+        )
+        (return (i32.const 42))
+      )
+      (return (i32.const 1337))
+    )
+    (unreachable)
+  )
+)
+
+(assert_return (invoke "run") (i32.const 42))
+