[StackSwitching] Add resume_throw_ref (#8140)

Implement it as a variation on the existing ResumeThrow class: if the
tag exists, it is a resume_throw with that tag; if the tag is null, it is
resume_throw_ref.

Remove an assert in visitTryTable in the interpreter: the code actually
works, and no TODO was necessary there: resuming "just works", like
with CallRef etc. (visitTry does have a problem, resuming from a
try's catch block - but try_table is fine.) This allows the test to pass.

Author: kripken

scripts/gen-s-parser.py

@@ -609,6 +609,7 @@
     ("suspend",              "makeSuspend()"),
     ("resume",               "makeResume()"),
     ("resume_throw",         "makeResumeThrow()"),
+    ("resume_throw_ref",     "makeResumeThrowRef()"),
     ("switch",               "makeStackSwitch()"),
     # GC
     ("ref.i31",              "makeRefI31(Unshared)"),

src/gen-s-parser.inc

@@ -4972,12 +4972,23 @@ switch (buf[0]) {
               return Ok{};
             }
             goto parse_error;
-          case '_':
-            if (op == "resume_throw"sv) {
-              CHECK_ERR(makeResumeThrow(ctx, pos, annotations));
-              return Ok{};
+          case '_': {
+            switch (buf[12]) {
+              case '\0':
+                if (op == "resume_throw"sv) {
+                  CHECK_ERR(makeResumeThrow(ctx, pos, annotations));
+                  return Ok{};
+                }
+                goto parse_error;
+              case '_':
+                if (op == "resume_throw_ref"sv) {
+                  CHECK_ERR(makeResumeThrowRef(ctx, pos, annotations));
+                  return Ok{};
+                }
+                goto parse_error;
+              default: goto parse_error;
             }
-            goto parse_error;
+          }
           default: goto parse_error;
         }
       }

src/ir/child-typer.h

@@ -1325,10 +1325,16 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
       ct = curr->cont->type.getHeapType();
     }
     assert(ct->isContinuation());
-    auto params = wasm.getTag(curr->tag)->params();
-    assert(params.size() == curr->operands.size());
-    for (size_t i = 0; i < params.size(); ++i) {
-      note(&curr->operands[i], params[i]);
+    if (curr->tag) {
+      // resume_throw
+      auto params = wasm.getTag(curr->tag)->params();
+      assert(params.size() == curr->operands.size());
+      for (size_t i = 0; i < params.size(); ++i) {
+        note(&curr->operands[i], params[i]);
+      }
+    } else {
+      // resume_throw_ref
+      note(&curr->operands[0], Type(HeapType::exn, Nullable));
     }
     note(&curr->cont, Type(*ct, Nullable));
   }

src/parser/contexts.h

@@ -935,6 +935,13 @@ struct NullInstrParserCtx {
     return Ok{};
   }
   template<typename HeapTypeT>
+  Result<> makeResumeThrowRef(Index,
+                              const std::vector<Annotation>&,
+                              HeapTypeT,
+                              const TagLabelListT&) {
+    return Ok{};
+  }
+  template<typename HeapTypeT>
   Result<>
   makeStackSwitch(Index, const std::vector<Annotation>&, HeapTypeT, TagIdxT) {
     return Ok{};
@@ -2889,24 +2896,41 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
     return withLoc(pos, irBuilder.makeResume(type, tags, labels));
   }
 
