Support br_on_cast null (#5397)

As well as br_on_cast_fail null. Unlike the existing br_on_cast* instructions,
these new instructions treat the cast as succeeding when the input is a null.
Update the internal representation of the cast type in `BrOn` expressions to be
a `Type` rather than a `HeapType` so it will include nullability information.

Also update and improve `RemoveUnusedBrs` to handle the new instructions
correctly and optimize in more cases.

20 files changed, 231 insertions, 76 deletions

diff --git a/src/binaryen-c.cpp b/src/binaryen-c.cpp
index b26ec3389..34a45540d 100644
--- a/src/binaryen-c.cpp
+++ b/src/binaryen-c.cpp
@@ -1757,12 +1757,10 @@ BinaryenExpressionRef BinaryenBrOn(BinaryenModuleRef module,
                                    BinaryenOp op,
                                    const char* name,
                                    BinaryenExpressionRef ref,
-                                   BinaryenHeapType intendedType) {
-  Builder builder(*(Module*)module);
+                                   BinaryenType castType) {
   return static_cast<Expression*>(
-    intendedType ? builder.makeBrOn(
-                     BrOnOp(op), name, (Expression*)ref, HeapType(intendedType))
-                 : builder.makeBrOn(BrOnOp(op), name, (Expression*)ref));
+    Builder(*(Module*)module)
+      .makeBrOn(BrOnOp(op), name, (Expression*)ref, Type(castType)));
 }
 BinaryenExpressionRef BinaryenStructNew(BinaryenModuleRef module,
                                         BinaryenExpressionRef* operands,
@@ -4119,16 +4117,16 @@ void BinaryenBrOnSetRef(BinaryenExpressionRef expr,
   assert(refExpr);
   static_cast<BrOn*>(expression)->ref = (Expression*)refExpr;
 }
-BinaryenHeapType BinaryenBrOnGetIntendedType(BinaryenExpressionRef expr) {
+BinaryenType BinaryenBrOnGetCastType(BinaryenExpressionRef expr) {
   auto* expression = (Expression*)expr;
   assert(expression->is<BrOn>());
-  return static_cast<BrOn*>(expression)->intendedType.getID();
+  return static_cast<BrOn*>(expression)->castType.getID();
 }
-void BinaryenBrOnSetIntendedType(BinaryenExpressionRef expr,
-                                 BinaryenHeapType intendedType) {
+void BinaryenBrOnSetCastType(BinaryenExpressionRef expr,
+                             BinaryenType castType) {
   auto* expression = (Expression*)expr;
   assert(expression->is<BrOn>());
-  static_cast<BrOn*>(expression)->intendedType = HeapType(intendedType);
+  static_cast<BrOn*>(expression)->castType = Type(castType);
 }
 // StructNew
 BinaryenIndex BinaryenStructNewGetNumOperands(BinaryenExpressionRef expr) {
diff --git a/src/binaryen-c.h b/src/binaryen-c.h
index 7894c1c24..357e1f18d 100644
--- a/src/binaryen-c.h
+++ b/src/binaryen-c.h
@@ -1048,7 +1048,7 @@ BINARYEN_API BinaryenExpressionRef BinaryenBrOn(BinaryenModuleRef module,
                                                 BinaryenOp op,
                                                 const char* name,
                                                 BinaryenExpressionRef ref,
-                                                BinaryenHeapType intendedType);
+                                                BinaryenType castType);
 BINARYEN_API BinaryenExpressionRef
 BinaryenStructNew(BinaryenModuleRef module,
                   BinaryenExpressionRef* operands,
@@ -2384,10 +2384,9 @@ BINARYEN_API BinaryenExpressionRef
 BinaryenBrOnGetRef(BinaryenExpressionRef expr);
 BINARYEN_API void BinaryenBrOnSetRef(BinaryenExpressionRef expr,
                                      BinaryenExpressionRef refExpr);
-BINARYEN_API BinaryenHeapType
-BinaryenBrOnGetIntendedType(BinaryenExpressionRef expr);
-BINARYEN_API void BinaryenBrOnSetIntendedType(BinaryenExpressionRef expr,
-                                              BinaryenHeapType intendedType);
+BINARYEN_API BinaryenType BinaryenBrOnGetCastType(BinaryenExpressionRef expr);
+BINARYEN_API void BinaryenBrOnSetCastType(BinaryenExpressionRef expr,
+                                          BinaryenType castType);
 
 // StructNew
 
diff --git a/src/ir/gc-type-utils.h b/src/ir/gc-type-utils.h
index abe5f2dfc..5e3de37e5 100644
--- a/src/ir/gc-type-utils.h
+++ b/src/ir/gc-type-utils.h
@@ -52,16 +52,20 @@ inline EvaluationResult evaluateKindCheck(Expression* curr) {
       // We don't check nullability here.
       case BrOnNull:
       case BrOnNonNull:
+        return Unknown;
       case BrOnCastFail:
         flip = true;
         [[fallthrough]];
       case BrOnCast:
-        // Note that the type must be non-nullable for us to succeed since a
-        // null would make us fail.
-        if (Type::isSubType(br->ref->type,
-                            Type(br->intendedType, NonNullable))) {
+        // If we already have a subtype of the cast type, the cast will succeed.
+        if (Type::isSubType(br->ref->type, br->castType)) {
           return flip ? Failure : Success;
         }
+        // If the cast type is unrelated to the type we have, the cast will
+        // certainly fail.
+        if (!Type::isSubType(br->castType, br->ref->type)) {
+          return flip ? Success : Failure;
+        }
         return Unknown;
       case BrOnNonFunc:
         flip = true;
diff --git a/src/ir/module-utils.cpp b/src/ir/module-utils.cpp
index 0cdff0955..9dfc5f123 100644
--- a/src/ir/module-utils.cpp
+++ b/src/ir/module-utils.cpp
@@ -77,7 +77,7 @@ struct CodeScanner
       counts.note(cast->castType);
     } else if (auto* cast = curr->dynCast<BrOn>()) {
       if (cast->op == BrOnCast || cast->op == BrOnCastFail) {
-        counts.note(cast->intendedType);
+        counts.note(cast->castType);
       }
     } else if (auto* get = curr->dynCast<StructGet>()) {
       counts.note(get->ref->type);
diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index 81e3ae989..c1064b7db 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -2143,13 +2143,19 @@ struct PrintExpressionContents
         printMedium(o, "br_on_cast ");
         printName(curr->name, o);
         o << ' ';
-        printHeapType(o, curr->intendedType, wasm);
+        if (curr->castType.isNullable()) {
+          printMedium(o, "null ");
+        }
+        printHeapType(o, curr->castType.getHeapType(), wasm);
         return;
       case BrOnCastFail:
         printMedium(o, "br_on_cast_fail ");
         printName(curr->name, o);
         o << ' ';
-        printHeapType(o, curr->intendedType, wasm);
+        if (curr->castType.isNullable()) {
+          printMedium(o, "null ");
+        }
+        printHeapType(o, curr->castType.getHeapType(), wasm);
         return;
       case BrOnFunc:
         printMedium(o, "br_on_func ");
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index f465d46c1..4edfb8b5e 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -713,12 +713,13 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         }
 
         // First, check for a possible null which would prevent all other
-        // optimizations.
+        // optimizations (except for br_on_cast variants).
         // TODO: Look into using BrOnNonNull here, to replace a br_on_func whose
         // input is (ref null func) with br_on_non_null (as only the null check
         // would be needed).
         auto refType = curr->ref->type;
-        if (refType.isNullable()) {
+        if (refType.isNullable() && curr->op != BrOnCast &&
+            curr->op != BrOnCastFail) {
           return;
         }
 
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index 996d9ab4c..303c81128 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -1127,8 +1127,8 @@ enum ASTNodes {
   BrOnCastStaticFail = 0x47,
   RefTestNull = 0x48,
   RefCastNull = 0x49,
-  // TODO: BrOnCastNull
-  // TODO: BrOnCastFailNull
+  BrOnCastNull = 0x4a,
+  BrOnCastFailNull = 0x4b,
   RefCastNop = 0x4c,
   RefIsFunc = 0x50,
   RefIsData = 0x51,
diff --git a/src/wasm-builder.h b/src/wasm-builder.h
index a83a0c0c2..49c9bb0ba 100644
--- a/src/wasm-builder.h
+++ b/src/wasm-builder.h
@@ -899,12 +899,12 @@ public:
     ret->finalize();
     return ret;
   }
-  BrOn* makeBrOn(BrOnOp op, Name name, Expression* ref, HeapType intendedType) {
+  BrOn* makeBrOn(BrOnOp op, Name name, Expression* ref, Type castType) {
     auto* ret = wasm.allocator.alloc<BrOn>();
     ret->op = op;
     ret->name = name;
     ret->ref = ref;
-    ret->intendedType = intendedType;
+    ret->castType = castType;
     ret->finalize();
     return ret;
   }
diff --git a/src/wasm-delegations-fields.def b/src/wasm-delegations-fields.def
index 9c966faec..f929def4a 100644
--- a/src/wasm-delegations-fields.def
+++ b/src/wasm-delegations-fields.def
@@ -629,7 +629,7 @@ switch (DELEGATE_ID) {
     DELEGATE_START(BrOn);
     DELEGATE_FIELD_INT(BrOn, op);
     DELEGATE_FIELD_SCOPE_NAME_USE(BrOn, name);
-    DELEGATE_FIELD_HEAPTYPE(BrOn, intendedType);
+    DELEGATE_FIELD_TYPE(BrOn, castType);
     DELEGATE_FIELD_CHILD(BrOn, ref);
     DELEGATE_END(BrOn);
     break;
diff --git a/src/wasm.h b/src/wasm.h
index f1dcd7b46..db10fe260 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -1553,12 +1553,11 @@ public:
   Name name;
   Expression* ref;
 
-  HeapType intendedType;
+  Type castType;
 
   void finalize();
 
-  // TODO: Support br_on_cast* null as well.
-  Type getCastType() { return Type(intendedType, NonNullable); }
+  Type getCastType() { return castType; }
 
   // Returns the type sent on the branch, if it is taken.
   Type getSentType();
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index c6fb361bd..61fd65eab 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -6931,6 +6931,7 @@ bool WasmBinaryBuilder::maybeVisitRefCast(Expression*& out, uint32_t code) {
 
 bool WasmBinaryBuilder::maybeVisitBrOn(Expression*& out, uint32_t code) {
   BrOnOp op;
+  auto nullability = NonNullable;
   switch (code) {
     case BinaryConsts::BrOnNull:
       op = BrOnNull;
@@ -6946,6 +6947,14 @@ bool WasmBinaryBuilder::maybeVisitBrOn(Expression*& out, uint32_t code) {
     case BinaryConsts::BrOnCastFail:
       op = BrOnCastFail;
       break;
+    case BinaryConsts::BrOnCastNull:
+      op = BrOnCast;
+      nullability = Nullable;
+      break;
+    case BinaryConsts::BrOnCastFailNull:
+      op = BrOnCastFail;
+      nullability = Nullable;
+      break;
     case BinaryConsts::BrOnFunc:
       op = BrOnFunc;
       break;
@@ -6968,14 +6977,15 @@ bool WasmBinaryBuilder::maybeVisitBrOn(Expression*& out, uint32_t code) {
       return false;
   }
   auto name = getBreakTarget(getU32LEB()).name;
-  HeapType intendedType;
+  Type castType = Type::none;
   if (op == BrOnCast || op == BrOnCastFail) {
     bool legacy = code == BinaryConsts::BrOnCastStatic ||
                   code == BinaryConsts::BrOnCastStaticFail;
-    intendedType = legacy ? getIndexedHeapType() : getHeapType();
+    auto type = legacy ? getIndexedHeapType() : getHeapType();
+    castType = Type(type, nullability);
   }
   auto* ref = popNonVoidExpression();
-  out = Builder(wasm).makeBrOn(op, name, ref, intendedType);
+  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 2f791609a..fd3cbcc23 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -2821,12 +2821,18 @@ Expression* SExpressionWasmBuilder::makeRefCastNop(Element& s) {
 Expression* SExpressionWasmBuilder::makeBrOn(Element& s, BrOnOp op) {
   int i = 1;
   auto name = getLabel(*s[i++]);
-  HeapType heapType;
+  Type castType = Type::none;
   if (op == BrOnCast || op == BrOnCastFail) {
-    heapType = parseHeapType(*s[i++]);
+    auto nullability = NonNullable;
+    if (s[i]->str().str == "null") {
+      nullability = Nullable;
+      ++i;
+    }
+    auto type = parseHeapType(*s[i++]);
+    castType = Type(type, nullability);
   }
   auto* ref = parseExpression(*s[i]);
-  return Builder(wasm).makeBrOn(op, name, ref, heapType);
+  return Builder(wasm).makeBrOn(op, name, ref, castType);
 }
 
 Expression* SExpressionWasmBuilder::makeStructNew(Element& s, bool default_) {
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index e24060f2a..fbf73b7c2 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -2058,11 +2058,19 @@ void BinaryInstWriter::visitBrOn(BrOn* curr) {
       break;
     case BrOnCast:
       o << int8_t(BinaryConsts::GCPrefix);
-      o << U32LEB(BinaryConsts::BrOnCast);
+      if (curr->castType.isNullable()) {
+        o << U32LEB(BinaryConsts::BrOnCastNull);
+      } else {
+        o << U32LEB(BinaryConsts::BrOnCast);
+      }
       break;
     case BrOnCastFail:
       o << int8_t(BinaryConsts::GCPrefix);
-      o << U32LEB(BinaryConsts::BrOnCastFail);
+      if (curr->castType.isNullable()) {
+        o << U32LEB(BinaryConsts::BrOnCastFailNull);
+      } else {
+        o << U32LEB(BinaryConsts::BrOnCastFail);
+      }
       break;
     case BrOnFunc:
       o << int8_t(BinaryConsts::GCPrefix) << U32LEB(BinaryConsts::BrOnFunc);
@@ -2087,7 +2095,7 @@ void BinaryInstWriter::visitBrOn(BrOn* curr) {
   }
   o << U32LEB(getBreakIndex(curr->name));
   if (curr->op == BrOnCast || curr->op == BrOnCastFail) {
-    parent.writeHeapType(curr->intendedType);
+    parent.writeHeapType(curr->castType.getHeapType());
   }
 }
 
diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
index bec9b0ad0..486f13d24 100644
--- a/src/wasm/wasm-validator.cpp
+++ b/src/wasm/wasm-validator.cpp
@@ -2554,14 +2554,19 @@ void FunctionValidator::visitBrOn(BrOn* curr) {
     return;
   }
   if (curr->op == BrOnCast || curr->op == BrOnCastFail) {
+    if (!shouldBeTrue(curr->castType.isRef(),
+                      curr,
+                      "br_on_cast must have reference cast type")) {
+      return;
+    }
     shouldBeEqual(
-      curr->intendedType.getBottom(),
+      curr->castType.getHeapType().getBottom(),
       curr->ref->type.getHeapType().getBottom(),
       curr,
       "br_on_cast* target type and ref type must have a common supertype");
   } else {
-    shouldBeEqual(curr->intendedType,
-                  HeapType(),
+    shouldBeEqual(curr->castType,
+                  Type(Type::none),
                   curr,
                   "non-cast br_on* must not set intendedType field");
   }
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index 682936461..de18966ff 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -968,18 +968,12 @@ void BrOn::finalize() {
       // us flow out the null, but it does not).
       type = Type::none;
       break;
-    case BrOnCast:
     case BrOnFunc:
     case BrOnData:
     case BrOnI31:
       // If we do not branch, we return the input in this case.
       type = ref->type;
       break;
-    case BrOnCastFail:
-      // If we do not branch, the cast worked, and we have something of the cast
-      // type.
-      type = Type(intendedType, NonNullable);
-      break;
     case BrOnNonFunc:
       type = Type(HeapType::func, NonNullable);
       break;
@@ -989,6 +983,26 @@ void BrOn::finalize() {
     case BrOnNonI31:
       type = Type(HeapType::i31, NonNullable);
       break;
+    case BrOnCast:
+      if (castType.isNullable()) {
+        // Nulls take the branch, so the result is non-nullable.
+        type = Type(ref->type.getHeapType(), NonNullable);
+      } else {
+        // Nulls do not take the branch, so the result is non-nullable only if
+        // the input is.
+        type = ref->type;
+      }
+      break;
+    case BrOnCastFail:
+      if (castType.isNullable()) {
+        // Nulls do not take the branch, so the result is non-nullable only if
+        // the input is.
+        type = Type(castType.getHeapType(), ref->type.getNullability());
+      } else {
+        // Nulls take the branch, so the result is non-nullable.
+        type = castType;
+      }
+      break;
     default:
       WASM_UNREACHABLE("invalid br_on_*");
   }
@@ -1007,22 +1021,30 @@ Type BrOn::getSentType() {
       }
       // BrOnNonNull sends the non-nullable type on the branch.
       return Type(ref->type.getHeapType(), NonNullable);
-    case BrOnCast:
-      if (ref->type == Type::unreachable) {
-        return Type::unreachable;
-      }
-      return Type(intendedType, NonNullable);
     case BrOnFunc:
       return Type(HeapType::func, NonNullable);
     case BrOnData:
       return Type(HeapType::data, NonNullable);
     case BrOnI31:
       return Type(HeapType::i31, NonNullable);
-    case BrOnCastFail:
     case BrOnNonFunc:
     case BrOnNonData:
     case BrOnNonI31:
       return ref->type;
+    case BrOnCast:
+      // The same as the result type of br_on_cast_fail.
+      if (castType.isNullable()) {
+        return Type(castType.getHeapType(), ref->type.getNullability());
+      } else {
+        return castType;
+      }
+    case BrOnCastFail:
+      // The same as the result type of br_on_cast.
+      if (castType.isNullable()) {
+        return Type(ref->type.getHeapType(), NonNullable);
+      } else {
+        return ref->type;
+      }
     default:
       WASM_UNREACHABLE("invalid br_on_*");
   }
diff --git a/test/lit/cast-to-basic.wast b/test/lit/cast-to-basic.wast
index c6b8ea9e7..1aa329bf5 100644
--- a/test/lit/cast-to-basic.wast
+++ b/test/lit/cast-to-basic.wast
@@ -56,11 +56,11 @@
     )
   )
 
-  ;; CHECK:      (func $br-fail (type $none_=>_none)
+  ;; CHECK:      (func $br-null (type $none_=>_none)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block $label$1 (result dataref)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (br_on_cast_fail $label$1 data
+  ;; CHECK-NEXT:     (br_on_cast $label$1 null data
   ;; CHECK-NEXT:      (ref.null none)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
@@ -68,11 +68,36 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
-  (func $br-fail
+  (func $br-null
     (drop
       (block $l (result structref)
         (drop
-          (br_on_cast_fail $l struct
+          (br_on_cast $l null struct
+            (ref.null none)
+          )
+        )
+        (ref.null none)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $br-fail-null (type $none_=>_none)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block $label$1 (result dataref)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (br_on_cast_fail $label$1 null data
+  ;; CHECK-NEXT:      (ref.null none)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $br-fail-null
+    (drop
+      (block $l (result structref)
+        (drop
+          (br_on_cast_fail $l null struct
             (ref.null none)
           )
         )
diff --git a/test/lit/passes/remove-unused-brs-gc.wast b/test/lit/passes/remove-unused-brs-gc.wast
index c448a47c2..5c45b9e7c 100644
--- a/test/lit/passes/remove-unused-brs-gc.wast
+++ b/test/lit/passes/remove-unused-brs-gc.wast
@@ -3,8 +3,13 @@
 ;; RUN:  | filecheck %s
 
 (module
- ;; CHECK:      (type $struct (struct ))
- (type $struct (struct ))
+ (rec
+  ;; CHECK:      (rec
+  ;; CHECK-NEXT:  (type $struct (struct ))
+  (type $struct (struct))
+  ;; CHECK:       (type $struct2 (struct ))
+  (type $struct2 (struct))
+ )
 
  ;; CHECK:      (func $br_on_non_i31-1 (type $none_=>_none)
  ;; CHECK-NEXT:  (drop
@@ -117,7 +122,6 @@
  )
 
  ;; CHECK:      (func $br_on_cast (type $none_=>_ref|$struct|) (result (ref $struct))
- ;; CHECK-NEXT:  (local $temp (ref null $struct))
  ;; CHECK-NEXT:  (block $block (result (ref $struct))
  ;; CHECK-NEXT:   (drop
  ;; CHECK-NEXT:    (br $block
@@ -128,7 +132,6 @@
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT: )
  (func $br_on_cast (result (ref $struct))
-  (local $temp (ref null $struct))
   (block $block (result (ref $struct))
    (drop
     ;; This static cast can be computed at compile time: it will definitely be
@@ -141,23 +144,66 @@
   )
  )
 
+ ;; CHECK:      (func $br_on_cast_unrelated (type $none_=>_ref|$struct|) (result (ref $struct))
+ ;; CHECK-NEXT:  (block $block
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (struct.new_default $struct2)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (unreachable)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $br_on_cast_unrelated (result (ref $struct))
+  (block $block (result (ref $struct))
+   (drop
+    ;; This cast can be computed at compile time: it will definitely fail, so we
+    ;; can remove it.
+    (br_on_cast $block $struct
+     (struct.new $struct2)
+    )
+   )
+   (unreachable)
+  )
+ )
+
  ;; CHECK:      (func $br_on_cast_no (type $none_=>_ref|$struct|) (result (ref $struct))
- ;; CHECK-NEXT:  (local $temp (ref null $struct))
+ ;; CHECK-NEXT:  (local $struct (ref null $struct))
  ;; CHECK-NEXT:  (block $block (result (ref $struct))
  ;; CHECK-NEXT:   (drop
  ;; CHECK-NEXT:    (br_on_cast $block $struct
- ;; CHECK-NEXT:     (ref.null none)
+ ;; CHECK-NEXT:     (local.get $struct)
  ;; CHECK-NEXT:    )
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:   (unreachable)
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT: )
  (func $br_on_cast_no (result (ref $struct))
-  (local $temp (ref null $struct))
+  (local $struct (ref null $struct))
   (block $block (result (ref $struct))
    (drop
     (br_on_cast $block $struct
      ;; As above, but now the type is nullable, so we cannot infer anything.
+     (local.get $struct)
+    )
+   )
+   (unreachable)
+  )
+ )
+
+ ;; CHECK:      (func $br_on_cast_nullable (type $none_=>_ref?|$struct|) (result (ref null $struct))
+ ;; CHECK-NEXT:  (block $block (result nullref)
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (br $block
+ ;; CHECK-NEXT:     (ref.null none)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (unreachable)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $br_on_cast_nullable (result (ref null $struct))
+  (block $block (result (ref null $struct))
+   (drop
+    (br_on_cast $block null $struct
+     ;; As above, but now the cast allows nulls, so we can optimize.
      (ref.null $struct)
     )
    )
@@ -166,7 +212,6 @@
  )
 
  ;; CHECK:      (func $br_on_cast_fail (type $none_=>_ref|$struct|) (result (ref $struct))
- ;; CHECK-NEXT:  (local $temp (ref null $struct))
  ;; CHECK-NEXT:  (block $block
  ;; CHECK-NEXT:   (drop
  ;; CHECK-NEXT:    (struct.new_default $struct)
@@ -175,7 +220,6 @@
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT: )
  (func $br_on_cast_fail (result (ref $struct))
-  (local $temp (ref null $struct))
   (block $block (result (ref $struct))
    (drop
     ;; As $br_on_cast, but this checks for a failing cast, so we know it will
@@ -189,6 +233,7 @@
  )
 
  ;; CHECK:      (func $casts-are-costly (type $i32_=>_none) (param $x i32)
+ ;; CHECK-NEXT:  (local $struct (ref null $struct))
  ;; CHECK-NEXT:  (drop
  ;; CHECK-NEXT:   (if (result i32)
  ;; CHECK-NEXT:    (local.get $x)
@@ -213,7 +258,7 @@
  ;; CHECK-NEXT:    (block $something (result anyref)
  ;; CHECK-NEXT:     (drop
  ;; CHECK-NEXT:      (br_on_cast $something $struct
- ;; CHECK-NEXT:       (ref.null none)
+ ;; CHECK-NEXT:       (local.get $struct)
  ;; CHECK-NEXT:      )
  ;; CHECK-NEXT:     )
  ;; CHECK-NEXT:     (ref.null none)
@@ -242,6 +287,7 @@
   ;; We never turn an if into a select if an arm has a cast of any kind, as
   ;; those things involve branches internally, so we'd be adding more than we
   ;; save.
+  (local $struct (ref null $struct))
   (drop
    (if (result i32)
     (local.get $x)
@@ -267,7 +313,7 @@
      (block $something (result anyref)
       (drop
        (br_on_cast $something $struct
-        (ref.null $struct)
+        (local.get $struct)
        )
       )
       (ref.null any)
diff --git a/test/passes/Oz_fuzz-exec_all-features.txt b/test/passes/Oz_fuzz-exec_all-features.txt
index af706cc23..30949def4 100644
--- a/test/passes/Oz_fuzz-exec_all-features.txt
+++ b/test/passes/Oz_fuzz-exec_all-features.txt
@@ -52,7 +52,7 @@
 [LoggingExternalInterface logging 0]
 [LoggingExternalInterface logging 1]
 [LoggingExternalInterface logging 0]
-[LoggingExternalInterface logging 0]
+[LoggingExternalInterface logging 1]
 [fuzz-exec] calling static-br_on_cast
 [LoggingExternalInterface logging 3]
 [fuzz-exec] calling static-br_on_cast_fail
@@ -62,7 +62,7 @@
  (type $void_func (func))
  (type $struct (struct (field (mut i32))))
  (type $i32_=>_none (func (param i32)))
- (type $extendedstruct (struct (field (mut i32)) (field f64)))
+ (type $extendedstruct (struct_subtype (field (mut i32)) (field f64) $struct))
  (type $int_func (func (result i32)))
  (import "fuzzing-support" "log-i32" (func $log (param i32)))
  (export "structs" (func $0))
@@ -329,7 +329,7 @@
    (i32.const 0)
   )
   (call $log
-   (i32.const 0)
+   (i32.const 1)
   )
  )
  (func $21 (type $void_func) (; has Stack IR ;)
@@ -391,7 +391,7 @@
 [LoggingExternalInterface logging 0]
 [LoggingExternalInterface logging 1]
 [LoggingExternalInterface logging 0]
-[LoggingExternalInterface logging 0]
+[LoggingExternalInterface logging 1]
 [fuzz-exec] calling static-br_on_cast
 [LoggingExternalInterface logging 3]
 [fuzz-exec] calling static-br_on_cast_fail
diff --git a/test/passes/Oz_fuzz-exec_all-features.wast b/test/passes/Oz_fuzz-exec_all-features.wast
index bf6767c73..0ce79aa4b 100644
--- a/test/passes/Oz_fuzz-exec_all-features.wast
+++ b/test/passes/Oz_fuzz-exec_all-features.wast
@@ -1,6 +1,6 @@
 (module
  (type $struct (struct (mut i32)))
- (type $extendedstruct (struct (mut i32) f64))
+ (type $extendedstruct (struct_subtype (mut i32) f64 $struct))
  (type $bytes (array (mut i8)))
 
  (type $void_func (func))
diff --git a/test/spec/ref_cast.wast b/test/spec/ref_cast.wast
index 704063ec6..b81671e30 100644
--- a/test/spec/ref_cast.wast
+++ b/test/spec/ref_cast.wast
@@ -103,6 +103,18 @@
     (i32.const 1)
   )
 
+  (func (export "test-br-on-cast-null-struct") (result i32)
+    (drop
+      (block $l (result (ref null struct))
+        (drop
+          (br_on_cast $l null struct (ref.null none))
+        )
+        (return (i32.const 0))
+      )
+    )
+    (i32.const 1)
+  )
+
   (func (export "test-br-on-cast-fail-struct") (result i32)
     (drop
       (block $l (result (ref struct))
@@ -115,6 +127,18 @@
     (i32.const 1)
   )
 
+  (func (export "test-br-on-cast-fail-null-struct") (result i32)
+    (drop
+      (block $l (result (ref struct))
+        (drop
+          (br_on_cast_fail $l null struct (ref.null none))
+        )
+        (return (i32.const 0))
+      )
+    )
+    (i32.const 1)
+  )
+
   (func (export "test-trap-null")
     (drop
       (ref.cast $t0
@@ -131,7 +155,9 @@
 (assert_return (invoke "test-ref-test-any") (i32.const 1))
 (assert_return (invoke "test-ref-cast-struct"))
 (assert_return (invoke "test-br-on-cast-struct") (i32.const 1))
+(assert_return (invoke "test-br-on-cast-null-struct") (i32.const 1))
 (assert_return (invoke "test-br-on-cast-fail-struct") (i32.const 0))
+(assert_return (invoke "test-br-on-cast-fail-null-struct") (i32.const 0))
 (assert_trap (invoke "test-trap-null"))
 
 (assert_invalid