Optimize if/select with one arm an EqZ and another a 0 or a 1 (#3822)

(select
  (i32.eqz (X))
  (i32.const 0|1)
  (Y)
)

=>

(i32.eqz
  (select
    (X)
    (i32.const 1|0)
    (Y)
  )
)

This is beneficial as the eqz may be folded into something on the outside. I
see this pattern in real-world code, both a GC benchmark (which is why I
noticed it) and it shrinks code size by tiny amounts on the emscripten
benchmark suite as well.

1 files changed, 61 insertions, 0 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 609caf588..b6ebb961a 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -864,6 +864,7 @@ struct OptimizeInstructions
     if (auto* ret = optimizeSelect(curr)) {
       return replaceCurrent(ret);
     }
+    optimizeTernary(curr, curr->condition, curr->ifTrue, curr->ifFalse);
   }
 
   void visitGlobalSet(GlobalSet* curr) {
@@ -914,6 +915,7 @@ struct OptimizeInstructions
           return replaceCurrent(ret);
         }
       }
+      optimizeTernary(curr, curr->condition, curr->ifTrue, curr->ifFalse);
     }
   }
 
@@ -2779,6 +2781,65 @@ private:
         return false;
     }
   }
+
+  // Optimize an if-else or a select, something with a condition and two
+  // arms with outputs.
+  void optimizeTernary(Expression* curr,
+                       Expression* condition,
+                       Expression*& ifTrue,
+                       Expression*& ifFalse) {
+    if (curr->type == Type::unreachable) {
+      return;
+    }
+
+    using namespace Match;
+    Builder builder(*getModule());
+
+    // If one arm is an operation and the other is an appropriate constant, we
+    // can move the operation outside (where it may be further optimized), e.g.
+    //
+    //  (select
+    //    (i32.eqz (X))
+    //    (i32.const 0|1)
+    //    (Y)
+    //  )
+    // =>
+    //  (i32.eqz
+    //    (select
+    //      (X)
+    //      (i32.const 1|0)
+    //      (Y)
+    //    )
+    //  )
+    {
+      Unary* un;
+      Expression* x;
+      Const* c;
+      auto check = [&](Expression* a, Expression* b) {
+        if (matches(a, unary(&un, EqZInt32, any(&x))) && matches(b, ival(&c))) {
+          auto value = c->value.geti32();
+          return value == 0 || value == 1;
+        }
+        return false;
+      };
+      if (check(ifTrue, ifFalse) || check(ifFalse, ifTrue)) {
+        auto updateArm = [&](Expression* arm) -> Expression* {
+          if (arm == un) {
+            // This is the arm that had the eqz, which we need to remove.
+            return un->value;
+          } else {
+            // This is the arm with the constant, which we need to flip.
+            c->value = Literal(int32_t(1 - c->value.geti32()));
+            return c;
+          }
+        };
+        ifTrue = updateArm(ifTrue);
+        ifFalse = updateArm(ifFalse);
+        un->value = curr;
+        return replaceCurrent(un);
+      }
+    }
+  }
 };
 
 Pass* createOptimizeInstructionsPass() { return new OptimizeInstructions; }