OptimizeInstructions: Optimize boolean selects (#4147)

If all a select's inputs are boolean, we can sometimes turn the select
into an AND or an OR operation,

x ? y : 0  =>  x & y
x ? 1 : y  =>  x | y

I believe LLVM aggressively canonicalizes to this form. It makes sense
to do here too as it is smaller (save the constant 0 or 1). It also allows
further optimizations (which is why LLVM does it) but I don't think we
have those yet.

4 files changed, 271 insertions, 4 deletions

diff --git a/src/ir/bits.h b/src/ir/bits.h
index 950864a23..dc02bd8f8 100644
--- a/src/ir/bits.h
+++ b/src/ir/bits.h
@@ -413,6 +413,8 @@ Index getMaxBits(Expression* curr,
     // a tee passes through the value
     return getMaxBits(set->value, localInfoProvider);
   } else if (auto* get = curr->dynCast<LocalGet>()) {
+    // TODO: Should this be optional?
+    assert(localInfoProvider);
     return localInfoProvider->getMaxBitsForLocal(get);
   } else if (auto* load = curr->dynCast<Load>()) {
     // if signed, then the sign-extension might fill all the bits
diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 357f172d9..768ea5702 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -1819,6 +1819,25 @@ private:
         }
       }
     }
+    if (curr->type == Type::i32 &&
+        Bits::getMaxBits(curr->condition, this) <= 1 &&
+        Bits::getMaxBits(curr->ifTrue, this) <= 1 &&
+        Bits::getMaxBits(curr->ifFalse, this) <= 1) {
+      // The condition and both arms are i32 booleans, which allows us to do
+      // boolean optimizations.
+      Expression* x;
+      Expression* y;
+
+      // x ? y : 0   ==>   x & y
+      if (matches(curr, select(any(&y), ival(0), any(&x)))) {
+        return builder.makeBinary(AndInt32, y, x);
+      }
+
+      // x ? 1 : y   ==>   x | y
+      if (matches(curr, select(ival(1), any(&y), any(&x)))) {
+        return builder.makeBinary(OrInt32, y, x);
+      }
+    }
     {
       // Sides are identical, fold
       Expression *ifTrue, *ifFalse, *c;
diff --git a/test/lit/passes/optimize-instructions.wast b/test/lit/passes/optimize-instructions.wast
index 7b0849352..22edf58e2 100644
--- a/test/lit/passes/optimize-instructions.wast
+++ b/test/lit/passes/optimize-instructions.wast
@@ -703,6 +703,252 @@
       )
     )
   )
