Typed continuations: suspend instructions

scripts/fuzz_opt.py

@@ -328,6 +328,7 @@ def is_git_repo():
     'typed_continuations_resume.wast',
     'typed_continuations_contnew.wast',
     'typed_continuations_contbind.wast',
+    'typed_continuations_suspend.wast',
     # New EH implementation is in progress
     'exception-handling.wast',
     'translate-eh-old-to-new.wast',

scripts/gen-s-parser.py

@@ -570,6 +570,7 @@
     ("cont.new",             "makeContNew(s)"),
     ("cont.bind",            "makeContBind(s)"),
     ("resume",               "makeResume(s)"),
+    ("suspend",              "makeSuspend(s)"),
     # GC
     ("i31.new",              "makeRefI31(s)"),  # deprecated
     ("ref.i31",              "makeRefI31(s)"),

src/gen-s-parser.inc

@@ -3369,6 +3369,9 @@ switch (buf[0]) {
           default: goto parse_error;
         }
       }
+      case 'u':
+        if (op == "suspend"sv) { return makeSuspend(s); }
+        goto parse_error;
       default: goto parse_error;
     }
   }
@@ -8653,6 +8656,12 @@ switch (buf[0]) {
           default: goto parse_error;
         }
       }
+      case 'u':
+        if (op == "suspend"sv) {
+          CHECK_ERR(makeSuspend(ctx, pos, annotations));
+          return Ok{};
+        }
+        goto parse_error;
       default: goto parse_error;
     }
   }

src/ir/ReFinalize.cpp

@@ -185,6 +185,7 @@ void ReFinalize::visitStringSliceIter(StringSliceIter* curr) {
 void ReFinalize::visitContNew(ContNew* curr) { curr->finalize(); }
 void ReFinalize::visitContBind(ContBind* curr) { curr->finalize(); }
 void ReFinalize::visitResume(Resume* curr) { curr->finalize(); }
+void ReFinalize::visitSuspend(Suspend* curr) { curr->finalize(getModule()); }
 
 void ReFinalize::visitExport(Export* curr) { WASM_UNREACHABLE("unimp"); }
 void ReFinalize::visitGlobal(Global* curr) { WASM_UNREACHABLE("unimp"); }

src/ir/cost.h

@@ -742,7 +742,18 @@ struct CostAnalyzer : public OverriddenVisitor<CostAnalyzer, CostType> {
   }
   CostType visitResume(Resume* curr) {
     // Inspired by indirect calls, but twice the cost.
-    return 12 + visit(curr->cont);
+    CostType ret = 12 + visit(curr->cont);
+    for (auto* arg : curr->operands) {
+      ret += visit(arg);
+    }
+    return ret;
+  }
+  CostType visitSuspend(Suspend* curr) {
+    CostType ret = 12;
+    for (auto* arg : curr->operands) {
+      ret += visit(arg);
+    }
+    return ret;
   }
 
 private:

src/ir/effects.h

@@ -994,6 +994,14 @@ class EffectAnalyzer {
         parent.throws_ = true;
       }
     }
+    void visitSuspend(Suspend* curr) {
+      // Similar to resume/call: Suspending means that we execute arbitrary
+      // other code before we may resume here.
+      parent.calls = true;
+      if (parent.features.hasExceptionHandling() && parent.tryDepth == 0) {
+        parent.throws_ = true;
+      }
+    }
   };
 
 public:

src/ir/possible-contents.cpp

@@ -1212,6 +1212,10 @@ struct InfoCollector
     // TODO: optimize when possible
     addRoot(curr);
   }
+  void visitSuspend(Suspend* curr) {
+    // TODO: optimize when possible
+    addRoot(curr);
+  }
 
   void visitFunction(Function* func) {
     // Functions with a result can flow a value out from their body.

src/ir/subtype-exprs.h

@@ -394,6 +394,7 @@ struct SubtypingDiscoverer : public OverriddenVisitor<SubType> {
   void visitContBind(ContBind* curr) { WASM_UNREACHABLE("not implemented"); }
   void visitContNew(ContNew* curr) { WASM_UNREACHABLE("not implemented"); }
   void visitResume(Resume* curr) { WASM_UNREACHABLE("not implemented"); }
+  void visitSuspend(Suspend* curr) { WASM_UNREACHABLE("not implemented"); }
 };
 
 } // namespace wasm

