Fix BrOn logic in RemoveUnusedBrs (#4062)

This only moves code around. visitBrOn was in the main part of the pass,
which was incorrect as it could interfere with other work being done there.
Specifically, we have a stack of Ifs there, and if we replace a BrOn with
an If, an assertion was hit. To fix this, run it like sinkBlocks(), in a separate
interleaved phase.

This fixes a bug reported by askeksa-google here:

https://github.com/WebAssembly/gc/issues/226#issuecomment-868739853

2 files changed, 105 insertions, 47 deletions

diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index d0950b485..f80f9ba39 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -381,52 +381,6 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
     //       later down, see visitLocalSet.
   }
 
-  void visitBrOn(BrOn* curr) {
-    // Ignore unreachable BrOns which we cannot improve anyhow.
-    if (curr->type == Type::unreachable) {
-      return;
-    }
-
-    // First, check for a possible null which would prevent all other
-    // optimizations.
-    // 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()) {
-      return;
-    }
-
-    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;
-    }
-    if (curr->op == BrOnNonNull) {
-      // This cannot be null, so the br is always taken.
-      replaceCurrent(Builder(*getModule()).makeBreak(curr->name, curr->ref));
-      anotherCycle = true;
-      return;
-    }
-
-    // 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;
-    }
-  }
-
   // override scan to add a pre and a post check task to all nodes
   static void scan(RemoveUnusedBrs* self, Expression** currp) {
     self->pushTask(visitAny, currp);
@@ -688,6 +642,77 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
     return false;
   }
 
+  // GC-specific optimizations. These are split out from the main code to keep
+  // things as simple as possible.
+  bool optimizeGC(Function* func) {
+    if (!getModule()->features.hasGC()) {
+      return false;
+    }
+
+    struct Optimizer : public PostWalker<Optimizer> {
+      bool worked = false;
+
+      void visitBrOn(BrOn* curr) {
+        // Ignore unreachable BrOns which we cannot improve anyhow.
+        if (curr->type == Type::unreachable) {
+          return;
+        }
+
+        // First, check for a possible null which would prevent all other
+        // optimizations.
+        // 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()) {
+          return;
+        }
+
+        if (curr->op == BrOnNull) {
+          // This cannot be null, so the br is never taken, and the non-null
+          // value flows through.
+          replaceCurrent(curr->ref);
+          worked = true;
+          return;
+        }
+        if (curr->op == BrOnNonNull) {
+          // This cannot be null, so the br is always taken.
+          replaceCurrent(
+            Builder(*getModule()).makeBreak(curr->name, curr->ref));
+          worked = true;
+          return;
+        }
+
+        // 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));
+          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.
+          replaceCurrent(curr->ref);
+          worked = true;
+        }
+      }
+    } optimizer;
+
+    optimizer.setModule(getModule());
+    optimizer.doWalkFunction(func);
+
+    // If we removed any BrOn instructions, that might affect the reachability
+    // of the things they used to break to, so update types.
+    if (optimizer.worked) {
+      ReFinalize().walkFunctionInModule(func, getModule());
+      return true;
+    }
+    return false;
+  }
+
   void doWalkFunction(Function* func) {
     // multiple cycles may be needed
     do {
@@ -720,10 +745,12 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
       if (anotherCycle) {
         ReFinalize().walkFunctionInModule(func, getModule());
       }
-      // sink blocks
       if (sinkBlocks(func)) {
         anotherCycle = true;
       }
+      if (optimizeGC(func)) {
+        anotherCycle = true;
+      }
     } while (anotherCycle);
 
     // thread trivial jumps
diff --git a/test/lit/passes/remove-unused-brs-gc.wast b/test/lit/passes/remove-unused-brs-gc.wast
index 7fe579c5a..e38eccebd 100644
--- a/test/lit/passes/remove-unused-brs-gc.wast
+++ b/test/lit/passes/remove-unused-brs-gc.wast
@@ -51,4 +51,35 @@
    )
   )
  )
+
+ ;; CHECK:      (func $br_on-if (param $0 dataref)
+ ;; CHECK-NEXT:  (block $label
+ ;; CHECK-NEXT:   (drop
+ ;; CHECK-NEXT:    (select (result dataref)
+ ;; CHECK-NEXT:     (local.get $0)
+ ;; CHECK-NEXT:     (local.get $0)
+ ;; CHECK-NEXT:     (i32.const 0)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $br_on-if (param $0 (ref data))
+  (block $label
+   (drop
+    ;; This br is never taken, as the input is non-nullable, so we can remove
+    ;; it. When we do so, we replace it with the if. We should not rescan that
+    ;; if, which has already been walked, as that would hit an assertion.
+    ;;
+    (br_on_null $label
+     ;; This if can also be turned into a select, separately from the above
+     ;; (that is not specifically intended to be tested here).
+     (if (result (ref data))
+      (i32.const 0)
+      (local.get $0)
+      (local.get $0)
+     )
+    )
+   )
+  )
+ )
 )