2 files changed, 33 insertions, 70 deletions

diff --git a/src/ir/gc-type-utils.h b/src/ir/gc-type-utils.h
index 318b1b3a3..6dbb3157a 100644
--- a/src/ir/gc-type-utils.h
+++ b/src/ir/gc-type-utils.h
@@ -39,6 +39,22 @@ enum EvaluationResult {
   SuccessOnlyIfNonNull,
 };
 
+inline EvaluationResult flipEvaluationResult(EvaluationResult result) {
+  switch (result) {
+    case Unknown:
+      return Unknown;
+    case Success:
+      return Failure;
+    case Failure:
+      return Success;
+    case SuccessOnlyIfNull:
+      return SuccessOnlyIfNonNull;
+    case SuccessOnlyIfNonNull:
+      return SuccessOnlyIfNull;
+  }
+  WASM_UNREACHABLE("unexpected result");
+}
+
 // Given the type of a reference and a type to attempt to cast it to, return
 // what we know about the result.
 inline EvaluationResult evaluateCastCheck(Type refType, Type castType) {
@@ -85,67 +101,6 @@ inline EvaluationResult evaluateCastCheck(Type refType, Type castType) {
   return Unknown;
 }
 
-// Given an instruction that checks if the child reference is of a certain kind
-// (like br_on_func checks if it is a function), see if type info lets us
-// determine that at compile time.
-// This ignores nullability - it just checks the kind.
-inline EvaluationResult evaluateKindCheck(Expression* curr) {
-  Kind expected;
-  Expression* child;
-
-  // Some operations flip the condition.
-  bool flip = false;
-
-  if (auto* br = curr->dynCast<BrOn>()) {
-    switch (br->op) {
-      // We don't check nullability here.
-      case BrOnNull:
-      case BrOnNonNull:
-        return Unknown;
-      case BrOnCastFail:
-        flip = true;
-        [[fallthrough]];
-      case BrOnCast: {
-        auto result =
-          GCTypeUtils::evaluateCastCheck(br->ref->type, br->castType);
-        if (result == Success) {
-          return flip ? Failure : Success;
-        } else if (result == Failure) {
-          return flip ? Success : Failure;
-        }
-        return Unknown;
-      }
-      default:
-        WASM_UNREACHABLE("unhandled BrOn");
-    }
-    child = br->ref;
-  } else {
-    WASM_UNREACHABLE("invalid input to evaluateKindCheck");
-  }
-
-  auto childType = child->type;
-
-  Kind actual;
-
-  if (childType == Type::unreachable) {
-    return Unknown;
-  } else if (childType.isFunction()) {
-    actual = Func;
-  } else if (childType.isData()) {
-    actual = Data;
-  } else if (childType.getHeapType() == HeapType::i31) {
-    actual = I31;
-  } else {
-    return Unknown;
-  }
-
-  auto success = actual == expected;
-  if (flip) {
-    success = !success;
-  }
-  return success ? Success : Failure;
-}
-
 } // namespace wasm::GCTypeUtils
 
 #endif // wasm_ir_gc_type_utils_h
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index 4edfb8b5e..a7011abe0 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -712,18 +712,21 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
           return;
         }
 
-        // First, check for a possible null which would prevent all other
-        // optimizations (except for br_on_cast variants).
+        // First, check for a possible null which would prevent optimizations on
+        // null checks.
         // 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).
+        // TODO: Use the fallthrough to determine in more cases that we
+        // definitely have a null.
         auto refType = curr->ref->type;
-        if (refType.isNullable() && curr->op != BrOnCast &&
-            curr->op != BrOnCastFail) {
+        if (refType.isNullable() &&
+            (curr->op == BrOnNull || curr->op == BrOnNonNull)) {
           return;
         }
 
         if (curr->op == BrOnNull) {
+          assert(refType.isNonNullable());
           // This cannot be null, so the br is never taken, and the non-null
           // value flows through.
           replaceCurrent(curr->ref);
@@ -731,6 +734,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
           return;
         }
         if (curr->op == BrOnNonNull) {
+          assert(refType.isNonNullable());
           // This cannot be null, so the br is always taken.
           replaceCurrent(
             Builder(*getModule()).makeBreak(curr->name, curr->ref));
@@ -739,20 +743,24 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         }
 
         // Check if the type is the kind we are checking for.
-        auto result = GCTypeUtils::evaluateKindCheck(curr);
+        auto result = GCTypeUtils::evaluateCastCheck(refType, curr->castType);
+        if (curr->op == BrOnCastFail) {
+          result = GCTypeUtils::flipEvaluationResult(result);
+        }
 
         if (result == GCTypeUtils::Success) {
-          // The type is what we are looking for, so we can switch from BrOn to
-          // a simple br which is always taken.
+          // The cast succeeds, so we can switch from BrOn to a simple br that
+          // is always taken.
           replaceCurrent(
             Builder(*getModule()).makeBreak(curr->name, curr->ref));
           worked = true;
         } else if (result == GCTypeUtils::Failure) {
-          // The type is not what we are looking for, so the branch is never
-          // taken, and the value just flows through.
+          // The cast fails, so the branch is never taken, and the value just
+          // flows through.
           replaceCurrent(curr->ref);
           worked = true;
         }
+        // TODO: Handle SuccessOnlyIfNull and SuccessOnlyIfNonNull.
       }
     } optimizer;