Fix unreachable br_on_cast parse bug

src/ir/child-typer.h

@@ -914,28 +914,31 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     note(&curr->ref, VarRef{Nullable, VarDefHeapType{VarExactness{0u}, *ht}});
   }
 
-  void visitBrOn(BrOn* curr, std::optional<Type> target = std::nullopt) {
+  void visitBrOn(BrOn* curr,
+                 std::optional<Type> in = std::nullopt,
+                 std::optional<Type> out = std::nullopt) {
     switch (curr->op) {
       case BrOnNull:
       case BrOnNonNull:
         // br_on(_non)_null is polymorphic over reference types and does not
         // take a type immediate.
-        assert(!target);
+        assert(!in && !out);
         // Polymorphic over heap types.
         note(&curr->ref, VarRef{Nullable, VarHeapType{0u}});
         return;
       case BrOnCast:
       case BrOnCastFail:
       case BrOnCastDescEq:
       case BrOnCastDescEqFail: {
-        if (!target) {
+        assert(!!in == !!out);
+        if (!out) {
           assert(curr->castType.isRef());
-          target = curr->castType;
+          out = curr->castType;
+          in = Type(out->getHeapType().getTop(), Nullable);
         }
-        auto top = target->getHeapType().getTop();
-        note(&curr->ref, Type(top, Nullable));
+        note(&curr->ref, *in);
         if (curr->op == BrOnCastDescEq || curr->op == BrOnCastDescEqFail) {
-          auto desc = target->getHeapType().getDescriptorType();
+          auto desc = out->getHeapType().getDescriptorType();
           assert(desc);
           note(&curr->desc, Type(*desc, Nullable));
         }

src/parser/contexts.h

@@ -2800,8 +2800,8 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
                     const std::vector<Annotation>& annotations,
                     Index label,
                     BrOnOp op,
-                    Type in = Type::none,
-                    Type out = Type::none) {
+                    std::optional<Type> in = std::nullopt,
+                    std::optional<Type> out = std::nullopt) {
     return withLoc(
       pos,
       irBuilder.makeBrOn(label, op, in, out, parseAnnotations(annotations)));

src/wasm-ir-builder.h

@@ -233,8 +233,8 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Result<> makeRefGetDesc(HeapType type);
   Result<> makeBrOn(Index label,
                     BrOnOp op,
-                    Type in = Type::none,
-                    Type out = Type::none,
+                    std::optional<Type> in = std::nullopt,
+                    std::optional<Type> out = std::nullopt,
                     const CodeAnnotation& annotations = {});
   Result<> makeStructNew(HeapType type, bool isDesc);
   Result<> makeStructNewDefault(HeapType type, bool isDesc);

src/wasm/wasm-ir-builder.cpp

@@ -23,6 +23,7 @@
 #include "ir/properties.h"
 #include "ir/utils.h"
 #include "wasm-ir-builder.h"
+#include "wasm.h"
 
 #ifndef IR_BUILDER_DEBUG
 #define IR_BUILDER_DEBUG 0
@@ -674,6 +675,14 @@ struct IRBuilder::ChildPopper
     return popConstrainedChildren(children);
   }
 
+  Result<> visitBrOn(BrOn* curr,
+                     std::optional<Type> in = std::nullopt,
+                     std::optional<Type> out = std::nullopt) {
+    std::vector<Child> children;
+    ConstraintCollector{builder, children}.visitBrOn(curr, in, out);
+    return popConstrainedChildren(children);
+  }
+
   Result<> visitSwitch(Switch* curr,
                        std::optional<Type> labelType = std::nullopt) {
     std::vector<Child> children;
@@ -2087,26 +2096,29 @@ Result<> IRBuilder::makeRefGetDesc(HeapType type) {
 
 Result<> IRBuilder::makeBrOn(Index label,
                              BrOnOp op,
-                             Type in,
-                             Type out,
+                             std::optional<Type> in,
+                             std::optional<Type> out,
                              const CodeAnnotation& annotations) {
+  bool isCast = op != BrOnNull && op != BrOnNonNull;
+  bool isDescCast = op == BrOnCastDescEq || op == BrOnCastDescEqFail;
+  assert(bool(in) == bool(out) && bool(in) == isCast);
   std::optional<HeapType> descriptor;
-  if (op == BrOnCastDescEq || op == BrOnCastDescEqFail) {
-    assert(out.isRef());
-    descriptor = out.getHeapType().getDescriptorType();
+  if (isDescCast) {
+    assert(out && out->isRef());
+    descriptor = out->getHeapType().getDescriptorType();
     if (!descriptor) {
       return Err{"cast target must have descriptor"};
     }
   }
 
   BrOn curr;
   curr.op = op;
-  curr.castType = out;
+  curr.castType = isCast ? *out : Type::none;
   curr.desc = nullptr;
-  if (op != BrOnNull && op != BrOnNonNull && !out.isCastable()) {
+  if (isCast && !out->isCastable()) {
     return Err{"br_on cannot cast to invalid type"};
   }
-  CHECK_ERR(visitBrOn(&curr));
+  CHECK_ERR(ChildPopper{*this}.visitBrOn(&curr, in, out));
 
   // Validate things that would cause errors later.
   if (curr.ref->type != Type::unreachable && !curr.ref->type.isRef()) {
@@ -2124,14 +2136,14 @@ Result<> IRBuilder::makeBrOn(Index label,
       break;
     case BrOnCastDescEq:
     case BrOnCastDescEqFail: {
-      CHECK_ERR(validateTypeAnnotation(out.with(*descriptor).with(Nullable),
+      CHECK_ERR(validateTypeAnnotation(out->with(*descriptor).with(Nullable),
                                        curr.desc));
     }
       [[fallthrough]];
     case BrOnCast:
     case BrOnCastFail:
-      assert(in.isRef());
-      CHECK_ERR(validateTypeAnnotation(in, curr.ref));
+      assert(in->isRef());
+      CHECK_ERR(validateTypeAnnotation(*in, curr.ref));
   }
 
   // Extra values need to be sent in a scratch local.
@@ -2168,7 +2180,7 @@ Result<> IRBuilder::makeBrOn(Index label,
     case BrOnCastFail:
     case BrOnCastDescEq:
     case BrOnCastDescEqFail:
-      testType = in;
+      testType = *in;
       break;
   }
 
@@ -2181,7 +2193,7 @@ Result<> IRBuilder::makeBrOn(Index label,
     auto name = getLabelName(label);
     CHECK_ERR(name);
 
-    auto* br = builder.makeBrOn(op, *name, curr.ref, out, curr.desc);
+    auto* br = builder.makeBrOn(op, *name, curr.ref, curr.castType, curr.desc);
     applyAnnotations(br, annotations);
     push(br);
     return Ok{};
@@ -2205,7 +2217,8 @@ Result<> IRBuilder::makeBrOn(Index label,
 
   // Perform the branch.
   CHECK_ERR(visitBrOn(&curr));
-  auto* br = builder.makeBrOn(op, extraLabel, curr.ref, out, curr.desc);
+  auto* br =
+    builder.makeBrOn(op, extraLabel, curr.ref, curr.castType, curr.desc);
   applyAnnotations(br, annotations);
   push(br);
 
@@ -2227,18 +2240,18 @@ Result<> IRBuilder::makeBrOn(Index label,
       WASM_UNREACHABLE("unexpected op");
     case BrOnCast:
     case BrOnCastDescEq:
-      if (out.isNullable()) {
-        resultType = Type(in.getHeapType(), NonNullable);
+      if (out->isNullable()) {
+        resultType = Type(in->getHeapType(), NonNullable);
       } else {
-        resultType = in;
+        resultType = *in;
       }
       break;
     case BrOnCastFail:
     case BrOnCastDescEqFail:
-      if (in.isNonNullable()) {
-        resultType = Type(out.getHeapType(), NonNullable);
+      if (in->isNonNullable()) {
+        resultType = Type(out->getHeapType(), NonNullable);
       } else {
-        resultType = out;
+        resultType = *out;
       }
       break;
   }

test/lit/basic/unreachable-br-on-cast-desc.wast

@@ -0,0 +1,41 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+
+;; RUN: wasm-opt %s -all -S -o - | filecheck %s
+
+;; Regression test for a parser bug where we would incorrectly pop a
+;; nullable `ref` operand for an unreachable br_on_cast_desc_eq{_fail} that
+;; expected a non-nullable `ref`, then immediately fail the relevant type check.
+
+(module
+  (rec
+    ;; CHECK:      (rec
+    ;; CHECK-NEXT:  (type $s (descriptor $d) (struct))
+    (type $s (descriptor $d) (struct))
+    ;; CHECK:       (type $d (describes $s) (struct))
+    (type $d (describes $s) (struct))
+  )
+  ;; CHECK:      (func $unreachable-br-on-cast-desc (type $2) (result (ref $s))
+  ;; CHECK-NEXT:  (block $label (result (ref $s))
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (block ;; (replaces unreachable BrOn we can't emit)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (unreachable)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (unreachable)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $unreachable-br-on-cast-desc (result (ref $s))
+    (br_on_cast_desc_eq 0 (ref $s) (ref $s)
+      ;; The input type is expected to be non-nullable, so this cannot be the
+      ;; input to the cast. We should create a new `unreachable` instead.
+      (ref.null none)
+      (unreachable)
+    )
+  )
+)