[Optimize Instructions] Refactor squared rules (#4840)

+ Move these rules to separate function;
+ Refactor them to use matches;
+ Add comments;
+ Handle rotational shifts as well;
+ Handle overflows for `<<`, `>>`, `>>>` shifts;
+ Add mixed rotate rules:
```rust
   rotl(rotr(x, C1), C2)   =>   rotr(x, C1 - C2)
   rotr(rotl(x, C1), C2)   =>   rotl(x, C1 - C2)
```

1 files changed, 127 insertions, 51 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 63f47abd1..55176a6e9 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -682,57 +682,8 @@ struct OptimizeInstructions
       if (auto* ret = optimizeWithConstantOnRight(curr)) {
         return replaceCurrent(ret);
       }
-      // the square of some operations can be merged
-      if (auto* left = curr->left->dynCast<Binary>()) {
-        if (left->op == curr->op) {
-          if (auto* leftRight = left->right->dynCast<Const>()) {
-            if (left->op == AndInt32 || left->op == AndInt64) {
-              leftRight->value = leftRight->value.and_(right->value);
-              return replaceCurrent(left);
-            } else if (left->op == OrInt32 || left->op == OrInt64) {
-              leftRight->value = leftRight->value.or_(right->value);
-              return replaceCurrent(left);
-            } else if (left->op == XorInt32 || left->op == XorInt64) {
-              leftRight->value = leftRight->value.xor_(right->value);
-              return replaceCurrent(left);
-            } else if (left->op == MulInt32 || left->op == MulInt64) {
-              leftRight->value = leftRight->value.mul(right->value);
-              return replaceCurrent(left);
-
-              // TODO:
-              // handle signed / unsigned divisions. They are more complex
-            } else if (left->op == ShlInt32 || left->op == ShrUInt32 ||
-                       left->op == ShrSInt32 || left->op == ShlInt64 ||
-                       left->op == ShrUInt64 || left->op == ShrSInt64) {
-              // shifts only use an effective amount from the constant, so
-              // adding must be done carefully
-              auto total = Bits::getEffectiveShifts(leftRight) +
-                           Bits::getEffectiveShifts(right);
-              if (total == Bits::getEffectiveShifts(total, right->type)) {
-                // no overflow, we can do this
-                leftRight->value = Literal::makeFromInt32(total, right->type);
-                return replaceCurrent(left);
-              } // TODO: handle overflows
-            }
-          }
-        }
-        if (left->op == Abstract::getBinary(left->type, Abstract::Shl) &&
-            curr->op == Abstract::getBinary(curr->type, Abstract::Mul)) {
-          if (auto* leftRight = left->right->dynCast<Const>()) {
-            left->op = Abstract::getBinary(left->type, Abstract::Mul);
-            // (x << C1) * C2   ->   x * (C2 << C1)
-            leftRight->value = right->value.shl(leftRight->value);
-            return replaceCurrent(left);
-          }
-        }
-        if (left->op == Abstract::getBinary(left->type, Abstract::Mul) &&
-            curr->op == Abstract::getBinary(curr->type, Abstract::Shl)) {
-          if (auto* leftRight = left->right->dynCast<Const>()) {
-            // (x * C1) << C2   ->   x * (C1 << C2)
-            leftRight->value = leftRight->value.shl(right->value);
-            return replaceCurrent(left);
-          }
-        }
+      if (auto* ret = optimizeDoubletonWithConstantOnRight(curr)) {
+        return replaceCurrent(ret);
       }
       if (right->type == Type::i32) {
         BinaryOp op;
@@ -3443,6 +3394,131 @@ private:
     return nullptr;
   }
 
+  // Folding two expressions into one with similar operations and
+  // constants on RHSs
+  Expression* optimizeDoubletonWithConstantOnRight(Binary* curr) {
+    using namespace Match;
+    using namespace Abstract;
+    {
+      Binary* inner;
+      Const *c1, *c2 = curr->right->cast<Const>();
+      if (matches(curr->left, binary(&inner, any(), ival(&c1))) &&
+          inner->op == curr->op) {
+        Type type = inner->type;
+        BinaryOp op = inner->op;
+        // (x & C1) & C2   =>   x & (C1 & C2)
+        if (op == getBinary(type, And)) {
+          c1->value = c1->value.and_(c2->value);
+          return inner;
+        }
+        // (x | C1) | C2   =>   x | (C1 | C2)
+        if (op == getBinary(type, Or)) {
+          c1->value = c1->value.or_(c2->value);
+          return inner;
+        }
+        // (x ^ C1) ^ C2   =>   x ^ (C1 ^ C2)
+        if (op == getBinary(type, Xor)) {
+          c1->value = c1->value.xor_(c2->value);
+          return inner;
+        }
+        // (x * C1) * C2   =>   x * (C1 * C2)
+        if (op == getBinary(type, Mul)) {
+          c1->value = c1->value.mul(c2->value);
+          return inner;
+        }
+        // TODO:
+        // handle signed / unsigned divisions. They are more complex
+
+        // (x <<>> C1) <<>> C2   =>   x <<>> (C1 + C2)
+        if (hasAnyShift(op)) {
+          // shifts only use an effective amount from the constant, so
+          // adding must be done carefully
+          auto total =
+            Bits::getEffectiveShifts(c1) + Bits::getEffectiveShifts(c2);
+          auto effectiveTotal = Bits::getEffectiveShifts(total, c1->type);
+          if (total == effectiveTotal) {
+            // no overflow, we can do this
+            c1->value = Literal::makeFromInt32(total, c1->type);
+            return inner;
+          } else {
+            // overflow. Handle different scenarious
+            if (hasAnyRotateShift(op)) {
+              // overflow always accepted in rotation shifts
+              c1->value = Literal::makeFromInt32(effectiveTotal, c1->type);
+              return inner;
+            }
+            // handle overflows for general shifts
+            //   x << C1 << C2    =>   0 or { drop(x), 0 }
+            //   x >>> C1 >>> C2  =>   0 or { drop(x), 0 }
+            // iff `C1 + C2` -> overflows
+            if ((op == getBinary(type, Shl) || op == getBinary(type, ShrU))) {
+              auto* x = inner->left;
+              c1->value = Literal::makeZero(c1->type);
+              if (!effects(x).hasSideEffects()) {
+                //  =>  0
+                return c1;
+              } else {
+                //  =>  { drop(x), 0 }
+                Builder builder(*getModule());
+                return builder.makeBlock({builder.makeDrop(x), c1});
+              }
+            }
+            //   i32(x) >> C1 >> C2   =>   x >> 31
+            //   i64(x) >> C1 >> C2   =>   x >> 63
+            // iff `C1 + C2` -> overflows
+            if (op == getBinary(type, ShrS)) {
+              c1->value = Literal::makeFromInt32(c1->type.getByteSize() * 8 - 1,
+                                                 c1->type);
+              return inner;
+            }
+          }
+        }
+      }
+    }
+    {
+      // (x << C1) * C2   =>   x * (C2 << C1)
+      Binary* inner;
+      Const *c1, *c2;
+      if (matches(
+            curr,
+            binary(Mul, binary(&inner, Shl, any(), ival(&c1)), ival(&c2)))) {
+        inner->op = getBinary(inner->type, Mul);
+        c1->value = c2->value.shl(c1->value);
+        return inner;
+      }
+    }
+    {
+      // (x * C1) << C2   =>   x * (C1 << C2)
+      Binary* inner;
+      Const *c1, *c2;
+      if (matches(
+            curr,
+            binary(Shl, binary(&inner, Mul, any(), ival(&c1)), ival(&c2)))) {
+        c1->value = c1->value.shl(c2->value);
+        return inner;
+      }
+    }
+    {
+      // TODO: Add canonicalization rotr to rotl and remove these rules.
+      // rotl(rotr(x, C1), C2)   =>   rotr(x, C1 - C2)
+      // rotr(rotl(x, C1), C2)   =>   rotl(x, C1 - C2)
+      Binary* inner;
+      Const *c1, *c2;
+      if (matches(
+            curr,
+            binary(RotL, binary(&inner, RotR, any(), ival(&c1)), ival(&c2))) ||
+          matches(
+            curr,
+            binary(RotR, binary(&inner, RotL, any(), ival(&c1)), ival(&c2)))) {
+        auto diff = Bits::getEffectiveShifts(c1) - Bits::getEffectiveShifts(c2);
+        c1->value = Literal::makeFromInt32(
+          Bits::getEffectiveShifts(diff, c2->type), c2->type);
+        return inner;
+      }
+    }
+    return nullptr;
+  }
+
   // optimize trivial math operations, given that the left side of a binary
   // is a constant. since we canonicalize constants to the right for symmetrical
   // operations, we only need to handle asymmetrical ones here