Typed continuations: resume instructions (#6083)

This PR is part of a series that adds basic support for the [typed continuations proposal](https://github.com/wasmfx/specfx).

This particular PR adds support for the `resume` instruction. The most notable missing feature is validation, which is not implemented, yet.

5 files changed, 163 insertions, 0 deletions

diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index bf14efeb1..06c6f3648 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -4046,6 +4046,10 @@ BinaryConsts::ASTNodes WasmBinaryReader::readExpression(Expression*& curr) {
       visitCallRef(call);
       break;
     }
+    case BinaryConsts::Resume: {
+      visitResume((curr = allocator.alloc<Resume>())->cast<Resume>());
+      break;
+    }
     case BinaryConsts::AtomicPrefix: {
       code = static_cast<uint8_t>(getU32LEB());
       if (maybeVisitLoad(curr, code, /*isAtomic=*/true)) {
@@ -7758,6 +7762,52 @@ void WasmBinaryReader::visitRefAs(RefAs* curr, uint8_t code) {
   curr->finalize();
 }
 
+void WasmBinaryReader::visitResume(Resume* curr) {
+  BYN_TRACE("zz node: Resume\n");
+
+  auto contTypeIndex = getU32LEB();
+  curr->contType = getTypeByIndex(contTypeIndex);
+  if (!curr->contType.isContinuation()) {
+    throwError("non-continuation type in resume instruction " +
+               curr->contType.toString());
+  }
+
+  auto numHandlers = getU32LEB();
+
+  // We *must* bring the handlerTags vector to an appropriate size to ensure
+  // that we do not invalidate the pointers we add to tagRefs. They need to stay
+  // valid until processNames ran.
+  curr->handlerTags.resize(numHandlers);
+  curr->handlerBlocks.resize(numHandlers);
+
+  BYN_TRACE("handler num: " << numHandlers << std::endl);
+  for (size_t i = 0; i < numHandlers; i++) {
+    BYN_TRACE("read one tag handler pair \n");
+    auto tagIndex = getU32LEB();
+    auto tag = getTagName(tagIndex);
+
+    auto handlerIndex = getU32LEB();
+    auto handler = getBreakTarget(handlerIndex).name;
+
+    curr->handlerTags[i] = tag;
+    curr->handlerBlocks[i] = handler;
+
+    // We don't know the final name yet
+    tagRefs[tagIndex].push_back(&curr->handlerTags[i]);
+  }
+
+  curr->cont = popNonVoidExpression();
+
+  auto numArgs =
+    curr->contType.getContinuation().type.getSignature().params.size();
+  curr->operands.resize(numArgs);
+  for (size_t i = 0; i < numArgs; i++) {
+    curr->operands[numArgs - i - 1] = popNonVoidExpression();
+  }
+
+  curr->finalize(&wasm);
+}
+
 void WasmBinaryReader::throwError(std::string text) {
   throw ParseException(text, 0, pos);
 }
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index 9dbfc88c2..cc8fe2273 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -2988,6 +2988,38 @@ Expression* SExpressionWasmBuilder::makeCallRef(Element& s, bool isReturn) {
     target, operands, sigType.getSignature().results, isReturn);
 }
 
+Expression* SExpressionWasmBuilder::makeResume(Element& s) {
+  auto ret = allocator.alloc<Resume>();
+
+  ret->contType = parseHeapType(*s[1]);
+  if (!ret->contType.isContinuation()) {
+    throw ParseException("expected continuation type", s[1]->line, s[1]->col);
+  }
+
+  Index i = 2;
+  while (i < s.size() && elementStartsWith(*s[i], "tag")) {
+    Element& inner = *s[i++];
+    if (inner.size() < 3) {
+      throw ParseException("invalid tag block", inner.line, inner.col);
+    }
+    Name tag = getTagName(*inner[1]);
+    if (!wasm.getTagOrNull(tag)) {
+      throw ParseException("bad tag name", inner[1]->line, inner[1]->col);
+    }
+    ret->handlerTags.push_back(tag);
+    ret->handlerBlocks.push_back(getLabel(*inner[2]));
+  }
+
+  while (i < s.size() - 1) {
+    ret->operands.push_back(parseExpression(s[i++]));
+  }
+
+  ret->cont = parseExpression(s[i]);
+
+  ret->finalize(&wasm);
+  return ret;
+}
+
 Expression* SExpressionWasmBuilder::makeRefI31(Element& s) {
   auto ret = allocator.alloc<RefI31>();
   ret->value = parseExpression(s[1]);
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index 3b14b3c35..6a2fd8468 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -2488,6 +2488,18 @@ void BinaryInstWriter::visitStringSliceIter(StringSliceIter* curr) {
     << U32LEB(BinaryConsts::StringViewIterSlice);
 }
 
+void BinaryInstWriter::visitResume(Resume* curr) {
+  o << int8_t(BinaryConsts::Resume);
+  parent.writeIndexedHeapType(curr->contType);
+
+  size_t handlerNum = curr->handlerTags.size();
+  o << U32LEB(handlerNum);
+  for (size_t i = 0; i < handlerNum; i++) {
+    o << U32LEB(parent.getTagIndex(curr->handlerTags[i]))
+      << U32LEB(getBreakIndex(curr->handlerBlocks[i]));
+  }
+}
+
 void BinaryInstWriter::emitScopeEnd(Expression* curr) {
   assert(!breakStack.empty());
   breakStack.pop_back();
diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
index 7ea8b0bdb..9930010a0 100644
--- a/src/wasm/wasm-validator.cpp
+++ b/src/wasm/wasm-validator.cpp
@@ -487,6 +487,7 @@ public:
   void visitStringIterMove(StringIterMove* curr);
   void visitStringSliceWTF(StringSliceWTF* curr);
   void visitStringSliceIter(StringSliceIter* curr);
+  void visitResume(Resume* curr);
 
   void visitFunction(Function* curr);
 
@@ -3285,6 +3286,24 @@ void FunctionValidator::visitStringSliceIter(StringSliceIter* curr) {
                "string operations require reference-types [--enable-strings]");
 }
 
+void FunctionValidator::visitResume(Resume* curr) {
+  // TODO implement actual type-checking
+  shouldBeTrue(
+    !getModule() || getModule()->features.hasTypedContinuations(),
+    curr,
+    "resume requires typed-continuatons [--enable-typed-continuations]");
+
+  shouldBeTrue(
+    curr->sentTypes.size() == curr->handlerBlocks.size(),
+    curr,
+    "sentTypes cache in Resume instruction has not been initialised");
+
+  shouldBeTrue((curr->contType.isContinuation() &&
+                curr->contType.getContinuation().type.isSignature()),
+               curr,
+               "invalid type in Resume expression");
+}
+
 void FunctionValidator::visitFunction(Function* curr) {
   if (curr->getResults().isTuple()) {
     shouldBeTrue(getModule()->features.hasMultivalue(),
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index ad44acda4..06e9de148 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -17,6 +17,7 @@
 #include "wasm.h"
 #include "ir/branch-utils.h"
 #include "wasm-traversal.h"
+#include "wasm-type.h"
 
 namespace wasm {
 
@@ -1350,6 +1351,55 @@ void StringSliceIter::finalize() {
   }
 }
 
+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) {
+  const Signature& contSig =
+    this->contType.getContinuation().type.getSignature();
+  type = contSig.results;
+
+  populateResumeSentTypes(this, wasm);
+}
+
 size_t Function::getNumParams() { return getParams().size(); }
 
 size_t Function::getNumVars() { return vars.size(); }