src/parser/contexts.h

@@ -824,6 +824,9 @@ struct NullInstrParserCtx {
                       const TagLabelListT&) {
     return Ok{};
   }
+  Result<> makeSuspend(Index, const std::vector<Annotation>&, TagIdxT) {
+    return Ok{};
+  }
 };
 
 struct NullCtx : NullTypeParserCtx, NullInstrParserCtx {
@@ -2594,6 +2597,11 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
     }
     return withLoc(pos, irBuilder.makeResume(type, tags, labels));
   }
+
+  Result<>
+  makeSuspend(Index pos, const std::vector<Annotation>& annotations, Name tag) {
+    return withLoc(pos, irBuilder.makeSuspend(tag));
+  }
 };
 
 } // namespace wasm::WATParser

src/parser/parsers.h

@@ -309,6 +309,8 @@ template<typename Ctx>
 Result<> makeContNew(Ctx*, Index, const std::vector<Annotation>&);
 template<typename Ctx>
 Result<> makeResume(Ctx&, Index, const std::vector<Annotation>&);
+template<typename Ctx>
+Result<> makeSuspend(Ctx&, Index, const std::vector<Annotation>&);
 
 // Modules
 template<typename Ctx> MaybeResult<Index> maybeTypeidx(Ctx& ctx);
@@ -2498,6 +2500,15 @@ makeResume(Ctx& ctx, Index pos, const std::vector<Annotation>& annotations) {
   return ctx.makeResume(pos, annotations, *type, tagLabels);
 }
 
+template<typename Ctx>
+Result<>
+makeSuspend(Ctx& ctx, Index pos, const std::vector<Annotation>& annotations) {
+  auto tag = tagidx(ctx);
+  CHECK_ERR(tag);
+
+  return ctx.makeSuspend(pos, annotations, *tag);
+}
+
 // =======
 // Modules
 // =======

src/passes/Print.cpp

@@ -2382,6 +2382,11 @@ struct PrintExpressionContents
       o << ')';
     }
   }