+  ;; CHECK:      (func $select-or (param $x i32) (param $y i32) (result i32)
+  ;; CHECK-NEXT:  (i32.or
+  ;; CHECK-NEXT:   (i32.and
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.eq
+  ;; CHECK-NEXT:    (local.get $y)
+  ;; CHECK-NEXT:    (i32.const 1337)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $select-or (param $x i32) (param $y i32) (result i32)
+    (select
+      (i32.const 1)
+      (i32.eq
+        (local.get $y)
+        (i32.const 1337)
+      )
+      (i32.and
+        (local.get $x)
+        (i32.const 1)
+      )
+    )
+  )
+  ;; CHECK:      (func $select-or-side-effects (param $x i32) (param $y i32) (result i32)
+  ;; CHECK-NEXT:  (i32.or
+  ;; CHECK-NEXT:   (i32.eq
+  ;; CHECK-NEXT:    (call $select-or-side-effects
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (local.get $y)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1337)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.and
+  ;; CHECK-NEXT:    (call $select-or-side-effects
+  ;; CHECK-NEXT:     (local.get $y)
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $select-or-side-effects (param $x i32) (param $y i32) (result i32)
+    ;; When there are side effects, the order of the operations must remain
+    ;; correct.
+    (select
+      (i32.const 1)
+      (i32.eq
+        (call $select-or-side-effects
+          (local.get $x)
+          (local.get $y)
+        )
+        (i32.const 1337)
+      )
+      (i32.and
+        (call $select-or-side-effects
+          (local.get $y)
+          (local.get $x)
+        )
+        (i32.const 1)
+      )
+    )
+  )
+  ;; CHECK:      (func $select-or-no-bits (param $x i32) (param $y i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (select
+  ;; CHECK-NEXT:    (i32.const 2)
+  ;; CHECK-NEXT:    (i32.eq
+  ;; CHECK-NEXT:     (local.get $y)
+  ;; CHECK-NEXT:     (i32.const 1337)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.eq
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (i32.const 42)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (select
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:    (local.get $y)
+  ;; CHECK-NEXT:    (i32.eq
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (i32.const 42)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (select
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:    (i32.eq
+  ;; CHECK-NEXT:     (local.get $y)
+  ;; CHECK-NEXT:     (i32.const 1337)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $select-or-no-bits (param $x i32) (param $y i32)
+    ;; The following cannot be optimized into an "or" operation due to maxBits
+    ;; not being known to be 1.
+    (drop
+      (select
+        ;; Too many bits in ifTrue
+        (i32.const 2)
+        (i32.eq
+          (local.get $y)
+          (i32.const 1337)
+        )
+        (i32.eq
+          (local.get $x)
+          (i32.const 42)
+        )
+      )
+    )
+    (drop
+      (select
+        (i32.const 1)
+        ;; Too many bits in ifFalse
+        (local.get $y)
+        (i32.eq
+          (local.get $x)
+          (i32.const 42)
+        )
+      )
+    )
+    (drop
+      (select
+        (i32.const 1)
+        (i32.eq
+          (local.get $y)
+          (i32.const 1337)
+        )
+        ;; Too many bits in condition
+        (local.get $x)
+      )
+    )
+  )
+  ;; CHECK:      (func $select-or-no-type (param $x i32) (param $y i64) (result i64)
+  ;; CHECK-NEXT:  (select
+  ;; CHECK-NEXT:   (i64.const 1)
+  ;; CHECK-NEXT:   (i64.and
+  ;; CHECK-NEXT:    (local.get $y)
+  ;; CHECK-NEXT:    (i64.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.and
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $select-or-no-type (param $x i32) (param $y i64) (result i64)
+    ;; An i64 result cannot be optimized into an "or" of the ifTrue and the
+    ;; condition due to their types being different.
+    (select
+      (i64.const 1)
+      (i64.and
+        (local.get $y)
+        (i64.const 1)
+      )
+      (i32.and
+        (local.get $x)
+        (i32.const 1)
+      )
+    )
+  )
+  ;; CHECK:      (func $select-or-no-const (param $x i32) (param $y i32) (result i32)
+  ;; CHECK-NEXT:  (select
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:   (i32.eq
+  ;; CHECK-NEXT:    (local.get $y)
+  ;; CHECK-NEXT:    (i32.const 1337)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.and
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $select-or-no-const (param $x i32) (param $y i32) (result i32)
+    (select
+      ;; The wrong const (should be 1).
+      (i32.const 0)
+      (i32.eq
+        (local.get $y)
+        (i32.const 1337)
+      )
+      (i32.and
+        (local.get $x)
+        (i32.const 1)
+      )
+    )
+  )
+  ;; CHECK:      (func $select-and (param $x i32) (param $y i32) (result i32)
+  ;; CHECK-NEXT:  (i32.and
+  ;; CHECK-NEXT:   (i32.eq
+  ;; CHECK-NEXT:    (local.get $y)
+  ;; CHECK-NEXT:    (i32.const 1337)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.eq
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $select-and (param $x i32) (param $y i32) (result i32)
+    (select
+      (i32.eq
+        (local.get $y)
+        (i32.const 1337)
+      )
+      (i32.const 0)
+      (i32.eq
+        (local.get $x)
+        (i32.const 42)
+      )
+    )
+  )
+  ;; CHECK:      (func $select-and-no-const (param $x i32) (param $y i32) (result i32)
+  ;; CHECK-NEXT:  (select
+  ;; CHECK-NEXT:   (i32.eq
+  ;; CHECK-NEXT:    (local.get $y)
+  ;; CHECK-NEXT:    (i32.const 1337)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.const 1)
+  ;; CHECK-NEXT:   (i32.eq
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $select-and-no-const (param $x i32) (param $y i32) (result i32)
+    (select
+      (i32.eq
+        (local.get $y)
+        (i32.const 1337)
+      )
+      ;; The wrong constant (should be 0).
+      (i32.const 1)
+      (i32.eq
+        (local.get $x)
+        (i32.const 42)
+      )
+    )
+  )
   ;; CHECK:      (func $load8_s-and-255 (result i32)
   ;; CHECK-NEXT:  (i32.load8_u
   ;; CHECK-NEXT:   (i32.const 0)
diff --git a/test/wasm2js/i64-lowering.2asm.js.opt b/test/wasm2js/i64-lowering.2asm.js.opt
index 395b6f442..6d90214d2 100644
--- a/test/wasm2js/i64-lowering.2asm.js.opt
+++ b/test/wasm2js/i64-lowering.2asm.js.opt
@@ -22,19 +22,19 @@ function asmFunc(env) {
  }
  
  function legalstub$3($0, $1, $2, $3) {
-  return ($1 | 0) > ($3 | 0) ? 1 : ($1 | 0) >= ($3 | 0) ? $0 >>> 0 >= $2 >>> 0 : 0;
+  return ($1 | 0) >= ($3 | 0) & $0 >>> 0 >= $2 >>> 0 | ($1 | 0) > ($3 | 0);
  }
  
  function legalstub$4($0, $1, $2, $3) {
-  return ($1 | 0) > ($3 | 0) ? 1 : ($1 | 0) >= ($3 | 0) ? $0 >>> 0 > $2 >>> 0 : 0;
+  return $0 >>> 0 > $2 >>> 0 & ($1 | 0) >= ($3 | 0) | ($1 | 0) > ($3 | 0);
  }
  
  function legalstub$5($0, $1, $2, $3) {
-  return ($1 | 0) < ($3 | 0) ? 1 : ($1 | 0) <= ($3 | 0) ? $0 >>> 0 <= $2 >>> 0 : 0;
+  return ($1 | 0) <= ($3 | 0) & $0 >>> 0 <= $2 >>> 0 | ($1 | 0) < ($3 | 0);
  }
  
  function legalstub$6($0, $1, $2, $3) {
-  return ($1 | 0) < ($3 | 0) ? 1 : ($1 | 0) <= ($3 | 0) ? $0 >>> 0 < $2 >>> 0 : 0;
+  return $0 >>> 0 < $2 >>> 0 & ($1 | 0) <= ($3 | 0) | ($1 | 0) < ($3 | 0);
  }
  
  function legalstub$7($0, $1, $2, $3) {