Stack switching proposal support (WebAssembly#7041)

This patch implements text and binary encoding/decoding support for the
stack switching proposal. It does so by adapting the previous
typed-continunations implementation. Particular changes:

* Support for new `resume` encoding.
* Added support for `resume_throw` and `switch`.
* Feature flag `typed-continuations` has been renamed to
`stack-switching`.

A small unfortunate implementation detail is that the internal name
`Switch` was already taken by the `br_table` instruction, so I opted to
give the `switch` instruction the internal name `StackSwitch`.

A minor detail is that I have reordered the declarations/definitions of
the stack switching instructions such that they appear in ascending
order according to their opcode value (this is the same order that the
stack-switching explainer document present them in).

I can look into adding validation support in a subsequent patch.

Author: dhil

scripts/gen-s-parser.py

@@ -596,11 +596,13 @@
     # Typed function references instructions
     ("call_ref",             "makeCallRef(/*isReturn=*/false)"),
     ("return_call_ref",      "makeCallRef(/*isReturn=*/true)"),
-    # Typed continuations instructions
+    # Stack switching instructions
     ("cont.new",             "makeContNew()"),
     ("cont.bind",            "makeContBind()"),
-    ("resume",               "makeResume()"),
     ("suspend",              "makeSuspend()"),
+    ("resume",               "makeResume()"),
+    ("resume_throw",         "makeResumeThrow()"),
+    ("switch",               "makeStackSwitch()"),
     # GC
     ("ref.i31",              "makeRefI31(Unshared)"),
     ("ref.i31_shared",       "makeRefI31(Shared)"),

scripts/test/fuzzing.py

@@ -90,6 +90,14 @@
     'names.wast',
     # huge amount of locals that make it extremely slow
     'too_much_for_liveness.wasm',
+    # TODO: fuzzer support for stack switching
+    'stack_switching.wast',
+    'stack_switching_contnew.wast',
+    'stack_switching_contbind.wast',
+    'stack_switching_suspend.wast',
+    'stack_switching_resume.wast',
+    'stack_switching_resume_throw.wast',
+    'stack_switching_switch.wast'
 ]
 
 

src/gen-s-parser.inc

@@ -4807,12 +4807,23 @@ switch (buf[0]) {
           default: goto parse_error;
         }
       }
-      case 's':
-        if (op == "resume"sv) {
-          CHECK_ERR(makeResume(ctx, pos, annotations));
-          return Ok{};
+      case 's': {
+        switch (buf[6]) {
+          case '\0':
+            if (op == "resume"sv) {
+              CHECK_ERR(makeResume(ctx, pos, annotations));
+              return Ok{};
+            }
+            goto parse_error;
+          case '_':
+            if (op == "resume_throw"sv) {
+              CHECK_ERR(makeResumeThrow(ctx, pos, annotations));
+              return Ok{};
+            }
+            goto parse_error;
+          default: goto parse_error;
         }
-        goto parse_error;
+      }
       case 't': {
         switch (buf[3]) {
           case 'h':
@@ -5172,6 +5183,12 @@ switch (buf[0]) {
           return Ok{};
         }
         goto parse_error;
+      case 'w':
+        if (op == "switch"sv) {
+          CHECK_ERR(makeStackSwitch(ctx, pos, annotations));
+          return Ok{};
+        }
+        goto parse_error;
       default: goto parse_error;
     }
   }

src/interpreter/interpreter.cpp

@@ -153,10 +153,12 @@ struct ExpressionInterpreter : OverriddenVisitor<ExpressionInterpreter, Flow> {
   Flow visitStringEq(StringEq* curr) { WASM_UNREACHABLE("TODO"); }
   Flow visitStringWTF16Get(StringWTF16Get* curr) { WASM_UNREACHABLE("TODO"); }
   Flow visitStringSliceWTF(StringSliceWTF* curr) { WASM_UNREACHABLE("TODO"); }
-  Flow visitContBind(ContBind* curr) { WASM_UNREACHABLE("TODO"); }
   Flow visitContNew(ContNew* curr) { WASM_UNREACHABLE("TODO"); }
-  Flow visitResume(Resume* curr) { WASM_UNREACHABLE("TODO"); }
+  Flow visitContBind(ContBind* curr) { WASM_UNREACHABLE("TODO"); }
   Flow visitSuspend(Suspend* curr) { WASM_UNREACHABLE("TODO"); }
+  Flow visitResume(Resume* curr) { WASM_UNREACHABLE("TODO"); }
+  Flow visitResumeThrow(ResumeThrow* curr) { WASM_UNREACHABLE("TODO"); }
+  Flow visitStackSwitch(StackSwitch* curr) { WASM_UNREACHABLE("TODO"); }
 };
 
 } // anonymous namespace

src/ir/ReFinalize.cpp

@@ -182,8 +182,10 @@ void ReFinalize::visitStringWTF16Get(StringWTF16Get* curr) { curr->finalize(); }
 void ReFinalize::visitStringSliceWTF(StringSliceWTF* curr) { curr->finalize(); }
 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::visitResume(Resume* curr) { curr->finalize(); }
