Typed continuations: resume instructions (WebAssembly#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.

scripts/fuzz_opt.py

@@ -315,6 +315,7 @@ def is_git_repo():
     'multi-memory-lowering-import-error.wast',
     # the fuzzer does not support typed continuations
     'typed_continuations.wast',
+    'typed_continuations_resume.wast',
     # New EH implementation is in progress
     'exception-handling.wast',
 ]

scripts/gen-s-parser.py

@@ -566,6 +566,8 @@
     # Typed function references instructions
     ("call_ref",             "makeCallRef(s, /*isReturn=*/false)"),
     ("return_call_ref",      "makeCallRef(s, /*isReturn=*/true)"),
+    # Typed continuations instructions
+    ("resume",               "makeResume(s)"),
     # GC
     ("i31.new",              "makeRefI31(s)"),  # deprecated
     ("ref.i31",              "makeRefI31(s)"),

src/gen-s-parser.inc

@@ -3004,6 +3004,9 @@ switch (buf[0]) {
           default: goto parse_error;
         }
       }
+      case 's':
+        if (op == "resume"sv) { return makeResume(s); }
+        goto parse_error;
       case 't': {
         switch (buf[3]) {
           case 'h':
@@ -8104,6 +8107,12 @@ switch (buf[0]) {
           default: goto parse_error;
         }
       }
+      case 's':
+        if (op == "resume"sv) {
+          CHECK_ERR(makeResume(ctx, pos));
+          return Ok{};
+        }
+        goto parse_error;
       case 't': {
         switch (buf[3]) {
           case 'h':

src/ir/ReFinalize.cpp

@@ -182,6 +182,8 @@ void ReFinalize::visitStringSliceIter(StringSliceIter* curr) {
   curr->finalize();
 }
 
+void ReFinalize::visitResume(Resume* curr) { curr->finalize(); }
+
 void ReFinalize::visitExport(Export* curr) { WASM_UNREACHABLE("unimp"); }
 void ReFinalize::visitGlobal(Global* curr) { WASM_UNREACHABLE("unimp"); }
 void ReFinalize::visitTable(Table* curr) { WASM_UNREACHABLE("unimp"); }

src/ir/branch-utils.h

@@ -82,6 +82,13 @@ void operateOnScopeNameUsesAndSentTypes(Expression* expr, T func) {
           func(name, tt->sentTypes[i]);
         }
       }
+    } else if (auto* r = expr->dynCast<Resume>()) {
+      for (Index i = 0; i < r->handlerTags.size(); i++) {
+        auto dest = r->handlerTags[i];
+        if (dest == name) {
+          func(name, r->sentTypes[i]);
+        }
+      }
     } else {
       assert(expr->is<Try>() || expr->is<Rethrow>()); // delegate or rethrow
     }
@@ -106,6 +113,10 @@ void operateOnScopeNameUsesAndSentValues(Expression* expr, T func) {
       // The values are supplied by throwing instructions, so we are unable to
       // know what they will be here.
       func(name, nullptr);
+    } else if (auto* res = expr->dynCast<Resume>()) {
+      // The values are supplied by suspend instructions executed while running
+      // the continuation, so we are unable to know what they will be here.
+      func(name, nullptr);
     } else {
       assert(expr->is<Try>() || expr->is<Rethrow>()); // delegate or rethrow
     }

src/ir/cost.h

@@ -726,6 +726,11 @@ struct CostAnalyzer : public OverriddenVisitor<CostAnalyzer, CostType> {
     return 8 + visit(curr->ref) + visit(curr->num);
   }
 
+  CostType visitResume(Resume* curr) {
+    // Inspired by indirect calls, but twice the cost.
+    return 12 + visit(curr->cont);
+  }
+
 private:
   CostType nullCheckCost(Expression* ref) {
     // A nullable type requires a bounds check in most VMs.

src/ir/effects.h

@@ -973,6 +973,19 @@ class EffectAnalyzer {
       // traps when ref is null.
       parent.implicitTrap = true;
     }
+
+    void visitResume(Resume* curr) {
+      // This acts as a kitchen sink effect.
+      parent.calls = true;
+
+      // resume instructions accept nullable continuation references and trap
+      // on null.
+      parent.implicitTrap = true;
+
+      if (parent.features.hasExceptionHandling() && parent.tryDepth == 0) {
+        parent.throws_ = true;
+      }
+    }
   };
 
 public:

src/ir/module-utils.cpp

@@ -341,6 +341,8 @@ struct CodeScanner
       counts.include(get->type);
     } else if (auto* set = curr->dynCast<ArraySet>()) {
       counts.note(set->ref->type);
+    } else if (auto* resume = curr->dynCast<Resume>()) {
+      counts.note(resume->contType);
     } else if (Properties::isControlFlowStructure(curr)) {
       counts.noteControlFlow(Signature(Type::none, curr->type));
     }

src/ir/possible-contents.cpp