+  struct ResumeThrowData {
+    std::vector<Name> tags;
+    std::vector<std::optional<Index>> labels;
+
+    ResumeThrowData(const std::vector<OnClauseInfo>& resumetable) {
+      tags.reserve(resumetable.size());
+      labels.reserve(resumetable.size());
+      for (const OnClauseInfo& info : resumetable) {
+        tags.push_back(info.tag);
+        if (info.isOnSwitch) {
+          labels.push_back(std::nullopt);
+        } else {
+          labels.push_back(std::optional<Index>(info.label));
+        }
+      }
+    }
+  };
+
   Result<> makeResumeThrow(Index pos,
                            const std::vector<Annotation>& annotations,
                            HeapType type,
                            Name tag,
                            const std::vector<OnClauseInfo>& resumetable) {
-    std::vector<Name> tags;
-    std::vector<std::optional<Index>> labels;
-    tags.reserve(resumetable.size());
-    labels.reserve(resumetable.size());
-    for (const OnClauseInfo& info : resumetable) {
-      tags.push_back(info.tag);
-      if (info.isOnSwitch) {
-        labels.push_back(std::nullopt);
-      } else {
-        labels.push_back(std::optional<Index>(info.label));
-      }
-    }
-    return withLoc(pos, irBuilder.makeResumeThrow(type, tag, tags, labels));
+    ResumeThrowData data(resumetable);
+    return withLoc(
+      pos, irBuilder.makeResumeThrow(type, tag, data.tags, data.labels));
+  }
+
+  Result<> makeResumeThrowRef(Index pos,
+                              const std::vector<Annotation>& annotations,
+                              HeapType type,
+                              const std::vector<OnClauseInfo>& resumetable) {
+    ResumeThrowData data(resumetable);
+    return withLoc(pos,
+                   irBuilder.makeResumeThrowRef(type, data.tags, data.labels));
   }
 
   Result<> makeStackSwitch(Index pos,

src/parser/parsers.h

@@ -2743,6 +2743,21 @@ Result<> makeResumeThrow(Ctx& ctx,
   return ctx.makeResumeThrow(pos, annotations, *type, *exnTag, *resumetable);
 }
 
+// resume_throw_ref ::= 'resume_throw' typeidx ('(' 'on' tagidx labelidx |
+// 'on' tagidx switch ')')*
+template<typename Ctx>
+Result<> makeResumeThrowRef(Ctx& ctx,
+                            Index pos,
+                            const std::vector<Annotation>& annotations) {
+  auto type = typeidx(ctx);
+  CHECK_ERR(type);
+
+  auto resumetable = makeResumeTable(ctx);
+  CHECK_ERR(resumetable);
+
+  return ctx.makeResumeThrowRef(pos, annotations, *type, *resumetable);
+}
+
 // switch ::= 'switch' typeidx tagidx
 template<typename Ctx>
 Result<> makeStackSwitch(Ctx& ctx,

src/passes/Print.cpp

@@ -2623,11 +2623,15 @@ struct PrintExpressionContents
   void visitResumeThrow(ResumeThrow* curr) {
     assert(curr->cont->type.isContinuation());
     printMedium(o, "resume_throw");
-
+    if (!curr->tag) {
+      printMedium(o, "_ref");
+    }
     o << ' ';
     printHeapTypeName(curr->cont->type.getHeapType());
-    o << ' ';
-    curr->tag.print(o);
+    if (curr->tag) {
+      o << ' ';
+      curr->tag.print(o);
+    }
 
     handleResumeTable(o, curr);
   }

src/wasm-binary.h

@@ -1246,7 +1246,8 @@ enum ASTNodes {
   Suspend = 0xe2,
   Resume = 0xe3,
   ResumeThrow = 0xe4,
-  Switch = 0xe5,  // NOTE(dhil): the internal class is known as
+  ResumeThrowRef = 0xe5,
+  Switch = 0xe6,  // NOTE(dhil): the internal class is known as
                   // StackSwitch to avoid conflict with the existing
                   // 'switch table'.
   OnLabel = 0x00, // (on $tag $label)

src/wasm-interpreter.h

@@ -273,9 +273,14 @@ struct ContData {
   // suspend).
   Literals resumeArguments;
 
-  // If set, this is the exception to be thrown at the resume point.
+  // If set, this is the tag for an exception to be thrown at the resume point
+  // (from resume_throw).
   Tag* exceptionTag = nullptr;
 
+  // If set, this is the exception ref to be thrown at the resume point (from
+  // resume_throw_ref).
+  Literal exception;
+
   // Whether we executed. Continuations are one-shot, so they may not be
   // executed a second time.
   bool executed = false;
@@ -4480,7 +4485,6 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     }
   }
   Flow visitTryTable(TryTable* curr) {
-    assert(!self()->isResuming()); // TODO
     try {
       return self()->visit(curr->body);
     } catch (const WasmException& e) {
@@ -4557,6 +4561,22 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     old->executed = true;
     return Literal(std::make_shared<ContData>(newData));
   }
+
+  void maybeThrowAfterResuming(std::shared_ptr<ContData>& currContinuation) {
+    // We may throw by creating a tag, or an exnref.
+    auto* tag = currContinuation->exceptionTag;
+    auto exnref = currContinuation->exception.type != Type::none;
+    assert(!(tag && exnref));
+    if (tag) {
+      // resume_throw
+      throwException(WasmException{
+        self()->makeExnData(tag, currContinuation->resumeArguments)});
+    } else if (exnref) {
+      // resume_throw_ref
+      throwException(WasmException{currContinuation->exception});
+    }
+  }
+
   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),
@@ -4579,11 +4599,7 @@ 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) {
-        throwException(WasmException{
-          self()->makeExnData(tag, currContinuation->resumeArguments)});
-      }
+      maybeThrowAfterResuming(currContinuation);
       return currContinuation->resumeArguments;
     }
 
