[OptimizeInstructions] Fix fuzz bug with shifts (#3376)

The code there looks for a "sign-extend": (x << a) >> b where the
right shift is signed. If a = b = 24 for example then that is a sign
extend of an 8-bit value (it works by shifting the 8-bit value's sign bit
to the position of the 32-bit value's sign bit, then shifting all the way
back, which fills everything above 8 bits with the sign bit). The tricky
thing is that in some cases we can handle a != b - but we forgot a
place to check that. Specifically, a repeated sign-extend is not
necessary, but if the outer one has extra shifts, we can't do it.

This is annoyingly complex code, but for purposes of reviewing this
PR, you can see (unless I messed up) that the only change is to
ensure that when we look for a repeated sign extend, then we
only optimize that case when there are no extra shifts. And a
repeated sign-extend is obviously ok to remove,

(((x << a) >> a) << a) >> a   =>   (x << a) >> a

This is an ancient bug, showing how hard it can be to find certain
patterns either by fuzzing or in the real world...

Fixes #3362

3 files changed, 71 insertions, 13 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 2ca670f89..2b0a95fc7 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -468,9 +468,9 @@ struct OptimizeInstructions
 
     if (auto* binary = curr->dynCast<Binary>()) {
       if (auto* ext = Properties::getAlmostSignExt(binary)) {
-        Index extraShifts;
-        auto bits = Properties::getAlmostSignExtBits(binary, extraShifts);
-        if (extraShifts == 0) {
+        Index extraLeftShifts;
+        auto bits = Properties::getAlmostSignExtBits(binary, extraLeftShifts);
+        if (extraLeftShifts == 0) {
           if (auto* load =
                 Properties::getFallthrough(ext, getPassOptions(), features)
                   ->dynCast<Load>()) {
@@ -488,11 +488,17 @@ struct OptimizeInstructions
             }
           }
         }
-        // if the sign-extend input cannot have a sign bit, we don't need it
-        // we also don't need it if it already has an identical-sized sign
-        // extend
-        if (Bits::getMaxBits(ext, this) + extraShifts < bits ||
-            isSignExted(ext, bits)) {
+        // We can in some cases remove part of a sign extend, that is,
+        //   (x << A) >> B   =>   x << (A - B)
+        // If the sign-extend input cannot have a sign bit, we don't need it.
+        if (Bits::getMaxBits(ext, this) + extraLeftShifts < bits) {
+          return removeAlmostSignExt(binary);
+        }
+        // We also don't need it if it already has an identical-sized sign
+        // extend applied to it. That is, if it is already a sign-extended
+        // value, then another sign extend will do nothing. We do need to be
+        // careful of the extra shifts, though.
+        if (isSignExted(ext, bits) && extraLeftShifts == 0) {
           return removeAlmostSignExt(binary);
         }
       } else if (binary->op == EqInt32 || binary->op == NeInt32) {
diff --git a/test/passes/optimize-instructions_fuzz-exec.txt b/test/passes/optimize-instructions_fuzz-exec.txt
index e8985222d..27d8795ed 100644
--- a/test/passes/optimize-instructions_fuzz-exec.txt
+++ b/test/passes/optimize-instructions_fuzz-exec.txt
@@ -257,19 +257,22 @@
 [LoggingExternalInterface logging 1]
 [LoggingExternalInterface logging 1]
 [LoggingExternalInterface logging 0]
+[fuzz-exec] calling do-shift
+[LoggingExternalInterface logging -64]
 [fuzz-exec] calling call-compare-maybe-signed-eq
 [fuzz-exec] note result: call-compare-maybe-signed-eq => 0
 [fuzz-exec] calling call-compare-maybe-signed-ne
 [fuzz-exec] note result: call-compare-maybe-signed-ne => 1
 (module
  (type $i32_=>_none (func (param i32)))
+ (type $none_=>_none (func))
  (type $none_=>_i32 (func (result i32)))
  (type $i32_=>_i32 (func (param i32) (result i32)))
- (type $none_=>_none (func))
  (import "fuzzing-support" "log-i32" (func $log (param i32)))
  (export "foo" (func $1))
- (export "call-compare-maybe-signed-eq" (func $3))
- (export "call-compare-maybe-signed-ne" (func $5))
+ (export "do-shift" (func $3))
+ (export "call-compare-maybe-signed-eq" (func $5))
+ (export "call-compare-maybe-signed-ne" (func $7))
  (func $signed-comparison-to-unsigned
   (call $log
    (block (result i32)
@@ -321,13 +324,35 @@
    )
   )
  )
+ (func $shift (param $0 i32)
+  (call $log
+   (i32.shr_s
+    (i32.shl
+     (i32.shr_s
+      (i32.shl
+       (local.get $0)
+       (i32.const 24)
+      )
+      (i32.const 24)
+     )
+     (i32.const 30)
+    )
+    (i32.const 24)
+   )
+  )
+ )
+ (func $3
+  (call $shift
+   (i32.const 65419)
+  )
+ )
  (func $compare-maybe-signed-eq (param $0 i32) (result i32)
   (drop
    (local.get $0)
   )
   (i32.const 0)
  )
- (func $3 (result i32)
+ (func $5 (result i32)
   (call $compare-maybe-signed-eq
    (i32.const 128)
   )
@@ -338,7 +363,7 @@
   )
   (i32.const 1)
  )
- (func $5 (result i32)
+ (func $7 (result i32)
   (call $compare-maybe-signed-ne
    (i32.const 128)
   )
@@ -348,6 +373,8 @@
 [LoggingExternalInterface logging 1]
 [LoggingExternalInterface logging 1]
 [LoggingExternalInterface logging 0]
+[fuzz-exec] calling do-shift
+[LoggingExternalInterface logging -64]
 [fuzz-exec] calling call-compare-maybe-signed-eq
 [fuzz-exec] note result: call-compare-maybe-signed-eq => 0
 [fuzz-exec] calling call-compare-maybe-signed-ne
diff --git a/test/passes/optimize-instructions_fuzz-exec.wast b/test/passes/optimize-instructions_fuzz-exec.wast
index 8bd43b694..87b6eb594 100644
--- a/test/passes/optimize-instructions_fuzz-exec.wast
+++ b/test/passes/optimize-instructions_fuzz-exec.wast
@@ -287,6 +287,31 @@
    )
   )
  )
+ (func $shift (param $0 i32)
+  (call $log
+   ;; x << 24 >> 24 << 30 >> 24 - the extra shifts make it invalid to do the
+   ;; optimization of not repeating a sign-extend. That is, this would be valid
+   ;; if the 30 were replaced by a 24.
+   (i32.shr_s
+    (i32.shl
+     (i32.shr_s
+      (i32.shl
+       (local.get $0)
+       (i32.const 24)
+      )
+      (i32.const 24)
+     )
+     (i32.const 30)
+    )
+    (i32.const 24)
+   )
+  )
+ )
+ (func "do-shift"
+  (call $shift
+   (i32.const 65419)
+  )
+ )
  ;; similar, but with the value compared to having the sign bit set but no
  ;; upper bits
  (func $compare-maybe-signed-eq (param $0 i32) (result i32)