+
+  void visitSuspend(Suspend* curr) {
+    printMedium(o, "suspend ");
+    curr->tag.print(o);
+  }
 };
 
 void PrintSExpression::setModule(Module* module) {

src/passes/TypeGeneralizing.cpp

@@ -878,6 +878,7 @@ struct TransferFn : OverriddenVisitor<TransferFn> {
   void visitContBind(ContBind* curr) { WASM_UNREACHABLE("TODO"); }
   void visitContNew(ContNew* curr) { WASM_UNREACHABLE("TODO"); }
   void visitResume(Resume* curr) { WASM_UNREACHABLE("TODO"); }
+  void visitSuspend(Suspend* curr) { WASM_UNREACHABLE("TODO"); }
 };
 
 struct TypeGeneralizing : WalkerPass<PostWalker<TypeGeneralizing>> {

src/wasm-binary.h

@@ -1179,6 +1179,7 @@ enum ASTNodes {
   // typed continuation opcodes
   ContNew = 0xe0,
   ContBind = 0xe1,
+  Suspend = 0xe2,
   Resume = 0xe3,
 
 };
@@ -1810,6 +1811,7 @@ class WasmBinaryReader {
   void visitContNew(ContNew* curr);
   void visitContBind(ContBind* curr);
   void visitResume(Resume* curr);
+  void visitSuspend(Suspend* curr);
 
   [[noreturn]] void throwError(std::string text);
 

src/wasm-builder.h

@@ -1231,6 +1231,13 @@ class Builder {
     ret->finalize(&wasm);
     return ret;
   }
+  Suspend* makeSuspend(Name tag, const std::vector<Expression*>& args) {
+    auto* ret = wasm.allocator.alloc<Suspend>();
+    ret->tag = tag;
+    ret->operands.set(args);
+    ret->finalize(&wasm);
+    return ret;
+  }
 
   // Additional helpers
 

src/wasm-delegations-fields.def

@@ -973,6 +973,13 @@ switch (DELEGATE_ID) {
     DELEGATE_END(Resume);
     break;
   }
+  case Expression::Id::SuspendId: {
+    DELEGATE_START(Suspend);
+    DELEGATE_FIELD_CHILD_VECTOR(Suspend, operands);
+    DELEGATE_FIELD_NAME_KIND(Suspend, tag, ModuleItemKind::Tag);
+    DELEGATE_END(Suspend);
+    break;
+  }
 }
 
 #undef DELEGATE_ID

src/wasm-delegations.def

@@ -108,5 +108,6 @@ DELEGATE(StringSliceIter);
 DELEGATE(ContBind);
 DELEGATE(ContNew);
 DELEGATE(Resume);
+DELEGATE(Suspend);
 
 #undef DELEGATE

src/wasm-interpreter.h

@@ -2404,6 +2404,7 @@ class ConstantExpressionRunner : public ExpressionRunner<SubType> {
   Flow visitContBind(ContBind* curr) { WASM_UNREACHABLE("unimplemented"); }
   Flow visitContNew(ContNew* curr) { WASM_UNREACHABLE("unimplemented"); }
   Flow visitResume(Resume* curr) { WASM_UNREACHABLE("unimplemented"); }
+  Flow visitSuspend(Suspend* curr) { WASM_UNREACHABLE("unimplemented"); }
 
   void trap(const char* why) override { throw NonconstantException(); }
 
@@ -3980,6 +3981,7 @@ class ModuleRunnerBase : public ExpressionRunner<SubType> {
   Flow visitContBind(ContBind* curr) { return Flow(NONCONSTANT_FLOW); }
   Flow visitContNew(ContNew* curr) { return Flow(NONCONSTANT_FLOW); }
   Flow visitResume(Resume* curr) { return Flow(NONCONSTANT_FLOW); }
+  Flow visitSuspend(Suspend* curr) { return Flow(NONCONSTANT_FLOW); }
 
   void trap(const char* why) override { externalInterface->trap(why); }
 

src/wasm-ir-builder.h

@@ -216,6 +216,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   [[nodiscard]] Result<> makeResume(HeapType ct,
                                     const std::vector<Name>& tags,
                                     const std::vector<Index>& labels);
+  [[nodiscard]] Result<> makeSuspend(Name tag);
 
   // Private functions that must be public for technical reasons.
   [[nodiscard]] Result<> visitExpression(Expression*);
@@ -256,6 +257,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   [[nodiscard]] Result<> visitStringEncode(StringEncode*);
   [[nodiscard]] Result<> visitContBind(ContBind*);
   [[nodiscard]] Result<> visitResume(Resume*);
+  [[nodiscard]] Result<> visitSuspend(Suspend*);
   [[nodiscard]] Result<> visitTupleMake(TupleMake*);
   [[nodiscard]] Result<>
   visitTupleExtract(TupleExtract*,

src/wasm-s-parser.h

@@ -331,6 +331,7 @@ class SExpressionWasmBuilder {
   Expression* makeContBind(Element& s);
   Expression* makeContNew(Element& s);
   Expression* makeResume(Element& s);
+  Expression* makeSuspend(Element& s);
 
   // Helper functions
   Type parseBlockType(Element& s, Index& i);

src/wasm.h

@@ -746,6 +746,7 @@ class Expression {
     ContBindId,
     ContNewId,
     ResumeId,
+    SuspendId,
     NumExpressionIds
   };
   Id _id;
@@ -2048,6 +2049,19 @@ class Resume : public SpecificExpression<Expression::ResumeId> {
   ArenaVector<Type> sentTypes;
 };
 
+class Suspend : public SpecificExpression<Expression::SuspendId> {
+public:
+  Suspend(MixedArena& allocator) : operands(allocator) {}
+
+  Name tag;
+  ExpressionList operands;
+
+  // We need access to the module to obtain the signature of the tag,
+  // which determines this node's type.
+  // If no module is given, then the type must have been set already.
+  void finalize(Module* wasm = nullptr);
+};
+
 // Globals
 
 struct Named {

src/wasm/wasm-binary.cpp

@@ -4060,6 +4060,10 @@ BinaryConsts::ASTNodes WasmBinaryReader::readExpression(Expression*& curr) {
       visitResume((curr = allocator.alloc<Resume>())->cast<Resume>());
       break;
     }
+    case BinaryConsts::Suspend: {
+      visitSuspend((curr = allocator.alloc<Suspend>())->cast<Suspend>());
+      break;
+    }
     case BinaryConsts::AtomicPrefix: {
       code = static_cast<uint8_t>(getU32LEB());
       if (maybeVisitLoad(curr, code, /*isAtomic=*/true)) {
@@ -7869,6 +7873,26 @@ void WasmBinaryReader::visitResume(Resume* curr) {
   curr->finalize(&wasm);
 }
 
+void WasmBinaryReader::visitSuspend(Suspend* curr) {
+  BYN_TRACE("zz node: Suspend\n");
+
+  auto tagIndex = getU32LEB();
+  if (tagIndex >= wasm.tags.size()) {
+    throwError("bad tag index");
+  }
+  auto* tag = wasm.tags[tagIndex].get();
+  curr->tag = tag->name;
+  tagRefs[tagIndex].push_back(&curr->tag);
+
+  auto numArgs = tag->sig.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);
 }

src/wasm/wasm-ir-builder.cpp

@@ -658,6 +658,19 @@ Result<> IRBuilder::visitResume(Resume* curr) {
   return Ok{};
 }
 
+Result<> IRBuilder::visitSuspend(Suspend* curr) {
+  auto tag = wasm.getTag(curr->tag);
+  auto sig = tag->sig;
+  auto size = sig.params.size();
+  curr->operands.resize(size);
+  for (size_t i = 0; i < size; ++i) {
+    auto val = pop();
+    CHECK_ERR(val);
+    curr->operands[size - i - 1] = *val;
+  }
+  return Ok{};
+}
+
 Result<> IRBuilder::visitTupleMake(TupleMake* curr) {
   assert(curr->operands.size() >= 2);
   for (size_t i = 0, size = curr->operands.size(); i < size; ++i) {
@@ -1930,4 +1943,14 @@ Result<> IRBuilder::makeResume(HeapType ct,
   return Ok{};
 }
 
+Result<> IRBuilder::makeSuspend(Name tag) {
+  Suspend curr(wasm.allocator);
+  curr.tag = tag;
+  CHECK_ERR(visitSuspend(&curr));
+
+  std::vector<Expression*> operands(curr.operands.begin(), curr.operands.end());
+  push(builder.makeSuspend(tag, operands));
+  return Ok{};
+}
+
 } // namespace wasm

src/wasm/wasm-s-parser.cpp

@@ -3059,6 +3059,20 @@ Expression* SExpressionWasmBuilder::makeResume(Element& s) {
   return ret;
 }
 
+Expression* SExpressionWasmBuilder::makeSuspend(Element& s) {
+  auto ret = allocator.alloc<Suspend>();
+
+  ret->tag = getTagName(*s[1]);
+
+  Index i = 2;
+  while (i < s.size()) {
+    ret->operands.push_back(parseExpression(s[i++]));
+  }
+
+  ret->finalize(&wasm);
+  return ret;
+}
+
 Expression* SExpressionWasmBuilder::makeRefI31(Element& s) {
   auto ret = allocator.alloc<RefI31>();
   ret->value = parseExpression(s[1]);

src/wasm/wasm-stack.cpp

@@ -2512,6 +2512,10 @@ void BinaryInstWriter::visitResume(Resume* curr) {
   }
 }
 
+void BinaryInstWriter::visitSuspend(Suspend* curr) {
+  o << int8_t(BinaryConsts::Suspend) << U32LEB(parent.getTagIndex(curr->tag));
+}
+
 void BinaryInstWriter::emitScopeEnd(Expression* curr) {
   assert(!breakStack.empty());
   breakStack.pop_back();