@@ -4616,7 +4632,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     new_->resumeExpr = curr;
     return Flow(SUSPEND_FLOW, tag, std::move(arguments));
   }
-  template<typename T> Flow doResume(T* curr, Tag* exceptionTag = nullptr) {
+  template<typename T> Flow doResume(T* curr) {
     Literals arguments;
     VISIT_ARGUMENTS(flow, curr->operands, arguments)
     VISIT_REUSE(flow, curr->cont);
@@ -4636,13 +4652,26 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
         trap("continuation already executed");
       }
       contData->executed = true;
+
       if (contData->resumeArguments.empty()) {
         // The continuation has no bound arguments. For now, we just handle the
         // simple case of binding all of them, so that means we can just use all
         // the immediate ones here. TODO
         contData->resumeArguments = arguments;
       }
-      contData->exceptionTag = exceptionTag;
+      // Fill in the continuation data. How we do this depends on whether we
+      // are resume or resume_throw*.
+      if (auto* resumeThrow = curr->template dynCast<ResumeThrow>()) {
+        if (resumeThrow->tag) {
+          // resume_throw
+          contData->exceptionTag =
+            self()->getModule()->getTag(resumeThrow->tag);
+        } else {
+          // resume_throw_ref
+          contData->exception = arguments[0];
+        }
+      }
+
       self()->pushCurrContinuation(contData);
       self()->continuationStore->resuming = true;
 #if WASM_INTERPRETER_DEBUG
@@ -4709,10 +4738,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     return ret;
   }
   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 visitResumeThrow(ResumeThrow* curr) { return doResume(curr); }
   Flow visitStackSwitch(StackSwitch* curr) { return Flow(NONCONSTANT_FLOW); }
 
   void trap(const char* why) override {
@@ -4784,10 +4810,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
         // to do is just start calling this function (with the arguments we've
         // set), so resuming is done. (And throw, if resume_throw.)
         self()->continuationStore->resuming = false;
-        if (auto* tag = currContinuation->exceptionTag) {
-          throwException(WasmException{
-            self()->makeExnData(tag, currContinuation->resumeArguments)});
-        }
+        maybeThrowAfterResuming(currContinuation);
       }
     }
 

src/wasm-ir-builder.h

@@ -273,6 +273,12 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
                            Name tag,
                            const std::vector<Name>& tags,
                            const std::vector<std::optional<Index>>& labels);
+  Result<> makeResumeThrowRef(HeapType ct,
+                              const std::vector<Name>& tags,
+                              const std::vector<std::optional<Index>>& labels) {
+    // resume_throw_ref has an empty tag.
+    return makeResumeThrow(ct, Name(), tags, labels);
+  }
   Result<> makeStackSwitch(HeapType ct, Name tag);
 
   // Private functions that must be public for technical reasons.

src/wasm.h

@@ -2114,6 +2114,9 @@ class ResumeThrow : public SpecificExpression<Expression::ResumeThrowId> {
     : handlerTags(allocator), handlerBlocks(allocator), operands(allocator),
       sentTypes(allocator) {}
 
+  // If tag is set to a non-null Name, this is a resume_throw and |operands|
+  // contains the values to be set in an exception of that tag. If tag is null,
+  // this is resume_throw_ref and |operands| contains a single item, the exnref.
   Name tag;
   // See the comment on `Resume` above.
   ArenaVector<Name> handlerTags;