Stack switching proposal support

src/parser/parsers.h

@@ -2454,7 +2454,8 @@ makeSuspend(Ctx& ctx, Index pos, const std::vector<Annotation>& annotations) {
 
 // resumetable ::= ('(' 'on' tagidx labelidx | 'on' tagidx switch ')')*
 template<typename Ctx>
-Result<> makeResumeTable(Ctx& ctx, typename Ctx::OnClauseListT& resumetable) {
+Result<typename Ctx::OnClauseListT> makeResumeTable(Ctx& ctx) {
+  auto resumetable = ctx.makeOnClauseList();
   while (ctx.in.takeSExprStart("on"sv)) {
     auto tag = tagidx(ctx);
     CHECK_ERR(tag);
@@ -2469,7 +2470,7 @@ Result<> makeResumeTable(Ctx& ctx, typename Ctx::OnClauseListT& resumetable) {
       return ctx.in.err("expected ')' at end of handler clause");
     }
   }
-  return Ok{};
+  return resumetable;
 }
 
 // resume ::= 'resume' typeidx resumetable
@@ -2479,11 +2480,10 @@ makeResume(Ctx& ctx, Index pos, const std::vector<Annotation>& annotations) {
   auto type = typeidx(ctx);
   CHECK_ERR(type);
 
-  auto resumetable = ctx.makeOnClauseList();
-  auto ans = makeResumeTable(ctx, resumetable);
-  CHECK_ERR(ans);
+  auto resumetable = makeResumeTable(ctx);
+  CHECK_ERR(resumetable);
 
-  return ctx.makeResume(pos, annotations, *type, resumetable);
+  return ctx.makeResume(pos, annotations, *type, *resumetable);
 }
 
 // resume_throw ::= 'resume_throw' typeidx tagidx ('(' 'on' tagidx labelidx |
@@ -2497,11 +2497,10 @@ Result<> makeResumeThrow(Ctx& ctx,
   auto exnTag = tagidx(ctx);
   CHECK_ERR(exnTag);
 
-  auto resumetable = ctx.makeOnClauseList();
-  auto ans = makeResumeTable(ctx, resumetable);
-  CHECK_ERR(ans);
+  auto resumetable = makeResumeTable(ctx);
+  CHECK_ERR(resumetable);
 
-  return ctx.makeResumeThrow(pos, annotations, *type, *exnTag, resumetable);
+  return ctx.makeResumeThrow(pos, annotations, *type, *exnTag, *resumetable);
 }
 
 // switch ::= 'switch' typeidx tagidx

src/passes/Print.cpp

@@ -380,6 +380,39 @@ struct PrintSExpression : public UnifiedExpressionVisitor<PrintSExpression> {
       visitExpression(curr);
     }
   }
+  void visitContNew(ContNew* curr) {
+    if (!maybePrintUnreachableReplacement(curr, curr->type)) {
+      visitExpression(curr);
+    }
+  }
+  void visitContBind(ContBind* curr) {
+    if (!maybePrintUnreachableOrNullReplacement(
+          curr, Type(curr->sourceType, Nullable)) ||
+        !maybePrintUnreachableOrNullReplacement(curr, curr->type)) {
+      visitExpression(curr);
+    }
+  }
+  void visitResume(Resume* curr) {
+    if (!maybePrintUnreachableOrNullReplacement(
+          curr, Type(curr->contType, Nullable)) ||
+        !maybePrintUnreachableOrNullReplacement(curr, curr->type)) {
+      visitExpression(curr);
+    }
+  }
+  void visitResumeThrow(ResumeThrow* curr) {
+    if (!maybePrintUnreachableOrNullReplacement(
+          curr, Type(curr->contType, Nullable)) ||
+        !maybePrintUnreachableOrNullReplacement(curr, curr->type)) {
+      visitExpression(curr);
+    }
+  }
+  void visitStackSwitch(StackSwitch* curr) {
+    if (!maybePrintUnreachableOrNullReplacement(
+          curr, Type(curr->contType, Nullable)) ||
+        !maybePrintUnreachableOrNullReplacement(curr, curr->type)) {
+      visitExpression(curr);
+    }
+  }
 
   // Module-level visitors
   void handleSignature(HeapType curr, Name name = Name());

src/wasm-builder.h

@@ -1194,6 +1194,51 @@ class Builder {
     ret->finalize(&wasm);
     return ret;
   }
+  template<typename ResumeOrResumeThrow>
+  std::vector<Type> computeResumeTableSentTypes(ResumeOrResumeThrow* curr) {
+    static_assert(std::is_base_of<Resume, ResumeOrResumeThrow>::value ||
+                  std::is_base_of<ResumeThrow, ResumeOrResumeThrow>::value);
+    assert(curr->handlerBlocks.size() == curr->handlerTags.size());
+
+    // Let $tag be a tag with type [tgp*] -> [tgr*]. Let $ct be a continuation
+    // type (cont $ft), where $ft is [ctp*] -> [ctr*]. Then an instruction
+    // (resume $ct ... (tag $tag $block) ... ) causes $block to receive values
+    // of the following types when suspending to $tag: tgp* (ref $ct') where ct'
+    // = (cont $ft') and ft' = [tgr*] -> [ctr*].
+    //
+    std::vector<Type> sentTypes;
+    sentTypes.reserve(curr->handlerTags.size());
+    auto contSig = curr->contType.getContinuation().type.getSignature();
+    auto& ctrs = contSig.params;
+    for (Index i = 0; i < curr->handlerTags.size(); i++) {
+      if (curr->handlerBlocks[i].isNull()) {
+        sentTypes[i] = Type::none;
+        continue;
+      }
+
+      auto& tag = curr->handlerTags[i];
+      auto& tagSig = wasm.getTag(tag)->sig;
+
+      auto& tgps = tagSig.params;
+      auto& tgrs = tagSig.results;
+
+      HeapType ftPrime{Signature(tgrs, ctrs)};
+      HeapType ctPrime{Continuation(ftPrime)};
+      Type ctPrimeRef(ctPrime, Nullability::NonNullable);
+
+      if (tgps.size() > 0) {
+        TypeList sentValueTypes;
+        sentValueTypes.reserve(tgps.size() + 1);
+
+        sentValueTypes.insert(sentValueTypes.begin(), tgps.begin(), tgps.end());
+        sentValueTypes.push_back(ctPrimeRef);
+        sentTypes[i] = Type(sentValueTypes);
+      } else {
+        sentTypes[i] = ctPrimeRef;
+      }
+    }
+    return sentTypes;
+  }
   Resume* makeResume(HeapType contType,
                      const std::vector<Name>& handlerTags,
                      const std::vector<Name>& handlerBlocks,
@@ -1203,6 +1248,7 @@ class Builder {
     ret->contType = contType;
     ret->handlerTags.set(handlerTags);
     ret->handlerBlocks.set(handlerBlocks);
+    ret->sentTypes.set(computeResumeTableSentTypes<Resume>(ret));
     ret->operands.set(operands);
     ret->cont = cont;
     ret->finalize(&wasm);
@@ -1219,6 +1265,7 @@ class Builder {
     ret->tag = tag;
     ret->handlerTags.set(handlerTags);
     ret->handlerBlocks.set(handlerBlocks);
+    ret->sentTypes.set(computeResumeTableSentTypes<ResumeThrow>(ret));
     ret->operands.set(operands);
     ret->cont = cont;
     ret->finalize(&wasm);

src/wasm/wasm-ir-builder.cpp

@@ -1953,8 +1953,9 @@ IRBuilder::makeResume(HeapType ct,
     return Err{"expected continuation type"};
   }
   Resume curr(wasm.allocator);
+  auto contSig = ct.getContinuation().type.getSignature();
   curr.contType = ct;
-  curr.operands.resize(ct.getContinuation().type.getSignature().params.size());
+  curr.operands.resize(contSig.params.size());
   CHECK_ERR(visitResume(&curr));
 
   std::vector<Name> labelNames;
@@ -1979,6 +1980,9 @@ IRBuilder::makeResumeThrow(HeapType ct,
                            Name tag,
                            const std::vector<Name>& tags,
                            const std::vector<std::optional<Index>>& labels) {
+  if (tags.size() != labels.size()) {
+    return Err{"the sizes of tags and labels must be equal"};
+  }
   if (!ct.isContinuation()) {
     return Err{"expected continuation type"};
   }

src/wasm/wasm-validator.cpp

@@ -3485,10 +3485,12 @@ void FunctionValidator::visitContNew(ContNew* curr) {
                curr,
                "cont.new requires stack-switching [--enable-stack-switching]");
 
-  shouldBeTrue((curr->type.isContinuation() &&
-                curr->type.getHeapType().getContinuation().type.isSignature()),
-               curr,
-               "invalid type in ContNew expression");
+  shouldBeTrue(
+    (curr->type.isContinuation() &&
+     curr->type.getHeapType().getContinuation().type.isSignature()) ||
+      curr->type == Type::unreachable,
+    curr,
+    "cont.new must be annotated with a continuation type");
 }
 
 void FunctionValidator::visitContBind(ContBind* curr) {
@@ -3497,15 +3499,19 @@ void FunctionValidator::visitContBind(ContBind* curr) {
                curr,
                "cont.bind requires stack-switching [--enable-stack-switching]");
 
-  shouldBeTrue((curr->sourceType.isContinuation() &&
-                curr->sourceType.getContinuation().type.isSignature()),
-               curr,
-               "invalid first type in ContBind expression");
+  shouldBeTrue(
+    (curr->sourceType.isContinuation() &&
+     curr->sourceType.getContinuation().type.isSignature()) ||
+      Type(curr->sourceType, Nullable) == Type::unreachable,
+    curr,
+    "the first type annotation on cont.bind must be a continuation type");
 
-  shouldBeTrue((curr->type.isContinuation() &&
-                curr->type.getHeapType().getContinuation().type.isSignature()),
-               curr,
-               "invalid second type in ContBind expression");
+  shouldBeTrue(
+    (curr->type.isContinuation() &&
+     curr->type.getHeapType().getContinuation().type.isSignature()) ||
+      curr->type == Type::unreachable,
+    curr,
+    "the second type annotation on cont.bind must be a continuation type");
 }
 
 void FunctionValidator::visitSuspend(Suspend* curr) {
@@ -3521,10 +3527,16 @@ void FunctionValidator::visitResume(Resume* curr) {
                curr,
                "resume requires stack-switching [--enable-stack-switching]");
 
+  shouldBeTrue(
+    curr->sentTypes.size() == curr->handlerBlocks.size(),
+    curr,
+    "sentTypes cache in resume instruction has not been initialized");
+
   shouldBeTrue((curr->contType.isContinuation() &&
-                curr->contType.getContinuation().type.isSignature()),
+                curr->contType.getContinuation().type.isSignature()) ||
+                 curr->type == Type::unreachable,
                curr,
-               "invalid type in Resume expression");
+               "resume must be annotated with a continuation type");
 }
 
 void FunctionValidator::visitResumeThrow(ResumeThrow* curr) {
@@ -3536,10 +3548,21 @@ void FunctionValidator::visitResumeThrow(ResumeThrow* curr) {
     "resume_throw requires exception handling [--enable-exception-handling] "
     "and stack-switching [--enable-stack-switching]");
 
+  shouldBeTrue(
+    curr->sentTypes.size() == curr->handlerBlocks.size(),
+    curr,
+    "sentTypes cache in resume_throw instruction has not been initialized");
+
   shouldBeTrue((curr->contType.isContinuation() &&
-                curr->contType.getContinuation().type.isSignature()),
+                curr->contType.getContinuation().type.isSignature()) ||
+                 curr->type == Type::unreachable,
                curr,
-               "invalid type in ResumeThrow expression");
+               "resume_throw must be annotated with a continuation type");
+
+  auto* tag = getModule()->getTagOrNull(curr->tag);
+  if (!shouldBeTrue(!!tag, curr, "resume_throw must be annotated with a tag")) {
+    return;
+  }
 }
 
 void FunctionValidator::visitStackSwitch(StackSwitch* curr) {
@@ -3549,9 +3572,15 @@ void FunctionValidator::visitStackSwitch(StackSwitch* curr) {
                "switch requires stack-switching [--enable-stack-switching]");
 
   shouldBeTrue((curr->contType.isContinuation() &&
-                curr->contType.getContinuation().type.isSignature()),
+                curr->contType.getContinuation().type.isSignature()) ||
+                 curr->type == Type::unreachable,
                curr,
-               "invalid type in Switch expression");
+               "switch must be annotated with a continuation type");
+
+  auto* tag = getModule()->getTagOrNull(curr->tag);
+  if (!shouldBeTrue(!!tag, curr, "switch must be annotated with a tag")) {
+    return;
+  }
 }
 
 void FunctionValidator::visitFunction(Function* curr) {

src/wasm/wasm.cpp

@@ -1387,6 +1387,47 @@ void Suspend::finalize(Module* wasm) {
   }
 }
 