@@ -1199,6 +1199,11 @@ struct InfoCollector
 
   void visitReturn(Return* curr) { addResult(curr->value); }
 
+  void visitResume(Resume* curr) {
+    // TODO: optimize when possible
+    addRoot(curr);
+  }
+
   void visitFunction(Function* func) {
     // Functions with a result can flow a value out from their body.
     addResult(func->body);

src/ir/subtype-exprs.h

@@ -333,6 +333,8 @@ struct SubtypingDiscoverer : public OverriddenVisitor<SubType> {
   void visitStringIterMove(StringIterMove* curr) {}
   void visitStringSliceWTF(StringSliceWTF* curr) {}
   void visitStringSliceIter(StringSliceIter* curr) {}
+
+  void visitResume(Resume* curr) { WASM_UNREACHABLE("not implemented"); }
 };
 
 } // namespace wasm

src/parser/parsers.h

@@ -164,6 +164,7 @@ Result<> makeStringIterMove(Ctx&, Index, StringIterMoveOp op);
 template<typename Ctx>
 Result<> makeStringSliceWTF(Ctx&, Index, StringSliceWTFOp op);
 template<typename Ctx> Result<> makeStringSliceIter(Ctx&, Index);
+template<typename Ctx> Result<> makeResume(Ctx&, Index);
 
 // Modules
 template<typename Ctx> MaybeResult<Index> maybeTypeidx(Ctx& ctx);
@@ -1812,6 +1813,10 @@ template<typename Ctx> Result<> makeStringSliceIter(Ctx& ctx, Index pos) {
   return ctx.makeStringSliceIter(pos);
 }
 
+template<typename Ctx> Result<> makeResume(Ctx& ctx, Index pos) {
+  return ctx.in.err("unimplemented instruction");
+}
+
 // =======
 // Modules
 // =======

src/passes/Print.cpp

