[custom-descriptors] Branching descriptor casts

scripts/gen-s-parser.py

@@ -613,8 +613,10 @@
     ("ref.get_desc",         "makeRefGetDesc()"),
     ("br_on_null",           "makeBrOnNull()"),
     ("br_on_non_null",       "makeBrOnNull(true)"),
-    ("br_on_cast",           "makeBrOnCast()"),
-    ("br_on_cast_fail",      "makeBrOnCast(true)"),
+    ("br_on_cast",           "makeBrOnCast(BrOnCast)"),
+    ("br_on_cast_fail",      "makeBrOnCast(BrOnCastFail)"),
+    ("br_on_cast_desc",      "makeBrOnCast(BrOnCastDesc)"),
+    ("br_on_cast_desc_fail", "makeBrOnCast(BrOnCastDescFail)"),
     ("struct.new",           "makeStructNew(false)"),
     ("struct.new_default",   "makeStructNew(true)"),
     ("struct.get",           "makeStructGet()"),

scripts/test/fuzzing.py

@@ -113,6 +113,8 @@
     'vacuum-stack-switching.wast',
     # TODO: fuzzer support for custom descriptors
     'custom-descriptors.wast',
+    'br_on_cast_desc.wast',
+    'ref.get_cast.wast',
     # TODO: fix split_wast() on tricky escaping situations like a string ending
     #       in \\" (the " is not escaped - there is an escaped \ before it)
     'string-lifting-section.wast',

src/gen-s-parser.inc

