RemoveUnusedBrs::tablify() improvements: handle EqZ and tee (#4144)

tablify() attempts to turns a sequence of br_ifs into a single
br_table. This PR adds some flexibility to the specific pattern it
looks for, specifically:

  *  Accept i32.eqz as a comparison to zero, and not just to look
     for i32.eq against a constant.
  *  Allow the first condition to be a tee. If it is, compare later
     conditions to local.get of that local.

This will allow more br_tables to be emitted in j2cl output.

4 files changed, 231 insertions, 9 deletions

diff --git a/scripts/fuzz_opt.py b/scripts/fuzz_opt.py
index 9ccdeb1e4..17d6301dd 100755
--- a/scripts/fuzz_opt.py
+++ b/scripts/fuzz_opt.py
@@ -172,6 +172,7 @@ IMPORTANT_INITIAL_CONTENTS = [
     # Recently-added or modified passes. These can be added to and pruned
     # frequently.
     os.path.join('lit', 'passes', 'once-reduction.wast'),
+    os.path.join('passes', 'remove-unused-brs_enable-multivalue.wast'),
 ]
 IMPORTANT_INITIAL_CONTENTS = [os.path.join(shared.get_test_dir('.'), t) for t in IMPORTANT_INITIAL_CONTENTS]
 
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index ea8a66612..9861695e0 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -1316,7 +1316,15 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
             // dce will clean it up
             return nullptr;
           }
-          auto* binary = br->condition->dynCast<Binary>();
+          auto* condition = br->condition;
+          // Also support eqz, which is the same as == 0.
+          if (auto* unary = condition->dynCast<Unary>()) {
+            if (unary->op == EqZInt32) {
+              return br;
+            }
+            return nullptr;
+          }
+          auto* binary = condition->dynCast<Binary>();
           if (!binary) {
             return nullptr;
           }
@@ -1342,16 +1350,29 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
           if (!br) {
             return nullptr;
           }
-          return br->condition->cast<Binary>()->left;
+          auto* condition = br->condition;
+          if (auto* binary = condition->dynCast<Binary>()) {
+            return binary->left;
+          } else if (auto* unary = condition->dynCast<Unary>()) {
+            assert(unary->op == EqZInt32);
+            return unary->value;
+          } else {
+            WASM_UNREACHABLE("invalid br_if condition");
+          }
         };
 
         // returns the constant value, as a uint32_t
         auto getProperBrIfConstant =
           [&getProperBrIf](Expression* curr) -> uint32_t {
-          return getProperBrIf(curr)
-            ->condition->cast<Binary>()
-            ->right->cast<Const>()
-            ->value.geti32();
+          auto* condition = getProperBrIf(curr)->condition;
+          if (auto* binary = condition->dynCast<Binary>()) {
+            return binary->right->cast<Const>()->value.geti32();
+          } else if (auto* unary = condition->dynCast<Unary>()) {
+            assert(unary->op == EqZInt32);
+            return 0;
+          } else {
+            WASM_UNREACHABLE("invalid br_if condition");
+          }
         };
         Index start = 0;
         while (start < list.size() - 1) {
@@ -1360,6 +1381,14 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
             start++;
             continue;
           }
+          // If the first condition value is a tee, that is ok, so long as the
+          // others afterwards are gets of the value that is tee'd.
+          LocalGet get;
+          if (auto* tee = conditionValue->dynCast<LocalSet>()) {
+            get.index = tee->index;
+            get.type = getFunction()->getLocalType(get.index);
+            conditionValue = &get;
+          }
           // if the condition has side effects, we can't replace many
           // appearances of it with a single one
           if (EffectAnalyzer(passOptions, *getModule(), conditionValue)
@@ -1427,13 +1456,15 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
               }
               Builder builder(*getModule());
               // the table and condition are offset by the min
+              auto* newCondition = getProperBrIfConditionValue(list[start]);
+
               if (min != 0) {
-                conditionValue = builder.makeBinary(
-                  SubInt32, conditionValue, builder.makeConst(int32_t(min)));
+                newCondition = builder.makeBinary(
+                  SubInt32, newCondition, builder.makeConst(int32_t(min)));
               }
               list[end - 1] = builder.makeBlock(
                 defaultName,
-                builder.makeSwitch(table, defaultName, conditionValue));
+                builder.makeSwitch(table, defaultName, newCondition));
               for (Index i = start; i < end - 1; i++) {
                 ExpressionManipulator::nop(list[i]);
               }
diff --git a/test/passes/remove-unused-brs_enable-multivalue.txt b/test/passes/remove-unused-brs_enable-multivalue.txt
index 29b5eca5e..46938f784 100644
--- a/test/passes/remove-unused-brs_enable-multivalue.txt
+++ b/test/passes/remove-unused-brs_enable-multivalue.txt
@@ -1669,6 +1669,109 @@
   )
   (unreachable)
  )