@@ -316,6 +316,7 @@ struct PrintSExpression : public UnifiedExpressionVisitor<PrintSExpression> {
   void visitLoop(Loop* curr);
   void visitTry(Try* curr);
   void visitTryTable(TryTable* curr);
+  void visitResume(Resume* curr);
   void maybePrintUnreachableReplacement(Expression* curr, Type type);
   void maybePrintUnreachableOrNullReplacement(Expression* curr, Type type);
   void visitCallRef(CallRef* curr) {
@@ -2438,6 +2439,24 @@ struct PrintExpressionContents
   void visitStringSliceIter(StringSliceIter* curr) {
     printMedium(o, "stringview_iter.slice");
   }
+
+  void visitResume(Resume* curr) {
+    printMedium(o, "resume");
+
+    o << ' ';
+    printHeapType(curr->contType);
+
+    // We deliberate keep all (tag ...) clauses on the same line as the resume
+    // itself to work around a quirk in update_lit_checks.py
+    for (Index i = 0; i < curr->handlerTags.size(); i++) {
+      o << " (";
+      printMedium(o, "tag ");
+      printName(curr->handlerTags[i], o);
+      o << ' ';
+      printName(curr->handlerBlocks[i], o);
+      o << ')';
+    }
+  }
 };
 
 void PrintSExpression::setModule(Module* module) {
@@ -2786,6 +2805,23 @@ void PrintSExpression::visitTryTable(TryTable* curr) {
   controlFlowDepth--;
 }
 
+void PrintSExpression::visitResume(Resume* curr) {
+  controlFlowDepth++;
+  o << '(';
+  printExpressionContents(curr);
+
+  incIndent();
+
+  for (Index i = 0; i < curr->operands.size(); i++) {
+    printFullLine(curr->operands[i]);
+  }
+
+  printFullLine(curr->cont);
+
+  controlFlowDepth--;
+  decIndent();
+}
+
 void PrintSExpression::maybePrintUnreachableReplacement(Expression* curr,
                                                         Type type) {
   // See the parallel function

src/passes/TypeGeneralizing.cpp

@@ -874,6 +874,8 @@ struct TransferFn : OverriddenVisitor<TransferFn> {
   void visitStringIterMove(StringIterMove* curr) { WASM_UNREACHABLE("TODO"); }
   void visitStringSliceWTF(StringSliceWTF* curr) { WASM_UNREACHABLE("TODO"); }
   void visitStringSliceIter(StringSliceIter* curr) { WASM_UNREACHABLE("TODO"); }
+
+  void visitResume(Resume* curr) { WASM_UNREACHABLE("TODO"); }
 };
 
 struct TypeGeneralizing : WalkerPass<PostWalker<TypeGeneralizing>> {

src/wasm-binary.h

@@ -1296,6 +1296,10 @@ enum ASTNodes {
   StringEncodeLossyUTF8Array = 0xb6,
   StringEncodeWTF8Array = 0xb7,
   StringNewUTF8ArrayTry = 0xb8,
+
+  // typed continuation opcodes
+  Resume = 0xe3,
+
 };
 
 enum MemoryAccess {
@@ -1922,6 +1926,7 @@ class WasmBinaryReader {
   void visitCallRef(CallRef* curr);
   void visitRefAsCast(RefCast* curr, uint32_t code);
   void visitRefAs(RefAs* curr, uint8_t code);
+  void visitResume(Resume* curr);
 
   [[noreturn]] void throwError(std::string text);
 

src/wasm-builder.h

@@ -1197,6 +1197,21 @@ class Builder {
     return ret;
   }
 
+  Resume* makeResume(HeapType contType,
+                     const std::vector<Name>& handlerTags,
+                     const std::vector<Name>& handlerBlocks,
+                     const std::vector<Expression*>& operands,
+                     Expression* cont) {
+    auto* ret = wasm.allocator.alloc<Resume>();
+    ret->contType = contType;
+    ret->handlerTags.set(handlerTags);
+    ret->handlerBlocks.set(handlerBlocks);
+    ret->operands.set(operands);
+    ret->cont = cont;
+    ret->finalize(&wasm);
+    return ret;
+  }
+
   // Additional helpers
 
   Drop* makeDrop(Expression* value) {

src/wasm-delegations-fields.def

@@ -945,6 +945,18 @@ switch (DELEGATE_ID) {
     DELEGATE_END(StringSliceIter);
     break;
   }
+
+  case Expression::Id::ResumeId: {
+    DELEGATE_START(Resume);
+    DELEGATE_FIELD_TYPE_VECTOR(Resume, sentTypes);
+    DELEGATE_FIELD_CHILD(Resume, cont);
+    DELEGATE_FIELD_CHILD_VECTOR(Resume, operands);
+    DELEGATE_FIELD_SCOPE_NAME_USE_VECTOR(Resume, handlerBlocks);
+    DELEGATE_FIELD_NAME_KIND_VECTOR(Resume, handlerTags, ModuleItemKind::Tag);
+    DELEGATE_FIELD_HEAPTYPE(Resume, contType);
+    DELEGATE_END(Resume);
+    break;
+  }
 }
 
 #undef DELEGATE_ID

src/wasm-delegations.def

@@ -105,5 +105,6 @@ DELEGATE(StringIterNext);
 DELEGATE(StringIterMove);
 DELEGATE(StringSliceWTF);
 DELEGATE(StringSliceIter);
+DELEGATE(Resume);
 
 #undef DELEGATE

src/wasm-interpreter.h

@@ -2349,6 +2349,7 @@ class ConstantExpressionRunner : public ExpressionRunner<SubType> {
     }
     return ExpressionRunner<SubType>::visitRefAs(curr);
   }
+  Flow visitResume(Resume* curr) { WASM_UNREACHABLE("unimplemented"); }
 
   void trap(const char* why) override { throw NonconstantException(); }
 
@@ -3915,6 +3916,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
     multiValues.pop_back();
     return ret;
   }
+  Flow visitResume(Resume* curr) { return Flow(NONCONSTANT_FLOW); }
 
   void trap(const char* why) override { externalInterface->trap(why); }
 

src/wasm-s-parser.h

@@ -328,6 +328,7 @@ class SExpressionWasmBuilder {
   Expression* makeStringIterMove(Element& s, StringIterMoveOp op);
   Expression* makeStringSliceWTF(Element& s, StringSliceWTFOp op);
   Expression* makeStringSliceIter(Element& s);
+  Expression* makeResume(Element& s);
 
   // Helper functions
   Type parseBlockType(Element& s, Index& i);

src/wasm.h

@@ -743,6 +743,7 @@ class Expression {
     StringIterMoveId,
     StringSliceWTFId,
     StringSliceIterId,
+    ResumeId,
     NumExpressionIds
   };
   Id _id;
@@ -1995,6 +1996,32 @@ class StringSliceIter
   void finalize();
 };
 
+class Resume : public SpecificExpression<Expression::ResumeId> {
+public:
+  Resume(MixedArena& allocator)
+    : handlerTags(allocator), handlerBlocks(allocator), operands(allocator),
+      sentTypes(allocator) {}
+
+  HeapType contType;
+  ArenaVector<Name> handlerTags;
+  ArenaVector<Name> handlerBlocks;
+
+  ExpressionList operands;
+  Expression* cont;
+
+  // When 'Module*' parameter is given, we populate the 'sentTypes' array, so
+  // that the types can be accessed in other analyses without accessing the
+  // module.
+  void finalize(Module* wasm = nullptr);
+
+  // sentTypes[i] contains the type of the values that will be sent to the block
+  // handlerBlocks[i] if suspending with tag handlerTags[i]. Not part of the
+  // instruction's syntax, but stored here for subsequent use.
+  // This information is cached here in order not to query the module
+  // every time we query the sent types.
+  ArenaVector<Type> sentTypes;
+};
+
 // Globals
 
 struct Named {