@@ -209,16 +209,38 @@ switch (buf[0]) {
                 switch (buf[10]) {
                   case '\0':
                     if (op == "br_on_cast"sv) {
-                      CHECK_ERR(makeBrOnCast(ctx, pos, annotations));
+                      CHECK_ERR(makeBrOnCast(ctx, pos, annotations, BrOnCast));
                       return Ok{};
                     }
                     goto parse_error;
-                  case '_':
-                    if (op == "br_on_cast_fail"sv) {
-                      CHECK_ERR(makeBrOnCast(ctx, pos, annotations, true));
-                      return Ok{};
+                  case '_': {
+                    switch (buf[11]) {
+                      case 'd': {
+                        switch (buf[15]) {
+                          case '\0':
+                            if (op == "br_on_cast_desc"sv) {
+                              CHECK_ERR(makeBrOnCast(ctx, pos, annotations, BrOnCastDesc));
+                              return Ok{};
+                            }
+                            goto parse_error;
+                          case '_':
+                            if (op == "br_on_cast_desc_fail"sv) {
+                              CHECK_ERR(makeBrOnCast(ctx, pos, annotations, BrOnCastDescFail));
+                              return Ok{};
+                            }
+                            goto parse_error;
+                          default: goto parse_error;
+                        }
+                      }
+                      case 'f':
+                        if (op == "br_on_cast_fail"sv) {
+                          CHECK_ERR(makeBrOnCast(ctx, pos, annotations, BrOnCastFail));
+                          return Ok{};
+                        }
+                        goto parse_error;
+                      default: goto parse_error;
                     }
-                    goto parse_error;
+                  }
                   default: goto parse_error;
                 }
               }

src/ir/ReFinalize.cpp

@@ -151,7 +151,7 @@ void ReFinalize::visitRefGetDesc(RefGetDesc* curr) { curr->finalize(); }
 void ReFinalize::visitBrOn(BrOn* curr) {
   curr->finalize();
   if (curr->type == Type::unreachable) {
-    replaceUntaken(curr->ref, nullptr);
+    replaceUntaken(curr->ref, curr->desc);
   } else {
     updateBreakValueType(curr->name, curr->getSentType());
   }
@@ -218,11 +218,11 @@ void ReFinalize::updateBreakValueType(Name name, Type type) {
 }
 
 // Replace an untaken branch/switch with an unreachable value.
-// A condition may also exist and may or may not be unreachable.
-void ReFinalize::replaceUntaken(Expression* value, Expression* condition) {
+// Another child may also exist and may or may not be unreachable.
+void ReFinalize::replaceUntaken(Expression* value, Expression* otherChild) {
   assert(value->type == Type::unreachable);
   auto* replacement = value;
-  if (condition) {
+  if (otherChild) {
     Builder builder(*getModule());
     // Even if we have
     //  (block
@@ -233,10 +233,10 @@ void ReFinalize::replaceUntaken(Expression* value, Expression* condition) {
     // the value is unreachable, and necessary since the type of
     // the condition did not have an impact before (the break/switch
     // type was unreachable), and might not fit in.
-    if (condition->type.isConcrete()) {
-      condition = builder.makeDrop(condition);
+    if (otherChild->type.isConcrete()) {
+      otherChild = builder.makeDrop(otherChild);
     }
-    replacement = builder.makeSequence(value, condition);
+    replacement = builder.makeSequence(value, otherChild);
     assert(replacement->type.isBasic() && "Basic type expected");
   }
   replaceCurrent(replacement);

src/ir/child-typer.h

@@ -865,16 +865,25 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     note(&curr->ref, Type(*ht, Nullable));
   }
 
-  void visitBrOn(BrOn* curr) {
+  void visitBrOn(BrOn* curr, std::optional<Type> target = std::nullopt) {
     switch (curr->op) {
       case BrOnNull:
       case BrOnNonNull:
         noteAnyReference(&curr->ref);
         return;
       case BrOnCast:
-      case BrOnCastFail: {
-        auto top = curr->castType.getHeapType().getTop();
+      case BrOnCastFail:
+      case BrOnCastDesc:
+      case BrOnCastDescFail: {
+        if (!target) {
+          target = curr->castType;
+        }
+        auto top = target->getHeapType().getTop();
         note(&curr->ref, Type(top, Nullable));
+        if (curr->op == BrOnCastDesc || curr->op == BrOnCastDescFail) {
+          auto descriptor = *target->getHeapType().getDescriptorType();
+          note(&curr->desc, Type(descriptor, Nullable));
+        }
         return;
       }
     }

src/ir/cost.h

@@ -673,9 +673,20 @@ struct CostAnalyzer : public OverriddenVisitor<CostAnalyzer, CostType> {
   }
   CostType visitBrOn(BrOn* curr) {
     // BrOn of a null can be fairly fast, but anything else is a cast check.
-    CostType base =
-      curr->op == BrOnNull || curr->op == BrOnNonNull ? 2 : CastCost;
-    return base + nullCheckCost(curr->ref) + maybeVisit(curr->ref);
+    switch (curr->op) {
+      case BrOnNull:
+      case BrOnNonNull:
+        return 2 + nullCheckCost(curr->ref) + visit(curr->ref);
+      case BrOnCast:
+      case BrOnCastFail:
+        return CastCost + visit(curr->ref);
+      case BrOnCastDesc:
+      case BrOnCastDescFail:
+        // These are not as expensive as full casts, since they just do a
+        // identity check on the descriptor.
+        return 2 + visit(curr->ref) + visit(curr->desc);
+    }
+    WASM_UNREACHABLE("unexpected op");
   }
   CostType visitStructNew(StructNew* curr) {
     CostType ret = AllocationCost + curr->operands.size();

src/ir/properties.h

@@ -521,7 +521,7 @@ inline bool hasUnwritableTypeImmediate(Expression* curr) {
 #define DELEGATE_ID curr->_id
 
 #define DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD(id, field)                        \
-  {                                                                            \
+  if (curr->cast<id>()->field) {                                               \
     auto type = curr->cast<id>()->field->type;                                 \
     if (type == Type::unreachable || type.isNull()) {                          \
       return true;                                                             \

src/parser/parsers.h

@@ -238,8 +238,7 @@ template<typename Ctx>
 Result<>
 makeBrOnNull(Ctx&, Index, const std::vector<Annotation>&, bool onFail = false);
 template<typename Ctx>
-Result<>
-makeBrOnCast(Ctx&, Index, const std::vector<Annotation>&, bool onFail = false);
+Result<> makeBrOnCast(Ctx&, Index, const std::vector<Annotation>&, BrOnOp op);
 template<typename Ctx>
 Result<>
 makeStructNew(Ctx&, Index, const std::vector<Annotation>&, bool default_);
@@ -2251,15 +2250,14 @@ template<typename Ctx>
 Result<> makeBrOnCast(Ctx& ctx,
                       Index pos,
                       const std::vector<Annotation>& annotations,
-                      bool onFail) {
+                      BrOnOp op) {
   auto label = labelidx(ctx);
   CHECK_ERR(label);
   auto in = reftype(ctx);
   CHECK_ERR(in);
   auto out = reftype(ctx);
   CHECK_ERR(out);
-  return ctx.makeBrOn(
-    pos, annotations, *label, onFail ? BrOnCastFail : BrOnCast, *in, *out);
+  return ctx.makeBrOn(pos, annotations, *label, op, *in, *out);
 }
 
 template<typename Ctx>

src/passes/Print.cpp

@@ -2226,7 +2226,26 @@ struct PrintExpressionContents
         curr->name.print(o);
         return;
       case BrOnCast:
-        printMedium(o, "br_on_cast ");
+      case BrOnCastDesc:
+      case BrOnCastFail:
+      case BrOnCastDescFail:
+        switch (curr->op) {
+          case BrOnCast:
+            printMedium(o, "br_on_cast");
+            break;
+          case BrOnCastFail:
+            printMedium(o, "br_on_cast_fail");
+            break;
+          case BrOnCastDesc:
+            printMedium(o, "br_on_cast_desc");
+            break;
+          case BrOnCastDescFail:
+            printMedium(o, "br_on_cast_desc_fail");
+            break;
+          default:
+            WASM_UNREACHABLE("unexpected op");
+        }
+        o << ' ';
         curr->name.print(o);
         o << ' ';
         if (curr->ref->type == Type::unreachable) {
@@ -2240,8 +2259,9 @@ struct PrintExpressionContents
         o << ' ';
         printType(curr->castType);
         return;
-      case BrOnCastFail:
-        printMedium(o, "br_on_cast_fail ");
+        printMedium(o,
+                    curr->op == BrOnCastFail ? "br_on_cast_fail "
+                                             : "br_on_cast_desc_fail ");
         curr->name.print(o);
         o << ' ';
         if (curr->ref->type == Type::unreachable) {

src/wasm-binary.h

@@ -1174,6 +1174,8 @@ enum ASTNodes {
   RefCastNull = 0x17,
   BrOnCast = 0x18,
   BrOnCastFail = 0x19,
+  BrOnCastDesc = 0x25,
+  BrOnCastDescFail = 0x26,
   AnyConvertExtern = 0x1a,
   ExternConvertAny = 0x1b,
   RefI31 = 0x1c,

src/wasm-builder.h

@@ -900,12 +900,18 @@ class Builder {
     ret->finalize();
     return ret;
   }
-  BrOn*
-  makeBrOn(BrOnOp op, Name name, Expression* ref, Type castType = Type::none) {
+  BrOn* makeBrOn(BrOnOp op,
+                 Name name,
+                 Expression* ref,
+                 Type castType = Type::none,
+                 Expression* desc = nullptr) {
+    assert((desc && (op == BrOnCastDesc || op == BrOnCastDescFail)) ||
+           (!desc && op != BrOnCastDesc && op != BrOnCastDescFail));
     auto* ret = wasm.allocator.alloc<BrOn>();
     ret->op = op;
     ret->name = name;
     ret->ref = ref;
+    ret->desc = desc;
     ret->castType = castType;
     ret->finalize();
     return ret;

src/wasm-delegations-fields.def

@@ -46,6 +46,14 @@
 //    present (like a Return's value). If you do not define this then
 //    DELEGATE_FIELD_CHILD is called.
 //
+// DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD(id, field) - The previous two
+// cases combined. If you do not define this, but you define exactly one of
+// DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD and DELEGATE_FIELD_OPTIONAL_CHILD, that
+// defined macro will be called. Defining both
+// DELEGATE_FIELD_IMMEDIATE_TYPE_CHILD and DELEGATE_FIELD_OPTIONAL_CHILD without
+// defining DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD is an error. If
+// neither of the other macros are defined, then DELEGATE_FIELD_CHILD is called.
+//
 // DELEGATE_FIELD_CHILD_VECTOR(id, field) - called for a variable-sized vector
 //   of  child pointers. If this is not defined, and DELEGATE_GET_FIELD is, then
 //    DELEGATE_FIELD_CHILD is called on them.
@@ -56,7 +64,7 @@
 // DELEGATE_FIELD_INT_ARRAY(id, field) - called for a std::array of fixed size
 //    of integer values (like a SIMD mask). If this is not defined, and
 //    DELEGATE_GET_FIELD is, then DELEGATE_FIELD_INT is called on them.
-
+//
 // DELEGATE_FIELD_INT_VECTOR(id, field) - called for a variable-sized vector
 //    of integer values. If this is not defined, and DELEGATE_GET_FIELD is, then
 //    DELEGATE_FIELD_INT is called on them.
@@ -113,6 +121,18 @@
 #error please define DELEGATE_FIELD_CHILD(id, field)
 #endif
 
+#ifndef DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD
+#if defined(DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD) && !defined(DELEGATE_FIELD_OPTIONAL_CHILD)
+#define DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD(id, field) DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD(id, field)
+#elif !defined(DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD) && defined(DELEGATE_FIELD_OPTIONAL_CHILD)
+#define DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD(id, field) DELEGATE_FIELD_OPTIONAL_CHILD(id, field)
+#elif !defined(DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD) && !defined(DELEGATE_FIELD_OPTIONAL_CHILD)
+#define DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD(id, field) DELEGATE_FIELD_CHILD(id, field)
+#else
+#error please define DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD(id, field)
+#endif
+#endif
+
 #ifndef DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD
 #define DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD(id, field) DELEGATE_FIELD_CHILD(id, field)
 #endif
@@ -641,6 +661,7 @@ DELEGATE_FIELD_CASE_START(BrOn)
 DELEGATE_FIELD_INT(BrOn, op)
 DELEGATE_FIELD_SCOPE_NAME_USE(BrOn, name)
 DELEGATE_FIELD_TYPE(BrOn, castType)
+DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD(BrOn, desc)
 DELEGATE_FIELD_CHILD(BrOn, ref)
 DELEGATE_FIELD_CASE_END(BrOn)
 
@@ -844,6 +865,7 @@ DELEGATE_FIELD_MAIN_END
 #undef DELEGATE_FIELD_CHILD
 #undef DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD
 #undef DELEGATE_FIELD_OPTIONAL_CHILD
+#undef DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD
 #undef DELEGATE_FIELD_CHILD_VECTOR
 #undef DELEGATE_FIELD_INT
 #undef DELEGATE_FIELD_INT_ARRAY