Stack switching proposal support

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).

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)"),

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':
@@ -5076,6 +5087,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/ir/ReFinalize.cpp

@@ -180,8 +180,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

@@ -89,6 +89,13 @@ void operateOnScopeNameUsesAndSentTypes(Expression* expr, T func) {
           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 == name) {
+          func(name, r->sentTypes[i]);
+        }
+      }
     } else {
       assert(expr->is<Try>() || expr->is<Rethrow>()); // delegate or rethrow
     }
@@ -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

@@ -1050,6 +1050,10 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     note(&curr->end, Type::i32);
   }
 
+  void visitContNew(ContNew* curr) {
+    note(&curr->func, Type(curr->contType.getContinuation().type, Nullable));
+  }
+
   void visitContBind(ContBind* curr) {
     auto paramsBefore =
       curr->contTypeBefore.getContinuation().type.getSignature().params;
@@ -1064,8 +1068,12 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     note(&curr->cont, Type(curr->contTypeBefore, Nullable));
   }
 
-  void visitContNew(ContNew* curr) {
-    note(&curr->func, Type(curr->contType.getContinuation().type, Nullable));
+  void visitSuspend(Suspend* curr) {
+    auto params = wasm.getTag(curr->tag)->sig.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) {
@@ -1077,12 +1085,23 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     note(&curr->cont, Type(curr->contType, Nullable));
   }
 
-  void visitSuspend(Suspend* curr) {
+  void visitResumeThrow(ResumeThrow* curr) {
     auto params = wasm.getTag(curr->tag)->sig.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));
+  }
+
+  void visitStackSwitch(StackSwitch* curr) {
+    auto params = curr->contType.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(curr->contType, Nullable));
   }
 };
 

src/ir/cost.h

@@ -756,6 +756,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.
@@ -766,9 +770,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.
@@ -778,8 +785,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

@@ -1008,13 +1008,21 @@ 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.
-      parent.implicitTrap = 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;
+      }
     }
     void visitResume(Resume* curr) {
       // This acts as a kitchen sink effect.
@@ -1028,10 +1036,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/possible-contents.cpp

@@ -1225,19 +1225,27 @@ struct InfoCollector
 
   void visitReturn(Return* curr) { addResult(curr->value); }
 
+  void visitContNew(ContNew* curr) {
+    // TODO: optimize when possible
+    addRoot(curr);
+  }
   void visitContBind(ContBind* curr) {
     // TODO: optimize when possible
     addRoot(curr);
   }
-  void visitContNew(ContNew* curr) {
+  void visitSuspend(Suspend* curr) {
     // TODO: optimize when possible
     addRoot(curr);
   }
   void visitResume(Resume* curr) {
     // TODO: optimize when possible
     addRoot(curr);
   }
-  void visitSuspend(Suspend* curr) {
+  void visitResumeThrow(ResumeThrow* curr) {
+    // TODO: optimize when possible
+    addRoot(curr);
+  }
+  void visitStackSwitch(StackSwitch* curr) {
     // TODO: optimize when possible
     addRoot(curr);
   }

src/ir/subtype-exprs.h

@@ -397,10 +397,16 @@ struct SubtypingDiscoverer : public OverriddenVisitor<SubType> {
   void visitStringWTF16Get(StringWTF16Get* curr) {}
   void visitStringSliceWTF(StringSliceWTF* curr) {}
 
-  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 visitContBind(ContBind* curr) { WASM_UNREACHABLE("not implemented"); }
   void visitSuspend(Suspend* curr) { WASM_UNREACHABLE("not implemented"); }
+  void visitResume(Resume* curr) { WASM_UNREACHABLE("not implemented"); }
+  void visitResumeThrow(ResumeThrow* curr) {
+    WASM_UNREACHABLE("not implemented");
+  }
+  void visitStackSwitch(StackSwitch* curr) {
+    WASM_UNREACHABLE("not implemented");
+  }
 };
 
 } // namespace wasm