3 files changed, 31 insertions, 32 deletions

diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 91309c906..133272a67 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -7034,14 +7034,10 @@ bool WasmBinaryBuilder::maybeVisitBrOn(Expression*& out, uint32_t code) {
     case BinaryConsts::BrOnNonNull:
       op = BrOnNonNull;
       break;
-    case BinaryConsts::BrOnCastStatic:
     case BinaryConsts::BrOnCast:
-    case BinaryConsts::BrOnCastNull:
       op = BrOnCast;
       break;
-    case BinaryConsts::BrOnCastStaticFail:
     case BinaryConsts::BrOnCastFail:
-    case BinaryConsts::BrOnCastFailNull:
       op = BrOnCastFail;
       break;
     case BinaryConsts::BrOnFunc:
@@ -7063,18 +7059,24 @@ bool WasmBinaryBuilder::maybeVisitBrOn(Expression*& out, uint32_t code) {
     default:
       return false;
   }
-  auto name = getBreakTarget(getU32LEB()).name;
-  if (castType == Type::none && (op == BrOnCast || op == BrOnCastFail)) {
-    auto nullability = (code == BinaryConsts::BrOnCastNull ||
-                        code == BinaryConsts::BrOnCastFailNull)
-                         ? Nullable
-                         : NonNullable;
-    bool legacy = code == BinaryConsts::BrOnCastStatic ||
-                  code == BinaryConsts::BrOnCastStaticFail;
-    auto type = legacy ? getIndexedHeapType() : getHeapType();
-    castType = Type(type, nullability);
+  uint8_t flags = 0;
+  if (op == BrOnCast || op == BrOnCastFail) {
+    flags = getInt8();
   }
+  auto name = getBreakTarget(getU32LEB()).name;
   auto* ref = popNonVoidExpression();
+  if (op == BrOnCast || op == BrOnCastFail) {
+    auto inputNullability = (flags & 1) ? Nullable : NonNullable;
+    auto castNullability = (flags & 2) ? Nullable : NonNullable;
+    auto inputHeapType = getHeapType();
+    auto castHeapType = getHeapType();
+    auto inputType = Type(inputHeapType, inputNullability);
+    castType = Type(castHeapType, castNullability);
+    if (!Type::isSubType(ref->type, inputType)) {
+      throwError(std::string("Invalid reference type for ") +
+                 ((op == BrOnCast) ? "br_on_cast" : "br_on_cast_fail"));
+    }
+  }
   out = Builder(wasm).makeBrOn(op, name, ref, castType);
   return true;
 }
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index 6adfb6149..a51ad3716 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -2892,16 +2892,16 @@ Expression* SExpressionWasmBuilder::makeBrOnCast(Element& s,
                                                  bool onFail) {
   int i = 1;
   auto name = getLabel(*s[i++]);
+  std::optional<Type> inputType;
   if (!castType) {
-    auto nullability = NonNullable;
-    if (s[i]->str().str == "null") {
-      nullability = Nullable;
-      ++i;
-    }
-    auto type = parseHeapType(*s[i++]);
-    castType = Type(type, nullability);
+    inputType = elementToType(*s[i++]);
+    castType = elementToType(*s[i++]);
   }
   auto* ref = parseExpression(*s[i]);
+  if (inputType && !Type::isSubType(ref->type, *inputType)) {
+    throw ParseException(
+      "br_on_cast* ref type does not match expected type", s.line, s.col);
+  }
   auto op = onFail ? BrOnCastFail : BrOnCast;
   return Builder(wasm).makeBrOn(op, name, ref, *castType);
 }
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index 7a4dd3983..ee457da69 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -2045,25 +2045,22 @@ void BinaryInstWriter::visitBrOn(BrOn* curr) {
       o << U32LEB(getBreakIndex(curr->name));
       return;
     case BrOnCast:
+    case BrOnCastFail: {
       o << int8_t(BinaryConsts::GCPrefix);
-      if (curr->castType.isNullable()) {
-        o << U32LEB(BinaryConsts::BrOnCastNull);
-      } else {
+      if (curr->op == BrOnCast) {
         o << U32LEB(BinaryConsts::BrOnCast);
-      }
-      o << U32LEB(getBreakIndex(curr->name));
-      parent.writeHeapType(curr->castType.getHeapType());
-      return;
-    case BrOnCastFail:
-      o << int8_t(BinaryConsts::GCPrefix);
-      if (curr->castType.isNullable()) {
-        o << U32LEB(BinaryConsts::BrOnCastFailNull);
       } else {
         o << U32LEB(BinaryConsts::BrOnCastFail);
       }
+      assert(curr->ref->type.isRef());
+      uint8_t flags = (curr->ref->type.isNullable() ? 1 : 0) |
+                      (curr->castType.isNullable() ? 2 : 0);
+      o << flags;
       o << U32LEB(getBreakIndex(curr->name));
+      parent.writeHeapType(curr->ref->type.getHeapType());
       parent.writeHeapType(curr->castType.getHeapType());
       return;
+    }
   }
   WASM_UNREACHABLE("invalid br_on_*");
 }