[Wasm GC] Optimize BrOn* of the wrong kind (#3724)

#3719 optimized the case where the kind is what we want, like br_on_func of
a function. This handles the opposite case, where we know the kind is wrong, and
so the break is not taken.

This seems equally useful for "polymorphic" code that does a bunch of checks
and routes to different code for each case, as after inlining or other optimizations
we may see which paths are taken and which are not.

Also refactor the checking code to a shared location as RefIs/As will also want to
use it.

3 files changed, 121 insertions, 41 deletions

diff --git a/src/ir/gc-type-utils.h b/src/ir/gc-type-utils.h
new file mode 100644
index 000000000..168c9529e
--- /dev/null
+++ b/src/ir/gc-type-utils.h
@@ -0,0 +1,94 @@
+/*
+ * Copyright 2021 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_ir_gc_type_utils_h
+#define wasm_ir_gc_type_utils_h
+
+#include "wasm.h"
+
+namespace wasm {
+
+namespace GCTypeUtils {
+
+// Helper code to evaluate a reference at compile time and check if it is of a
+// certain kind. Various wasm instructions check if something is a function or
+// data etc., and that code is shared here.
+
+enum Kind { Func, Data, I31, Other };
+
+enum EvaluationResult {
+  // The result is not known at compile time.
+  Unknown,
+  // The evaluation is known to succeed (i.e., we find what we are looking
+  // for), or fail, at compile time.
+  Success,
+  Failure
+};
+
+// 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.
+EvaluationResult evaluateKindCheck(Expression* curr) {
+  Kind expected;
+  Expression* child;
+
+  if (auto* br = curr->dynCast<BrOn>()) {
+    switch (br->op) {
+      // We don't check nullability here.
+      case BrOnNull:
+      // Casts can only be known at runtime using RTTs.
+      case BrOnCast:
+        return Unknown;
+      case BrOnFunc:
+        expected = Func;
+        break;
+      case BrOnData:
+        expected = Data;
+        break;
+      case BrOnI31:
+        expected = I31;
+        break;
+      default:
+        WASM_UNREACHABLE("unhandled BrOn");
+    }
+    child = br->ref;
+  } else {
+    WASM_UNREACHABLE("invalid input to evaluateKindCheck");
+  }
+
+  auto childType = child->type;
+
+  Kind actual;
+
+  if (childType.isFunction()) {
+    actual = Func;
+  } else if (childType.isData()) {
+    actual = Data;
+  } else if (childType.getHeapType() == HeapType::i31) {
+    actual = I31;
+  } else {
+    return Unknown;
+  }
+
+  return actual == expected ? Success : Failure;
+}
+
+} // namespace GCTypeUtils
+
+} // namespace wasm
+
+#endif // wasm_ir_gc_type_utils_h
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index 3bf56838e..cacd24402 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -21,6 +21,7 @@
 #include <ir/branch-utils.h>
 #include <ir/cost.h>
 #include <ir/effects.h>
+#include <ir/gc-type-utils.h>
 #include <ir/literal-utils.h>
 #include <ir/utils.h>
 #include <parsing.h>
@@ -386,51 +387,38 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
       return;
     }
 
-    // If the type provides enough information we may be able to know if this
-    // br is taken or not. If so, the br_on* may be unneeded. First, check for a
-    // possible null which would prevent such an optimization.
+    // First, check for a possible null which would prevent all other
+    // optimizations.
+    // (Note: if the spec had BrOnNonNull, instead of BrOnNull, then we could
+    // replace a br_on_func whose input is (ref null func) with br_on_non_null,
+    // as only the null check would be needed. But as things are, we cannot do
+    // such a thing.)
     auto refType = curr->ref->type;
     if (refType.isNullable()) {
       return;
     }
 
-    // Nulls are not possible, so specialization may be achievable, either
-    // removing the br_on* entirely or replacing it with a br.
-    auto replaceWithBr = [&]() {
-      replaceCurrent(Builder(*getModule()).makeBreak(curr->name, curr->ref));
+    if (curr->op == BrOnNull) {
+      // This cannot be null, so the br is never taken, and the non-null value
+      // flows through.
+      replaceCurrent(curr->ref);
       anotherCycle = true;
-    };
+      return;
+    }
 
-    switch (curr->op) {
-      case BrOnNull: {
-        // This cannot be null, so the br is never taken, and the non-null value
-        // flows through.
-        replaceCurrent(curr->ref);
-        anotherCycle = true;
-        break;
-      }
-      case BrOnCast: {
-        // Casts can only be done at runtime, using RTTs.
-        break;
-      }
-      case BrOnFunc: {
-        if (refType.isFunction()) {
-          replaceWithBr();
-        }
-        break;
-      }
-      case BrOnData: {
-        if (refType.isData()) {
-          replaceWithBr();
-        }
-        break;
-      }
-      case BrOnI31: {
-        if (refType.getHeapType() == HeapType::i31) {
-          replaceWithBr();
-        }
-        break;
-      }
+    // Check if the type is the kind we are checking for.
+    auto result = GCTypeUtils::evaluateKindCheck(curr);
+
+    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.
+      replaceCurrent(Builder(*getModule()).makeBreak(curr->name, curr->ref));
+      anotherCycle = 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.
+      replaceCurrent(curr->ref);
+      anotherCycle = true;
     }
   }
 
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index ae38e7f5c..a7040b959 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -276,9 +276,7 @@ Type markTemp(Type type) {
   return type;
 }
 
-bool isTemp(Type type) {
-  return !type.isBasic() && getTypeInfo(type)->isTemp;
-}
+bool isTemp(Type type) { return !type.isBasic() && getTypeInfo(type)->isTemp; }
 
 bool isTemp(HeapType type) {
   return !type.isBasic() && getHeapTypeInfo(type)->isTemp;