Typed continuations: resume instructions (#6083)

frank-emrich · web-flow · commit b4dee3dc0583 · 2024-01-11T13:22:43.000-08:00
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.
diff --git a/scripts/fuzz_opt.py b/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',
 ]
diff --git a/scripts/gen-s-parser.py b/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)"),
diff --git a/src/gen-s-parser.inc b/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':
diff --git a/src/ir/ReFinalize.cpp b/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"); }
diff --git a/src/ir/branch-utils.h b/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
     }
diff --git a/src/ir/cost.h b/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.
diff --git a/src/ir/effects.h b/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:
diff --git a/src/ir/module-utils.cpp b/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));
     }
diff --git a/src/ir/possible-contents.cpp b/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);
diff --git a/src/ir/subtype-exprs.h b/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
diff --git a/src/parser/parsers.h b/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
 // =======
diff --git a/src/passes/Print.cpp b/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
diff --git a/src/passes/TypeGeneralizing.cpp b/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>> {
diff --git a/src/wasm-binary.h b/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);
 
diff --git a/src/wasm-builder.h b/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) {
diff --git a/src/wasm-delegations-fields.def b/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
diff --git a/src/wasm-delegations.def b/src/wasm-delegations.def
@@ -105,5 +105,6 @@ DELEGATE(StringIterNext);
 DELEGATE(StringIterMove);
 DELEGATE(StringSliceWTF);
 DELEGATE(StringSliceIter);
+DELEGATE(Resume);
 
 #undef DELEGATE
diff --git a/src/wasm-interpreter.h b/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); }
 
diff --git a/src/wasm-s-parser.h b/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);
diff --git a/src/wasm.h b/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 {
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
diff --git a/src/wasm2js.h b/src/wasm2js.h
diff --git a/test/lit/basic/typed_continuations_resume.wast b/test/lit/basic/typed_continuations_resume.wast

Original file line number	Diff line number	Diff line change
`@@ -315,6 +315,7 @@ def is_git_repo():`
`315`	`315`	`'multi-memory-lowering-import-error.wast',`
`316`	`316`	`# the fuzzer does not support typed continuations`
`317`	`317`	`'typed_continuations.wast',`
	`318`	`+ 'typed_continuations_resume.wast',`
`318`	`319`	`# New EH implementation is in progress`
`319`	`320`	`'exception-handling.wast',`
`320`	`321`	`]`
Original file line number	Diff line number	Diff line change
`@@ -3004,6 +3004,9 @@ switch (buf[0]) {`
`3004`	`3004`	`default: goto parse_error;`
`3005`	`3005`	`}`
`3006`	`3006`	`}`
	`3007`	`+ case 's':`
	`3008`	`+ if (op == "resume"sv) { return makeResume(s); }`
	`3009`	`+ goto parse_error;`
`3007`	`3010`	`case 't': {`
`3008`	`3011`	`switch (buf[3]) {`
`3009`	`3012`	`case 'h':`
`@@ -8104,6 +8107,12 @@ switch (buf[0]) {`
`8104`	`8107`	`default: goto parse_error;`
`8105`	`8108`	`}`
`8106`	`8109`	`}`
	`8110`	`+ case 's':`
	`8111`	`+ if (op == "resume"sv) {`
	`8112`	`+ CHECK_ERR(makeResume(ctx, pos));`
	`8113`	`+ return Ok{};`
	`8114`	`+ }`
	`8115`	`+ goto parse_error;`
`8107`	`8116`	`case 't': {`
`8108`	`8117`	`switch (buf[3]) {`
`8109`	`8118`	`case 'h':`
Original file line number	Diff line number	Diff line change
`@@ -182,6 +182,8 @@ void ReFinalize::visitStringSliceIter(StringSliceIter* curr) {`
`182`	`182`	`curr->finalize();`
`183`	`183`	`}`
`184`	`184`
	`185`	`+void ReFinalize::visitResume(Resume* curr) { curr->finalize(); }`
	`186`	`+`
`185`	`187`	`void ReFinalize::visitExport(Export* curr) { WASM_UNREACHABLE("unimp"); }`
`186`	`188`	`void ReFinalize::visitGlobal(Global* curr) { WASM_UNREACHABLE("unimp"); }`
`187`	`189`	`void ReFinalize::visitTable(Table* curr) { WASM_UNREACHABLE("unimp"); }`
Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,13 @@ void operateOnScopeNameUsesAndSentTypes(Expression* expr, T func) {`
`82`	`82`	`func(name, tt->sentTypes[i]);`
`83`	`83`	`}`
`84`	`84`	`}`
	`85`	`+ } else if (auto* r = expr->dynCast<Resume>()) {`
	`86`	`+ for (Index i = 0; i < r->handlerTags.size(); i++) {`
	`87`	`+ auto dest = r->handlerTags[i];`
	`88`	`+ if (dest == name) {`
	`89`	`+ func(name, r->sentTypes[i]);`
	`90`	`+ }`
	`91`	`+ }`
`85`	`92`	`} else {`
`86`	`93`	`assert(expr->is<Try>() \|\| expr->is<Rethrow>()); // delegate or rethrow`
`87`	`94`	`}`
`@@ -106,6 +113,10 @@ void operateOnScopeNameUsesAndSentValues(Expression* expr, T func) {`
`106`	`113`	`// The values are supplied by throwing instructions, so we are unable to`
`107`	`114`	`// know what they will be here.`
`108`	`115`	`func(name, nullptr);`
	`116`	`+ } else if (auto* res = expr->dynCast<Resume>()) {`
	`117`	`+ // The values are supplied by suspend instructions executed while running`
	`118`	`+ // the continuation, so we are unable to know what they will be here.`
	`119`	`+ func(name, nullptr);`
`109`	`120`	`} else {`
`110`	`121`	`assert(expr->is<Try>() \|\| expr->is<Rethrow>()); // delegate or rethrow`
`111`	`122`	`}`
Original file line number	Diff line number	Diff line change
`@@ -341,6 +341,8 @@ struct CodeScanner`
`341`	`341`	`counts.include(get->type);`
`342`	`342`	`} else if (auto* set = curr->dynCast<ArraySet>()) {`
`343`	`343`	`counts.note(set->ref->type);`
	`344`	`+ } else if (auto* resume = curr->dynCast<Resume>()) {`
	`345`	`+ counts.note(resume->contType);`
`344`	`346`	`} else if (Properties::isControlFlowStructure(curr)) {`
`345`	`347`	`counts.noteControlFlow(Signature(Type::none, curr->type));`
`346`	`348`	`}`