constant_propagation_pass: Propagate constants across commutative operations

e.g. (a & b) & c == a & (b & c) where b and c are constants

src/frontend/ir/value.cpp

@@ -213,23 +213,7 @@ bool Value::IsUnsignedImmediate(u64 value) const {
 }
 
 bool Value::HasAllBitsSet() const {
-    ASSERT(IsImmediate());
+    return IsSignedImmediate(-1);
-
-    switch (GetType()) {
-    case IR::Type::U1:
-        return GetU1();
-    case IR::Type::U8:
-        return GetU8() == 0xFF;
-    case IR::Type::U16:
-        return GetU16() == 0xFFFF;
-    case IR::Type::U32:
-        return GetU32() == 0xFFFFFFFF;
-    case IR::Type::U64:
-        return GetU64() == 0xFFFFFFFFFFFFFFFF;
-    default:
-        ASSERT_MSG(false, "HasAllBitsSet called on an incompatible Value type.");
-        return false;
-    }
 }
 
 bool Value::IsZero() const {

src/ir_opt/constant_propagation_pass.cpp

@@ -23,6 +23,49 @@ void ReplaceUsesWith(IR::Inst& inst, bool is_32_bit, u64 value) {
     }
 }
 
+IR::Value Value(bool is_32_bit, u64 value) {
+    return is_32_bit ? IR::Value{static_cast<u32>(value)} : IR::Value{value};
+}
+
+template <typename ImmFn>
+bool FoldCommutative(IR::Inst& inst, bool is_32_bit, ImmFn imm_fn) {
+    const auto lhs = inst.GetArg(0);
+    const auto rhs = inst.GetArg(1);
+
+    const bool is_lhs_immediate = lhs.IsImmediate();
+    const bool is_rhs_immediate = rhs.IsImmediate();
+
+    if (is_lhs_immediate && is_rhs_immediate) {
+        const u64 result = imm_fn(lhs.GetImmediateAsU64(), rhs.GetImmediateAsU64());
+        ReplaceUsesWith(inst, is_32_bit, result);
+        return false;
+    }
+
+    if (is_lhs_immediate && !is_rhs_immediate) {
+        const IR::Inst* rhs_inst = rhs.GetInstRecursive();
+        if (rhs_inst->GetOpcode() == inst.GetOpcode() && rhs_inst->GetArg(1).IsImmediate()) {
+            const u64 combined = imm_fn(lhs.GetImmediateAsU64(), rhs_inst->GetArg(1).GetImmediateAsU64());
+            inst.SetArg(0, rhs_inst->GetArg(0));
+            inst.SetArg(1, Value(is_32_bit, combined));
+        } else {
+            // Normalize
+            inst.SetArg(0, rhs);
+            inst.SetArg(1, lhs);
+        }
+    }
+
+    if (!is_lhs_immediate && is_rhs_immediate) {
+        const IR::Inst* lhs_inst = lhs.GetInstRecursive();
+        if (lhs_inst->GetOpcode() == inst.GetOpcode() && lhs_inst->GetArg(1).IsImmediate()) {
+            const u64 combined = imm_fn(rhs.GetImmediateAsU64(), lhs_inst->GetArg(1).GetImmediateAsU64());
+            inst.SetArg(0, lhs_inst->GetArg(0));
+            inst.SetArg(1, Value(is_32_bit, combined));
+        }
+    }
+
+    return true;
+}
+
 // Folds AND operations based on the following:
 //
 // 1. imm_x & imm_y -> result
@@ -32,21 +75,13 @@ void ReplaceUsesWith(IR::Inst& inst, bool is_32_bit, u64 value) {
 // 5. x & y -> x (where y has all bits set to 1)
 //
 void FoldAND(IR::Inst& inst, bool is_32_bit) {
-    const auto lhs = inst.GetArg(0);
+    if (FoldCommutative(inst, is_32_bit, [](u64 a, u64 b) { return a & b; })) {
         const auto rhs = inst.GetArg(1);
-
+        if (rhs.IsZero()) {
-    const bool is_lhs_immediate = lhs.IsImmediate();
-    const bool is_rhs_immediate = rhs.IsImmediate();
-
-    if (is_lhs_immediate && is_rhs_immediate) {
-        const u64 result = lhs.GetImmediateAsU64() & rhs.GetImmediateAsU64();
-        ReplaceUsesWith(inst, is_32_bit, result);
-    } else if (lhs.IsZero() || rhs.IsZero()) {
             ReplaceUsesWith(inst, is_32_bit, 0);
-    } else if (is_lhs_immediate && lhs.HasAllBitsSet()) {
+        } else if (rhs.HasAllBitsSet()) {
-        inst.ReplaceUsesWith(rhs);
+            inst.ReplaceUsesWith(inst.GetArg(0));
-    } else if (is_rhs_immediate && rhs.HasAllBitsSet()) {
+        }
-        inst.ReplaceUsesWith(lhs);
     }
 }
 
@@ -108,16 +143,11 @@ void FoldDivide(IR::Inst& inst, bool is_32_bit, bool is_signed) {
 // 3. 0 ^ y -> y
 //
 void FoldEOR(IR::Inst& inst, bool is_32_bit) {
-    const auto lhs = inst.GetArg(0);
+    if (FoldCommutative(inst, is_32_bit, [](u64 a, u64 b) { return a ^ b; })) {
         const auto rhs = inst.GetArg(1);
-
+        if (rhs.IsZero()) {
-    if (lhs.IsImmediate() && rhs.IsImmediate()) {
+            inst.ReplaceUsesWith(inst.GetArg(0));
-        const u64 result = lhs.GetImmediateAsU64() ^ rhs.GetImmediateAsU64();
+        }
-        ReplaceUsesWith(inst, is_32_bit, result);
-    } else if (lhs.IsZero()) {
-        inst.ReplaceUsesWith(rhs);
-    } else if (rhs.IsZero()) {
-        inst.ReplaceUsesWith(lhs);
     }
 }
 
@@ -180,18 +210,13 @@ void FoldMostSignificantWord(IR::Inst& inst) {
 // 5. 1 * y -> y
 //
 void FoldMultiply(IR::Inst& inst, bool is_32_bit) {
-    const auto lhs = inst.GetArg(0);
+    if (FoldCommutative(inst, is_32_bit, [](u64 a, u64 b) { return a * b; })) {
         const auto rhs = inst.GetArg(1);
-
+        if (rhs.IsZero()) {
-    if (lhs.IsImmediate() && rhs.IsImmediate()) {
-        const u64 result = lhs.GetImmediateAsU64() * rhs.GetImmediateAsU64();
-        ReplaceUsesWith(inst, is_32_bit, result);
-    } else if (lhs.IsZero() || rhs.IsZero()) {
             ReplaceUsesWith(inst, is_32_bit, 0);
-    } else if (lhs.IsUnsignedImmediate(1)) {
-        inst.ReplaceUsesWith(rhs);
         } else if (rhs.IsUnsignedImmediate(1)) {
-        inst.ReplaceUsesWith(lhs);
+            inst.ReplaceUsesWith(inst.GetArg(0));
+        }
     }
 }