[Optimize Instructions] Simplify sign extentions (#4004)

requiring sign-extension:
```
i64(x) << 56 >> 56   ==>   i64.extend8_s(x)
i64(x) << 48 >> 48   ==>   i64.extend16_s(x)
i64(x) << 32 >> 32   ==>   i64.extend32_s(x)
i64.extend_i32_s(i32.wrap_i64(x))   ==>   i64.extend32_s(x)
```
general:
```
i32.wrap_i64(i64.extend_i32_s(x))   ==>   x
i32.wrap_i64(i64.extend_i32_u(x))   ==>   x
```

3 files changed, 115 insertions, 12 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 106bfea19..59b1130d1 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -406,6 +406,31 @@ struct OptimizeInstructions
         }
       }
       {
+        if (getModule()->features.hasSignExt()) {
+          Const *c1, *c2;
+          Expression* x;
+          // i64(x) << 56 >> 56   ==>   i64.extend8_s(x)
+          // i64(x) << 48 >> 48   ==>   i64.extend16_s(x)
+          // i64(x) << 32 >> 32   ==>   i64.extend32_s(x)
+          if (matches(curr,
+                      binary(ShrSInt64,
+                             binary(ShlInt64, any(&x), i64(&c1)),
+                             i64(&c2))) &&
+              Bits::getEffectiveShifts(c1) == Bits::getEffectiveShifts(c2)) {
+            switch (64 - Bits::getEffectiveShifts(c1)) {
+              case 8:
+                return replaceCurrent(builder.makeUnary(ExtendS8Int64, x));
+              case 16:
+                return replaceCurrent(builder.makeUnary(ExtendS16Int64, x));
+              case 32:
+                return replaceCurrent(builder.makeUnary(ExtendS32Int64, x));
+              default:
+                break;
+            }
+          }
+        }
+      }
+      {
         // unsigned(x) >= 0   =>   i32(1)
         Const* c;
         Expression* x;
@@ -792,6 +817,18 @@ struct OptimizeInstructions
         }
       }
       {
+        // i32.wrap_i64(i64.extend_i32_s(x))  =>  x
+        // i32.wrap_i64(i64.extend_i32_u(x))  =>  x
+        Unary* inner;
+        Expression* x;
+        if (matches(curr,
+                    unary(WrapInt64, unary(&inner, ExtendSInt32, any(&x)))) ||
+            matches(curr,
+                    unary(WrapInt64, unary(&inner, ExtendUInt32, any(&x))))) {
+          return replaceCurrent(x);
+        }
+      }
+      {
         // i32.eqz(i32.wrap_i64(x))  =>  i64.eqz(x)
         //   where maxBits(x) <= 32
         Unary* inner;
@@ -803,6 +840,20 @@ struct OptimizeInstructions
           return replaceCurrent(inner);
         }
       }
