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;