5 files changed, 258 insertions, 43 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;
diff --git a/test/passes/remove-unused-brs_all-features.txt b/test/passes/remove-unused-brs_all-features.txt
index 0c830cbc0..5cf83a0d7 100644
--- a/test/passes/remove-unused-brs_all-features.txt
+++ b/test/passes/remove-unused-brs_all-features.txt
@@ -1,14 +1,15 @@
 (module
- (type $struct (struct (field (ref null $vector))))
  (type $vector (array (mut i32)))
+ (type $struct (struct (field (ref null $vector))))
  (type $i32_=>_none (func (param i32)))
+ (type $none_=>_none (func))
  (type $ref|func|_=>_none (func (param (ref func))))
  (type $none_=>_i32 (func (result i32)))
  (type $none_=>_f64 (func (result f64)))
  (type $i32_=>_funcref (func (param i32) (result funcref)))
  (type $none_=>_ref?|$struct| (func (result (ref null $struct))))
  (import "out" "log" (func $log (param i32)))
- (elem declare func $br_on-to-br $i32_=>_none $none_=>_i32)
+ (elem declare func $br_on-to-br $br_on-to-flow $i32_=>_none $none_=>_i32)
  (func $foo (result (ref null $struct))
   (if (result (ref null $struct))
    (i32.const 1)
@@ -123,4 +124,64 @@
    )
   )
  )
+ (func $br_on-to-flow
+  (drop
+   (block $data (result (ref null data))
+    (drop
+     (ref.func $br_on-to-flow)
+    )
+    (ref.null data)
+   )
+  )
+  (drop
+   (block $datab (result (ref null data))
+    (drop
+     (i31.new
+      (i32.const 1337)
+     )
+    )
+    (ref.null data)
+   )
+  )
+  (drop
+   (block $func (result funcref)
+    (drop
+     (array.new_default_with_rtt $vector
+      (i32.const 2)
+      (rtt.canon $vector)
+     )
+    )
+    (ref.null func)
+   )
+  )
+  (drop
+   (block $funcb (result funcref)
+    (drop
+     (i31.new
+      (i32.const 1337)
+     )
+    )
+    (ref.null func)
+   )
+  )
+  (drop
+   (block $i31 (result (ref null i31))
+    (drop
+     (array.new_default_with_rtt $vector
+      (i32.const 2)
+      (rtt.canon $vector)
+     )
+    )
+    (ref.null i31)
+   )
+  )
+  (drop
+   (block $i31b (result (ref null i31))
+    (drop
+     (ref.func $br_on-to-flow)
+    )
+    (ref.null i31)
+   )
+  )
+ )
 )
diff --git a/test/passes/remove-unused-brs_all-features.wast b/test/passes/remove-unused-brs_all-features.wast
index 4f66f2fdd..4ee676589 100644
--- a/test/passes/remove-unused-brs_all-features.wast
+++ b/test/passes/remove-unused-brs_all-features.wast
@@ -108,4 +108,78 @@
    )
   )
  )
+
+ ;; a br_on of the obviously incorrect kind can just flow out the value as the
+ ;; break is never taken
+ (func $br_on-to-flow
+  ;; brs to data
+  (drop
+   (block $data (result (ref null data))
+    (drop
+     (br_on_data $data
+      (ref.func $br_on-to-flow)
+     )
+    )
+    (ref.null data)
+   )
+  )
+  (drop
+   (block $datab (result (ref null data))
+    (drop
+     (br_on_data $datab
+      (i31.new (i32.const 1337))
+     )
+    )
+    (ref.null data)
+   )
+  )
+  ;; brs to func
+  (drop
+   (block $func (result (ref null func))
+    (drop
+     (br_on_func $func
+      (array.new_default_with_rtt $vector
+       (i32.const 2)
+       (rtt.canon $vector)
+      )
+     )
+    )
+    (ref.null func)
+   )
+  )
+  (drop
+   (block $funcb (result (ref null func))
+    (drop
+     (br_on_func $funcb
+      (i31.new (i32.const 1337))
+     )
+    )
+    (ref.null func)
+   )
+  )
+  ;; brs to i31
+  (drop
+   (block $i31 (result (ref null i31))
+    (drop
+     (br_on_i31 $i31
+      (array.new_default_with_rtt $vector
+       (i32.const 2)
+       (rtt.canon $vector)
+      )
+     )
+    )
+    (ref.null i31)
+   )
+  )
+  (drop
+   (block $i31b (result (ref null i31))
+    (drop
+     (br_on_i31 $i31b
+      (ref.func $br_on-to-flow)
+     )
+    )
+    (ref.null i31)
+   )
+  )
+ )
 )