+void ReFinalize::visitResumeThrow(ResumeThrow* curr) { curr->finalize(); }
+void ReFinalize::visitStackSwitch(StackSwitch* curr) { curr->finalize(); }
 
 void ReFinalize::visitExport(Export* curr) { WASM_UNREACHABLE("unimp"); }
 void ReFinalize::visitGlobal(Global* curr) { WASM_UNREACHABLE("unimp"); }

src/ir/branch-utils.h

@@ -85,7 +85,14 @@ void operateOnScopeNameUsesAndSentTypes(Expression* expr, T func) {
     } else if (auto* r = expr->dynCast<Resume>()) {
       for (Index i = 0; i < r->handlerTags.size(); i++) {
         auto dest = r->handlerTags[i];
-        if (dest == name) {
+        if (!dest.isNull() && dest == name) {
+          func(name, r->sentTypes[i]);
+        }
+      }
+    } else if (auto* r = expr->dynCast<ResumeThrow>()) {
+      for (Index i = 0; i < r->handlerTags.size(); i++) {
+        auto dest = r->handlerTags[i];
+        if (!dest.isNull() && dest == name) {
           func(name, r->sentTypes[i]);
         }
       }
@@ -118,6 +125,8 @@ void operateOnScopeNameUsesAndSentValues(Expression* expr, T func) {
       // 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 if (expr->is<ResumeThrow>()) {
+      func(name, nullptr);
     } else {
       assert(expr->is<Try>() || expr->is<Rethrow>()); // delegate or rethrow
     }

src/ir/child-typer.h

@@ -1085,39 +1085,76 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     note(&curr->end, Type::i32);
   }
 
-  void visitContBind(ContBind* curr) {
-    auto paramsBefore =
-      curr->contTypeBefore.getContinuation().type.getSignature().params;
-    auto paramsAfter =
-      curr->contTypeAfter.getContinuation().type.getSignature().params;
-    assert(paramsBefore.size() >= paramsAfter.size());
-    auto n = paramsBefore.size() - paramsAfter.size();
+  void visitContNew(ContNew* curr) { note(&curr->func, curr->type); }
+
+  void visitContBind(ContBind* curr,
+                     std::optional<HeapType> src = std::nullopt,
+                     std::optional<HeapType> dest = std::nullopt) {
+    if (!src.has_value()) {
+      src = curr->cont->type.getHeapType();
+    }
+    if (!dest.has_value()) {
+      dest = curr->type.getHeapType();
+    }
+    auto sourceParams = src->getContinuation().type.getSignature().params;
+    auto targetParams = dest->getContinuation().type.getSignature().params;
+    assert(sourceParams.size() >= targetParams.size());
+    auto n = sourceParams.size() - targetParams.size();
     assert(curr->operands.size() == n);
     for (size_t i = 0; i < n; ++i) {
-      note(&curr->operands[i], paramsBefore[i]);
+      note(&curr->operands[i], sourceParams[i]);
     }
-    note(&curr->cont, Type(curr->contTypeBefore, Nullable));
+    note(&curr->cont, Type(*src, Nullable));
   }
 
-  void visitContNew(ContNew* curr) {
-    note(&curr->func, Type(curr->contType.getContinuation().type, Nullable));
+  void visitSuspend(Suspend* curr) {
+    auto params = wasm.getTag(curr->tag)->params();
+    assert(params.size() == curr->operands.size());
+    for (size_t i = 0; i < params.size(); ++i) {
+      note(&curr->operands[i], params[i]);
+    }
   }
 
-  void visitResume(Resume* curr) {
-    auto params = curr->contType.getContinuation().type.getSignature().params;
+  void visitResume(Resume* curr, std::optional<HeapType> ct = std::nullopt) {
+    if (!ct.has_value()) {
+      ct = curr->cont->type.getHeapType();
+    }
+    assert(ct->isContinuation());
+    auto params = ct->getContinuation().type.getSignature().params;
     assert(params.size() == curr->operands.size());
     for (size_t i = 0; i < params.size(); ++i) {
       note(&curr->operands[i], params[i]);
     }
-    note(&curr->cont, Type(curr->contType, Nullable));
+    note(&curr->cont, Type(*ct, Nullable));
   }
 
-  void visitSuspend(Suspend* curr) {
+  void visitResumeThrow(ResumeThrow* curr,
+                        std::optional<HeapType> ct = std::nullopt) {
+    if (!ct.has_value()) {
+      ct = curr->cont->type.getHeapType();
+    }
+    assert(ct->isContinuation());
     auto params = wasm.getTag(curr->tag)->params();
     assert(params.size() == curr->operands.size());
     for (size_t i = 0; i < params.size(); ++i) {
       note(&curr->operands[i], params[i]);
     }
+    note(&curr->cont, Type(*ct, Nullable));
+  }
+
+  void visitStackSwitch(StackSwitch* curr,
+                        std::optional<HeapType> ct = std::nullopt) {
+    if (!ct.has_value()) {
+      ct = curr->cont->type.getHeapType();
+    }
+    assert(ct->isContinuation());
+    auto params = ct->getContinuation().type.getSignature().params;
+    assert(params.size() >= 1 &&
+           ((params.size() - 1) == curr->operands.size()));
+    for (size_t i = 0; i < params.size() - 1; ++i) {
+      note(&curr->operands[i], params[i]);
+    }
+    note(&curr->cont, Type(*ct, Nullable));
   }
 };
 

src/ir/cost.h

@@ -764,6 +764,10 @@ struct CostAnalyzer : public OverriddenVisitor<CostAnalyzer, CostType> {
     return 8 + visit(curr->ref) + visit(curr->start) + visit(curr->end);
   }
 
+  CostType visitContNew(ContNew* curr) {
+    // Some arbitrary "high" value, reflecting that this may allocate a stack
+    return 14 + visit(curr->func);
+  }
   CostType visitContBind(ContBind* curr) {
     // Inspired by struct.new: The only cost of cont.bind is that it may need to
     // allocate a buffer to hold the arguments.
@@ -774,9 +778,12 @@ struct CostAnalyzer : public OverriddenVisitor<CostAnalyzer, CostType> {
     }
     return ret;
   }
-  CostType visitContNew(ContNew* curr) {
-    // Some arbitrary "high" value, reflecting that this may allocate a stack
-    return 14 + visit(curr->func);
+  CostType visitSuspend(Suspend* curr) {
+    CostType ret = 12;
+    for (auto* arg : curr->operands) {
+      ret += visit(arg);
+    }
+    return ret;
   }
   CostType visitResume(Resume* curr) {
     // Inspired by indirect calls, but twice the cost.
@@ -786,8 +793,17 @@ struct CostAnalyzer : public OverriddenVisitor<CostAnalyzer, CostType> {
     }
     return ret;
   }
-  CostType visitSuspend(Suspend* curr) {
-    CostType ret = 12;
+  CostType visitResumeThrow(ResumeThrow* curr) {
+    // Inspired by indirect calls, but twice the cost.
+    CostType ret = 12 + visit(curr->cont);
+    for (auto* arg : curr->operands) {
+      ret += visit(arg);
+    }
+    return ret;
+  }
+  CostType visitStackSwitch(StackSwitch* curr) {
+    // Inspired by indirect calls, but twice the cost.
+    CostType ret = 12 + visit(curr->cont);
     for (auto* arg : curr->operands) {
       ret += visit(arg);
     }

src/ir/effects.h

@@ -1049,12 +1049,23 @@ class EffectAnalyzer {
       // traps when ref is null.
       parent.implicitTrap = true;
     }
+    void visitContNew(ContNew* curr) {
+      // traps when curr->func is null ref.
+      parent.implicitTrap = true;
+    }
     void visitContBind(ContBind* curr) {
       // traps when curr->cont is null ref.
       parent.implicitTrap = true;
     }
-    void visitContNew(ContNew* curr) {
-      // traps when curr->func is null ref.
+    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;
+      }
+
+      // A suspend may go unhandled and therefore trap.
       parent.implicitTrap = true;
     }
     void visitResume(Resume* curr) {
@@ -1069,10 +1080,26 @@ 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.
+    void visitResumeThrow(ResumeThrow* curr) {
+      // This acts as a kitchen sink effect.
       parent.calls = true;
+
+      // resume_throw instructions accept nullable continuation
+      // references and trap on null.
+      parent.implicitTrap = true;
+
+      if (parent.features.hasExceptionHandling() && parent.tryDepth == 0) {
+        parent.throws_ = true;
+      }
+    }
+    void visitStackSwitch(StackSwitch* curr) {
+      // This acts as a kitchen sink effect.
+      parent.calls = true;
+
+      // switch instructions accept nullable continuation references
+      // and trap on null.
+      parent.implicitTrap = true;
+
       if (parent.features.hasExceptionHandling() && parent.tryDepth == 0) {
         parent.throws_ = true;
       }

src/ir/module-utils.cpp

@@ -448,12 +448,19 @@ struct CodeScanner
     } else if (auto* set = curr->dynCast<ArraySet>()) {
       info.note(set->ref->type);
     } else if (auto* contBind = curr->dynCast<ContBind>()) {
-      info.note(contBind->contTypeBefore);
-      info.note(contBind->contTypeAfter);
+      info.note(contBind->cont->type);
+      info.note(contBind->type);
     } else if (auto* contNew = curr->dynCast<ContNew>()) {
-      info.note(contNew->contType);
+      info.note(contNew->type);
     } else if (auto* resume = curr->dynCast<Resume>()) {
-      info.note(resume->contType);
+      info.note(resume->cont->type);
+      info.note(resume->type);
+    } else if (auto* resumeThrow = curr->dynCast<ResumeThrow>()) {
+      info.note(resumeThrow->cont->type);
+      info.note(resumeThrow->type);
+    } else if (auto* switch_ = curr->dynCast<StackSwitch>()) {
+      info.note(switch_->cont->type);
+      info.note(switch_->type);
     } else if (Properties::isControlFlowStructure(curr)) {
       info.noteControlFlow(Signature(Type::none, curr->type));
     }