Typed continuations: cont.new instructions (#6308)

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

This particular PR adds support for the `cont.new` instruction for creating
continuations, documented [here(https://github.com/wasmfx/specfx/blob/main/proposals/continuations/Overview.md#instructions).

In short, these instructions are of the form `(cont.new $ct)` where `$ct` must
be a continuation type. The instruction takes a single (nullable) function
reference as its argument, which means that the folded representation of the
instruction is of the form `(cont.new $ct (foo ...))`. 

Support for the instruction is implemented in both the old and the new wat
parser.

Note that this PR does not implement validation of the new instruction.

26 files changed, 187 insertions, 28 deletions

diff --git a/src/gen-s-parser.inc b/src/gen-s-parser.inc
index 8387dbb8e..89063ec81 100644
--- a/src/gen-s-parser.inc
+++ b/src/gen-s-parser.inc
@@ -149,21 +149,29 @@ switch (buf[0]) {
     }
   }
   case 'c': {
-    switch (buf[4]) {
-      case '\0':
-        if (op == "call"sv) { return makeCall(s, /*isReturn=*/false); }
-        goto parse_error;
-      case '_': {
-        switch (buf[5]) {
-          case 'i':
-            if (op == "call_indirect"sv) { return makeCallIndirect(s, /*isReturn=*/false); }
-            goto parse_error;
-          case 'r':
-            if (op == "call_ref"sv) { return makeCallRef(s, /*isReturn=*/false); }
+    switch (buf[1]) {
+      case 'a': {
+        switch (buf[4]) {
+          case '\0':
+            if (op == "call"sv) { return makeCall(s, /*isReturn=*/false); }
             goto parse_error;
+          case '_': {
+            switch (buf[5]) {
+              case 'i':
+                if (op == "call_indirect"sv) { return makeCallIndirect(s, /*isReturn=*/false); }
+                goto parse_error;
+              case 'r':
+                if (op == "call_ref"sv) { return makeCallRef(s, /*isReturn=*/false); }
+                goto parse_error;
+              default: goto parse_error;
+            }
+          }
           default: goto parse_error;
         }
       }
+      case 'o':
+        if (op == "cont.new"sv) { return makeContNew(s); }
+        goto parse_error;
       default: goto parse_error;
     }
   }
@@ -3816,30 +3824,41 @@ switch (buf[0]) {
     }
   }
   case 'c': {
-    switch (buf[4]) {
-      case '\0':
-        if (op == "call"sv) {
-          CHECK_ERR(makeCall(ctx, pos, /*isReturn=*/false));
-          return Ok{};
-        }
-        goto parse_error;
-      case '_': {
-        switch (buf[5]) {
-          case 'i':
-            if (op == "call_indirect"sv) {
-              CHECK_ERR(makeCallIndirect(ctx, pos, /*isReturn=*/false));
+    switch (buf[1]) {
+      case 'a': {
+        switch (buf[4]) {
+          case '\0':
+            if (op == "call"sv) {
+              CHECK_ERR(makeCall(ctx, pos, /*isReturn=*/false));
               return Ok{};
             }
             goto parse_error;
-          case 'r':
-            if (op == "call_ref"sv) {
-              CHECK_ERR(makeCallRef(ctx, pos, /*isReturn=*/false));
-              return Ok{};
+          case '_': {
+            switch (buf[5]) {
+              case 'i':
+                if (op == "call_indirect"sv) {
+                  CHECK_ERR(makeCallIndirect(ctx, pos, /*isReturn=*/false));
+                  return Ok{};
+                }
+                goto parse_error;
+              case 'r':
+                if (op == "call_ref"sv) {
+                  CHECK_ERR(makeCallRef(ctx, pos, /*isReturn=*/false));
+                  return Ok{};
+                }
+                goto parse_error;
+              default: goto parse_error;
             }
-            goto parse_error;
+          }
           default: goto parse_error;
         }
       }
+      case 'o':
+        if (op == "cont.new"sv) {
+          CHECK_ERR(makeContNew(ctx, pos));
+          return Ok{};
+        }
+        goto parse_error;
       default: goto parse_error;
     }
   }
diff --git a/src/ir/ReFinalize.cpp b/src/ir/ReFinalize.cpp
index 0338f630c..7c81b0d0d 100644
--- a/src/ir/ReFinalize.cpp
+++ b/src/ir/ReFinalize.cpp
@@ -182,6 +182,7 @@ void ReFinalize::visitStringSliceIter(StringSliceIter* curr) {
   curr->finalize();
 }
 
+void ReFinalize::visitContNew(ContNew* curr) { curr->finalize(); }
 void ReFinalize::visitResume(Resume* curr) { curr->finalize(); }
 
 void ReFinalize::visitExport(Export* curr) { WASM_UNREACHABLE("unimp"); }
diff --git a/src/ir/cost.h b/src/ir/cost.h
index 821e46524..7a8776c00 100644
--- a/src/ir/cost.h
+++ b/src/ir/cost.h
@@ -726,6 +726,10 @@ struct CostAnalyzer : public OverriddenVisitor<CostAnalyzer, CostType> {
     return 8 + visit(curr->ref) + visit(curr->num);
   }
 
+  CostType visitContNew(ContNew* curr) {
+    // Some arbitrary "high" value, reflecting that this may allocate a stack
+    return 14 + visit(curr->func);
+  }
   CostType visitResume(Resume* curr) {
     // Inspired by indirect calls, but twice the cost.
     return 12 + visit(curr->cont);
diff --git a/src/ir/effects.h b/src/ir/effects.h
index 3ff320a88..d0af932bd 100644
--- a/src/ir/effects.h
+++ b/src/ir/effects.h
@@ -974,6 +974,10 @@ private:
       parent.implicitTrap = true;
     }
 
+    void visitContNew(ContNew* curr) {
+      // traps when curr->func is null ref.
+      parent.implicitTrap = true;
+    }
     void visitResume(Resume* curr) {
       // This acts as a kitchen sink effect.
       parent.calls = true;
diff --git a/src/ir/module-utils.cpp b/src/ir/module-utils.cpp
index 304ddb04d..b61fddd8b 100644
--- 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* contNew = curr->dynCast<ContNew>()) {
+      counts.note(contNew->contType);
     } else if (auto* resume = curr->dynCast<Resume>()) {
       counts.note(resume->contType);
     } else if (Properties::isControlFlowStructure(curr)) {
diff --git a/src/ir/possible-contents.cpp b/src/ir/possible-contents.cpp
index 8d6891017..d0b0ea3ae 100644
--- a/src/ir/possible-contents.cpp
+++ b/src/ir/possible-contents.cpp
@@ -1200,6 +1200,10 @@ struct InfoCollector
 
   void visitReturn(Return* curr) { addResult(curr->value); }
 
+  void visitContNew(ContNew* curr) {
+    // TODO: optimize when possible
+    addRoot(curr);
+  }
   void visitResume(Resume* curr) {
     // TODO: optimize when possible
     addRoot(curr);
diff --git a/src/ir/subtype-exprs.h b/src/ir/subtype-exprs.h
index 056b373c2..6b8348319 100644
--- a/src/ir/subtype-exprs.h
+++ b/src/ir/subtype-exprs.h
@@ -368,6 +368,7 @@ struct SubtypingDiscoverer : public OverriddenVisitor<SubType> {
   void visitStringSliceWTF(StringSliceWTF* curr) {}
   void visitStringSliceIter(StringSliceIter* curr) {}
 
+  void visitContNew(ContNew* curr) { WASM_UNREACHABLE("not implemented"); }
   void visitResume(Resume* curr) { WASM_UNREACHABLE("not implemented"); }
 };
 
diff --git a/src/parser/contexts.h b/src/parser/contexts.h
index 50a2abd96..d2f0ea2d0 100644
--- a/src/parser/contexts.h
+++ b/src/parser/contexts.h
@@ -570,6 +570,9 @@ struct NullInstrParserCtx {
   Result<> makeStringIterMove(Index, StringIterMoveOp) { return Ok{}; }
   Result<> makeStringSliceWTF(Index, StringSliceWTFOp) { return Ok{}; }
   Result<> makeStringSliceIter(Index) { return Ok{}; }
+  template<typename HeapTypeT> Result<> makeContNew(Index, HeapTypeT) {
+    return Ok{};
+  }
   template<typename HeapTypeT>
   Result<> makeResume(Index, HeapTypeT, const TagLabelListT&) {
     return Ok{};
@@ -2010,6 +2013,10 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
     return withLoc(pos, irBuilder.makeStringSliceIter());
   }
 
+  Result<> makeContNew(Index pos, HeapType type) {
+    return withLoc(pos, irBuilder.makeContNew(type));
+  }
+
   Result<>
   makeResume(Index pos, HeapType type, const TagLabelListT& tagLabels) {
     std::vector<Name> tags;
diff --git a/src/parser/parsers.h b/src/parser/parsers.h
index 1069a39bd..007f815ac 100644
--- a/src/parser/parsers.h
+++ b/src/parser/parsers.h
@@ -167,6 +167,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<> makeContNew(Ctx&, Index);
 template<typename Ctx> Result<> makeResume(Ctx&, Index);
 
 // Modules
@@ -1990,6 +1991,13 @@ template<typename Ctx> Result<> makeStringSliceIter(Ctx& ctx, Index pos) {
   return ctx.makeStringSliceIter(pos);
 }
 
+template<typename Ctx> Result<> makeContNew(Ctx& ctx, Index pos) {
+  auto type = typeidx(ctx);
+  CHECK_ERR(type);
+
+  return ctx.makeContNew(pos, *type);
+}
+
 // resume ::= 'resume' typeidx ('(' 'tag' tagidx labelidx ')')*
 template<typename Ctx> Result<> makeResume(Ctx& ctx, Index pos) {
   auto type = typeidx(ctx);
diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index 6e4c1452c..2288b6914 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -2379,6 +2379,11 @@ struct PrintExpressionContents
     printMedium(o, "stringview_iter.slice");
   }
 
+  void visitContNew(ContNew* curr) {
+    printMedium(o, "cont.new ");
+    printHeapType(curr->contType);
+  }
+
   void visitResume(Resume* curr) {
     printMedium(o, "resume");
 
diff --git a/src/passes/TypeGeneralizing.cpp b/src/passes/TypeGeneralizing.cpp
index 293d17cdf..c19766ae3 100644
--- a/src/passes/TypeGeneralizing.cpp
+++ b/src/passes/TypeGeneralizing.cpp
@@ -875,6 +875,7 @@ struct TransferFn : OverriddenVisitor<TransferFn> {
   void visitStringSliceWTF(StringSliceWTF* curr) { WASM_UNREACHABLE("TODO"); }
   void visitStringSliceIter(StringSliceIter* curr) { WASM_UNREACHABLE("TODO"); }
 
+  void visitContNew(ContNew* curr) { WASM_UNREACHABLE("TODO"); }
   void visitResume(Resume* curr) { WASM_UNREACHABLE("TODO"); }
 };
 
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index b86474a8e..c1d79ec2a 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -1298,6 +1298,7 @@ enum ASTNodes {
   StringNewUTF8ArrayTry = 0xb8,
 
   // typed continuation opcodes
+  ContNew = 0xe0,
   Resume = 0xe3,
 
 };
@@ -1926,6 +1927,7 @@ public:
   void visitCallRef(CallRef* curr);
   void visitRefAsCast(RefCast* curr, uint32_t code);
   void visitRefAs(RefAs* curr, uint8_t code);
+  void visitContNew(ContNew* curr);
   void visitResume(Resume* curr);
 
   [[noreturn]] void throwError(std::string text);
diff --git a/src/wasm-builder.h b/src/wasm-builder.h
index 3096ada07..03fddb178 100644
--- a/src/wasm-builder.h
+++ b/src/wasm-builder.h
@@ -1198,6 +1198,13 @@ public:
     return ret;
   }
 
+  ContNew* makeContNew(HeapType contType, Expression* func) {
+    auto* ret = wasm.allocator.alloc<ContNew>();
+    ret->contType = contType;
+    ret->func = func;
+    ret->finalize();
+    return ret;
+  }
   Resume* makeResume(HeapType contType,
                      const std::vector<Name>& handlerTags,
                      const std::vector<Name>& handlerBlocks,
diff --git a/src/wasm-delegations-fields.def b/src/wasm-delegations-fields.def
index aeb32e689..c895c31dd 100644
--- a/src/wasm-delegations-fields.def
+++ b/src/wasm-delegations-fields.def
@@ -946,6 +946,13 @@ switch (DELEGATE_ID) {
     break;
   }
 
+  case Expression::Id::ContNewId: {
+    DELEGATE_START(ContNew);
+    DELEGATE_FIELD_CHILD(ContNew, func);
+    DELEGATE_FIELD_HEAPTYPE(ContNew, contType);
+    DELEGATE_END(ContNew);
+    break;
+  }
   case Expression::Id::ResumeId: {
     DELEGATE_START(Resume);
     DELEGATE_FIELD_TYPE_VECTOR(Resume, sentTypes);
diff --git a/src/wasm-delegations.def b/src/wasm-delegations.def
index 721c56db9..5a57e399b 100644
--- a/src/wasm-delegations.def
+++ b/src/wasm-delegations.def
@@ -105,6 +105,7 @@ DELEGATE(StringIterNext);
 DELEGATE(StringIterMove);
 DELEGATE(StringSliceWTF);
 DELEGATE(StringSliceIter);
+DELEGATE(ContNew);
 DELEGATE(Resume);
 
 #undef DELEGATE
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index b65b909c7..409bec9d8 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -2401,6 +2401,7 @@ public:
     }
     return ExpressionRunner<SubType>::visitRefAs(curr);
   }
+  Flow visitContNew(ContNew* curr) { WASM_UNREACHABLE("unimplemented"); }
   Flow visitResume(Resume* curr) { WASM_UNREACHABLE("unimplemented"); }
 
   void trap(const char* why) override { throw NonconstantException(); }
@@ -3968,6 +3969,7 @@ public:
     multiValues.pop_back();
     return ret;
   }
+  Flow visitContNew(ContNew* curr) { return Flow(NONCONSTANT_FLOW); }
   Flow visitResume(Resume* curr) { return Flow(NONCONSTANT_FLOW); }
 
   void trap(const char* why) override { externalInterface->trap(why); }
diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index 66dfd6dbe..05c59ca95 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -206,6 +206,7 @@ public:
   [[nodiscard]] Result<> makeStringIterMove(StringIterMoveOp op);
   [[nodiscard]] Result<> makeStringSliceWTF(StringSliceWTFOp op);
   [[nodiscard]] Result<> makeStringSliceIter();
+  [[nodiscard]] Result<> makeContNew(HeapType ct);
   [[nodiscard]] Result<> makeResume(HeapType ct,
                                     const std::vector<Name>& tags,
                                     const std::vector<Index>& labels);
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index b8f9d6c65..953599a1e 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -328,6 +328,7 @@ private:
   Expression* makeStringIterMove(Element& s, StringIterMoveOp op);
   Expression* makeStringSliceWTF(Element& s, StringSliceWTFOp op);
   Expression* makeStringSliceIter(Element& s);
+  Expression* makeContNew(Element& s);
   Expression* makeResume(Element& s);
 
   // Helper functions
diff --git a/src/wasm.h b/src/wasm.h
index 366aa7302..3da245547 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -743,6 +743,7 @@ public:
     StringIterMoveId,
     StringSliceWTFId,
     StringSliceIterId,
+    ContNewId,
     ResumeId,
     NumExpressionIds
   };
@@ -1997,6 +1998,17 @@ public:
   void finalize();
 };
 
+class ContNew : public SpecificExpression<Expression::ContNewId> {
+public:
+  ContNew() = default;
+  ContNew(MixedArena& allocator) {}
+
+  HeapType contType;
+  Expression* func;
+
+  void finalize();
+};
+
 class Resume : public SpecificExpression<Expression::ResumeId> {
 public:
   Resume(MixedArena& allocator)
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 06c6f3648..6f41c3c87 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -4046,6 +4046,12 @@ BinaryConsts::ASTNodes WasmBinaryReader::readExpression(Expression*& curr) {
       visitCallRef(call);
       break;
     }
+    case BinaryConsts::ContNew: {
+      auto contNew = allocator.alloc<ContNew>();
+      curr = contNew;
+      visitContNew(contNew);
+      break;
+    }
     case BinaryConsts::Resume: {
       visitResume((curr = allocator.alloc<Resume>())->cast<Resume>());
       break;
@@ -7762,6 +7768,20 @@ void WasmBinaryReader::visitRefAs(RefAs* curr, uint8_t code) {
   curr->finalize();
 }
 
+void WasmBinaryReader::visitContNew(ContNew* curr) {
+  BYN_TRACE("zz node: ContNew\n");
+
+  auto contTypeIndex = getU32LEB();
+  curr->contType = getTypeByIndex(contTypeIndex);
+  if (!curr->contType.isContinuation()) {
+    throwError("non-continuation type in cont.new instruction " +
+               curr->contType.toString());
+  }
+
+  curr->func = popNonVoidExpression();
+  curr->finalize();
+}
+
 void WasmBinaryReader::visitResume(Resume* curr) {
   BYN_TRACE("zz node: Resume\n");
 
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index bd3a3ad11..6676d1cb4 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -1783,6 +1783,17 @@ Result<> IRBuilder::makeStringSliceIter() {
   return Ok{};
 }
 
+Result<> IRBuilder::makeContNew(HeapType ct) {
+  if (!ct.isContinuation()) {
+    return Err{"expected continuation type"};
+  }
+  ContNew curr;
+  CHECK_ERR(visitContNew(&curr));
+
+  push(builder.makeContNew(ct, curr.func));
+  return Ok{};
+}
+
 Result<> IRBuilder::makeResume(HeapType ct,
                                const std::vector<Name>& tags,
                                const std::vector<Index>& labels) {
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index 95cb453fb..b7113d2b2 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -2990,6 +2990,20 @@ Expression* SExpressionWasmBuilder::makeCallRef(Element& s, bool isReturn) {
     target, operands, sigType.getSignature().results, isReturn);
 }
 
+Expression* SExpressionWasmBuilder::makeContNew(Element& s) {
+  auto ret = allocator.alloc<ContNew>();
+
+  ret->contType = parseHeapType(*s[1]);
+  if (!ret->contType.isContinuation()) {
+    throw ParseException("expected continuation type", s[1]->line, s[1]->col);
+  }
+
+  ret->func = parseExpression(s[2]);
+
+  ret->finalize();
+  return ret;
+}
+
 Expression* SExpressionWasmBuilder::makeResume(Element& s) {
   auto ret = allocator.alloc<Resume>();
 
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index 12ce68324..0857e9a0d 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -2489,6 +2489,11 @@ void BinaryInstWriter::visitStringSliceIter(StringSliceIter* curr) {
     << U32LEB(BinaryConsts::StringViewIterSlice);
 }
 
+void BinaryInstWriter::visitContNew(ContNew* curr) {
+  o << int8_t(BinaryConsts::ContNew);
+  parent.writeIndexedHeapType(curr->contType);
+}
+
 void BinaryInstWriter::visitResume(Resume* curr) {
   o << int8_t(BinaryConsts::Resume);
   parent.writeIndexedHeapType(curr->contType);
diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
index 3fbcb0bfb..d745585e5 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 visitContNew(ContNew* curr);
   void visitResume(Resume* curr);
 
   void visitFunction(Function* curr);
@@ -3294,6 +3295,19 @@ void FunctionValidator::visitStringSliceIter(StringSliceIter* curr) {
                "string operations require reference-types [--enable-strings]");
 }
 
+void FunctionValidator::visitContNew(ContNew* curr) {
+  // TODO implement actual type-checking
+  shouldBeTrue(
+    !getModule() || getModule()->features.hasTypedContinuations(),
+    curr,
+    "cont.new requires typed-continuatons [--enable-typed-continuations]");
+
+  shouldBeTrue((curr->contType.isContinuation() &&
+                curr->contType.getContinuation().type.isSignature()),
+               curr,
+               "invalid type in ContNew expression");
+}
+
 void FunctionValidator::visitResume(Resume* curr) {
   // TODO implement actual type-checking
   shouldBeTrue(
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index 4785b22ab..ce5fd7006 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -1352,6 +1352,8 @@ void StringSliceIter::finalize() {
   }
 }
 
+void ContNew::finalize() { type = Type(contType, NonNullable); }
+
 static void populateResumeSentTypes(Resume* curr, Module* wasm) {
   if (!wasm) {
     return;
diff --git a/src/wasm2js.h b/src/wasm2js.h
index ae3f60e18..92ba397be 100644
--- a/src/wasm2js.h
+++ b/src/wasm2js.h
@@ -2428,6 +2428,10 @@ Ref Wasm2JSBuilder::processFunctionBody(Module* m,
       WASM_UNREACHABLE("unimp");
     }
 
+    Ref visitContNew(ContNew* curr) {
+      unimplemented(curr);
+      WASM_UNREACHABLE("unimp");
+    }
     Ref visitResume(Resume* curr) {
       unimplemented(curr);
       WASM_UNREACHABLE("unimp");