+static void populateResumeSentTypes(Resume* curr, Module* wasm) {
+  if (!wasm) {
+    return;
+  }
+
+  const Signature& contSig =
+    curr->contType.getContinuation().type.getSignature();
+
+  // Let $tag be a tag with type [tgp*] -> [tgr*]. Let $ct be a continuation
+  // type (cont $ft), where $ft is [ctp*] -> [ctr*]. Then an instruction (resume
+  // $ct ... (tag $tag $block) ... ) causes $block to receive values of the
+  // following types when suspending to $tag: tgp* (ref $ct') where ct' = (cont
+  // $ft') and ft' = [tgr*] -> [ctr*].
+  //
+  auto& ctrs = contSig.results;
+  curr->sentTypes.clear();
+  curr->sentTypes.resize(curr->handlerTags.size());
+  for (Index i = 0; i < curr->handlerTags.size(); i++) {
+    auto& tag = curr->handlerTags[i];
+    auto& tagSig = wasm->getTag(tag)->sig;
+
+    auto& tgps = tagSig.params;
+    auto& tgrs = tagSig.results;
+
+    HeapType ftPrime{Signature(tgrs, ctrs)};
+    HeapType ctPrime{Continuation(ftPrime)};
+    Type ctPrimeRef(ctPrime, Nullability::NonNullable);
+
+    if (tgps.size() > 0) {
+      TypeList sentValueTypes;
+      sentValueTypes.reserve(tgps.size() + 1);
+
+      sentValueTypes.insert(sentValueTypes.begin(), tgps.begin(), tgps.end());
+      sentValueTypes.push_back(ctPrimeRef);
+      curr->sentTypes[i] = Type(sentValueTypes);
+    } else {
+      curr->sentTypes[i] = ctPrimeRef;
+    }
+  }
+}
+
 void Resume::finalize(Module* wasm) {
   if (cont->type == Type::unreachable) {
     type = Type::unreachable;
@@ -1395,6 +1436,49 @@ void Resume::finalize(Module* wasm) {
       this->contType.getContinuation().type.getSignature();
     type = contSig.results;
   }
+
+  populateResumeSentTypes(this, wasm);
+}
+
+static void populateResumeThrowSentTypes(ResumeThrow* curr, Module* wasm) {
+  if (!wasm) {
+    return;
+  }
+
+  const Signature& contSig =
+    curr->contType.getContinuation().type.getSignature();
+
+  // Let $tag be a tag with type [tgp*] -> [tgr*]. Let $ct be a continuation
+  // type (cont $ft), where $ft is [ctp*] -> [ctr*]. Then an instruction (resume
+  // $ct ... (tag $tag $block) ... ) causes $block to receive values of the
+  // following types when suspending to $tag: tgp* (ref $ct') where ct' = (cont
+  // $ft') and ft' = [tgr*] -> [ctr*].
+  //
+  auto& ctrs = contSig.results;
+  curr->sentTypes.clear();
+  curr->sentTypes.resize(curr->handlerTags.size());
+  for (Index i = 0; i < curr->handlerTags.size(); i++) {
+    auto& tag = curr->handlerTags[i];
+    auto& tagSig = wasm->getTag(tag)->sig;
+
+    auto& tgps = tagSig.params;
+    auto& tgrs = tagSig.results;
+
+    HeapType ftPrime{Signature(tgrs, ctrs)};
+    HeapType ctPrime{Continuation(ftPrime)};
+    Type ctPrimeRef(ctPrime, Nullability::NonNullable);
+
+    if (tgps.size() > 0) {
+      TypeList sentValueTypes;
+      sentValueTypes.reserve(tgps.size() + 1);
+
+      sentValueTypes.insert(sentValueTypes.begin(), tgps.begin(), tgps.end());
+      sentValueTypes.push_back(ctPrimeRef);
+      curr->sentTypes[i] = Type(sentValueTypes);
+    } else {
+      curr->sentTypes[i] = ctPrimeRef;
+    }
+  }
 }
 
 void ResumeThrow::finalize(Module* wasm) {
@@ -1405,6 +1489,8 @@ void ResumeThrow::finalize(Module* wasm) {
       this->contType.getContinuation().type.getSignature();
     type = contSig.results;
   }
+
+  populateResumeThrowSentTypes(this, wasm);
 }
 
 void StackSwitch::finalize(Module* wasm) {