+      {
+        Unary* inner;
+        Expression* x;
+        if (getModule()->features.hasSignExt()) {
+          // i64.extend_i32_s(i32.wrap_i64(x))  =>  i64.extend32_s(x)
+          if (matches(curr,
+                      unary(ExtendSInt32, unary(&inner, WrapInt64, any(&x))))) {
+            inner->op = ExtendS32Int64;
+            inner->type = Type::i64;
+            inner->value = x;
+            return replaceCurrent(inner);
+          }
+        }
+      }
     }
 
     if (curr->op == EqZInt32) {
diff --git a/test/lit/passes/optimize-instructions-sign-ext.wast b/test/lit/passes/optimize-instructions-sign-ext.wast
index 142c2d2a4..0b3396735 100644
--- a/test/lit/passes/optimize-instructions-sign-ext.wast
+++ b/test/lit/passes/optimize-instructions-sign-ext.wast
@@ -2,7 +2,7 @@
 ;; RUN: wasm-opt %s --optimize-instructions --enable-sign-ext -S -o - | filecheck %s
 
 (module
-  ;; CHECK:      (func $duplicate-elimination (param $x i32) (param $y i32) (param $z i32) (param $w f64)
+  ;; CHECK:      (func $duplicate-elimination (param $x i32)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (i32.extend8_s
   ;; CHECK-NEXT:    (local.get $x)
@@ -14,8 +14,52 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
-  (func $duplicate-elimination (param $x i32) (param $y i32) (param $z i32) (param $w f64)
+  (func $duplicate-elimination (param $x i32)
     (drop (i32.extend8_s (i32.extend8_s (local.get $x))))
     (drop (i32.extend16_s (i32.extend16_s (local.get $x))))
   )
+
+  ;; i64(x) << 56 >> 56   ==>   i64.extend8_s(x)
+  ;; i64(x) << 48 >> 48   ==>   i64.extend16_s(x)
+  ;; i64(x) << 32 >> 32   ==>   i64.extend32_s(x)
+  ;; i64.extend_i32_s(i32.wrap_i64(x))   ==>   i64.extend32_s(x)
+
+  ;; CHECK:      (func $i64-sign-extentions (param $x i64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i64.extend8_s
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i64.extend16_s
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i64.extend32_s
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i64.shr_s
+  ;; CHECK-NEXT:    (i64.shl
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (i64.const 16)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i64.const 16)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i64.extend32_s
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $i64-sign-extentions (param $x i64)
+    (drop (i64.shr_s (i64.shl (local.get $x) (i64.const 56)) (i64.const 56)))
+    (drop (i64.shr_s (i64.shl (local.get $x) (i64.const 48)) (i64.const 48)))
+    (drop (i64.shr_s (i64.shl (local.get $x) (i64.const 32)) (i64.const 32)))
+    (drop (i64.shr_s (i64.shl (local.get $x) (i64.const 16)) (i64.const 16))) ;; skip
+    (drop (i64.extend_i32_s (i32.wrap_i64 (local.get $x))))
+  )
 )
diff --git a/test/lit/passes/optimize-instructions.wast b/test/lit/passes/optimize-instructions.wast
index e8296e8cc..5fa703bd2 100644
--- a/test/lit/passes/optimize-instructions.wast
+++ b/test/lit/passes/optimize-instructions.wast
@@ -2707,11 +2707,7 @@
   ;; CHECK-NEXT:  (i32.shr_s
   ;; CHECK-NEXT:   (i32.shl
   ;; CHECK-NEXT:    (i32.shr_u
-  ;; CHECK-NEXT:     (i32.wrap_i64
-  ;; CHECK-NEXT:      (i64.extend_i32_s
-  ;; CHECK-NEXT:       (i32.const -1)
-  ;; CHECK-NEXT:      )
-  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (i32.const -1)
   ;; CHECK-NEXT:     (i32.const 24)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (i32.const 24)
@@ -2737,11 +2733,7 @@
   )
   ;; CHECK:      (func $sext-24-shr_u-wrap-extend (result i32)
   ;; CHECK-NEXT:  (i32.shr_u
-  ;; CHECK-NEXT:   (i32.wrap_i64
-  ;; CHECK-NEXT:    (i64.extend_i32_s
-  ;; CHECK-NEXT:     (i32.const -1)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.const -1)
   ;; CHECK-NEXT:   (i32.const 25)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
@@ -11256,6 +11248,22 @@
       )
     )
   )
+
+  ;; i32.wrap_i64(i64.extend_i32_s(x))  ==>  x
+  ;; i32.wrap_i64(i64.extend_i32_u(x))  ==>  x
+
+  ;; CHECK:      (func $sign-and-zero-extention-elimination (param $x i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $sign-and-zero-extention-elimination (param $x i32)
+    (drop (i32.wrap_i64 (i64.extend_i32_s (local.get $x))))
+    (drop (i32.wrap_i64 (i64.extend_i32_u (local.get $x))))
+  )
   ;; CHECK:      (func $optimize-shifts (param $x i32) (param $y i32) (param $z i64) (param $w i64)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (local.get $x)