+ (func $br-to-table-initial-tee (param $a i32)
+  (block $x
+   (block $y
+    (block $z
+     (nop)
+     (nop)
+     (block $tablify|0
+      (br_table $x $y $z $tablify|0
+       (i32.sub
+        (local.tee $a
+         (i32.add
+          (i32.const 10)
+          (i32.const 1)
+         )
+        )
+        (i32.const 10)
+       )
+      )
+     )
+     (unreachable)
+    )
+    (unreachable)
+   )
+   (unreachable)
+  )
+  (unreachable)
+ )
+ (func $br-to-table-initial-tee-wrong-index (param $a i32)
+  (local $b i32)
+  (block $x
+   (block $y
+    (block $z
+     (br_if $x
+      (i32.eq
+       (local.tee $a
+        (i32.const 99)
+       )
+       (i32.const 10)
+      )
+     )
+     (br_if $y
+      (i32.eq
+       (local.get $b)
+       (i32.const 11)
+      )
+     )
+     (br_if $z
+      (i32.eq
+       (local.get $b)
+       (i32.const 12)
+      )
+     )
+     (unreachable)
+    )
+    (unreachable)
+   )
+   (unreachable)
+  )
+  (unreachable)
+ )
+ (func $br-to-table-eqz (param $a i32)
+  (block $x
+   (block $y
+    (block $z
+     (nop)
+     (nop)
+     (block $tablify|0
+      (br_table $x $y $z $tablify|0
+       (local.get $a)
+      )
+     )
+     (unreachable)
+    )
+    (unreachable)
+   )
+   (unreachable)
+  )
+  (unreachable)
+ )
+ (func $br-to-table-tee-eqz (param $a i32)
+  (block $x
+   (block $y
+    (block $z
+     (nop)
+     (nop)
+     (block $tablify|0
+      (br_table $x $y $z $tablify|0
+       (local.tee $a
+        (i32.add
+         (i32.const 0)
+         (i32.const 1)
+        )
+       )
+      )
+     )
+     (unreachable)
+    )
+    (unreachable)
+   )
+   (unreachable)
+  )
+  (unreachable)
+ )
  (func $tiny-switch
   (if
    (i32.const 0)
diff --git a/test/passes/remove-unused-brs_enable-multivalue.wast b/test/passes/remove-unused-brs_enable-multivalue.wast
index afc8c3836..ae8878e00 100644
--- a/test/passes/remove-unused-brs_enable-multivalue.wast
+++ b/test/passes/remove-unused-brs_enable-multivalue.wast
@@ -1353,6 +1353,93 @@
     )
     (unreachable)
   )
+  (func $br-to-table-initial-tee (param $a i32)
+    (block $x
+      (block $y
+        (block $z
+          (br_if $x
+            (i32.eq
+              (local.tee $a
+                (i32.add
+                  (i32.const 10)
+                  (i32.const 1)
+                )
+              )
+              (i32.const 10)
+            )
+          )
+          (br_if $y (i32.eq (local.get $a) (i32.const 11)))
+          (br_if $z (i32.eq (local.get $a) (i32.const 12)))
+          (unreachable)
+        )
+        (unreachable)
+      )
+      (unreachable)
+    )
+    (unreachable)
+  )
+  (func $br-to-table-initial-tee-wrong-index (param $a i32)
+    (local $b i32)
+    (block $x
+      (block $y
+        (block $z
+          (br_if $x
+            (i32.eq
+              (local.tee $a (i32.const 99))
+              (i32.const 10)
+            )
+          )
+          ;; The subsequent conditions use a different local, $b, so we cannot
+          ;; optimize here.
+          (br_if $y (i32.eq (local.get $b) (i32.const 11)))
+          (br_if $z (i32.eq (local.get $b) (i32.const 12)))
+          (unreachable)
+        )
+        (unreachable)
+      )
+      (unreachable)
+    )
+    (unreachable)
+  )
+  (func $br-to-table-eqz (param $a i32)
+    (block $x
+      (block $y
+        (block $z
+          (br_if $x (i32.eqz (local.get $a)))
+          (br_if $y (i32.eq (local.get $a) (i32.const 1)))
+          (br_if $z (i32.eq (local.get $a) (i32.const 2)))
+          (unreachable)
+        )
+        (unreachable)
+      )
+      (unreachable)
+    )
+    (unreachable)
+  )
+  (func $br-to-table-tee-eqz (param $a i32)
+    (block $x
+      (block $y
+        (block $z
+          (br_if $x
+            (i32.eqz
+              (local.tee $a
+                (i32.add
+                  (i32.const 0)
+                  (i32.const 1)
+                )
+              )
+            )
+          )
+          (br_if $y (i32.eq (local.get $a) (i32.const 1)))
+          (br_if $z (i32.eq (local.get $a) (i32.const 2)))
+          (unreachable)
+        )
+        (unreachable)
+      )
+      (unreachable)
+    )
+    (unreachable)
+  )
   (func $tiny-switch
     (block $x
       (block $y