From eff2d9ff5a852c582c5406b641263297e376972e Mon Sep 17 00:00:00 2001
From: TomAFrench <tom@tomfren.ch>
Date: Thu, 8 Feb 2024 15:01:57 +0000
Subject: [PATCH] feat: convert constant right shifts into divisions

---
 .../src/ssa/ir/instruction/binary.rs          | 23 +++++++++++++++++--
 .../src/ssa/opt/remove_bit_shifts.rs          |  1 +
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction/binary.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction/binary.rs
index 4767da75ee..552be9420d 100644
--- a/compiler/noirc_evaluator/src/ssa/ir/instruction/binary.rs
+++ b/compiler/noirc_evaluator/src/ssa/ir/instruction/binary.rs
@@ -221,8 +221,27 @@ impl Binary {
                     return SimplifyResult::SimplifiedTo(zero);
                 }
             }
-            BinaryOp::Shl | BinaryOp::Shr => (),
-        }
+            BinaryOp::Shl => return SimplifyResult::None,
+            BinaryOp::Shr => {
+                // Bit shifts by constants can be treated as divisions.
+                if let Some(rhs_const) = rhs {
+                    if rhs_const >= FieldElement::from(operand_type.bit_size() as u128) {
+                        // Shifting by the full width of the operand type, any `lhs` goes to zero.
+                        let zero = dfg.make_constant(FieldElement::zero(), operand_type);
+                        return SimplifyResult::SimplifiedTo(zero);
+                    }
+
+                    // `two_pow_rhs` is limited to be at most `2 ^ {operand_bitsize - 1}` so it fits in `operand_type`.
+                    let two_pow_rhs = FieldElement::from(2u128).pow(&rhs_const);
+                    let two_pow_rhs = dfg.make_constant(two_pow_rhs, operand_type);
+                    return SimplifyResult::SimplifiedToInstruction(Instruction::binary(
+                        BinaryOp::Div,
+                        self.lhs,
+                        two_pow_rhs,
+                    ));
+                }
+            }
+        };
         SimplifyResult::None
     }
 }
diff --git a/compiler/noirc_evaluator/src/ssa/opt/remove_bit_shifts.rs b/compiler/noirc_evaluator/src/ssa/opt/remove_bit_shifts.rs
index 684f888160..a71a42d575 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/remove_bit_shifts.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/remove_bit_shifts.rs
@@ -140,6 +140,7 @@ impl Context<'_> {
         rhs: ValueId,
         bit_size: u32,
     ) -> ValueId {
+        let lhs_typ = self.function.dfg.type_of_value(lhs);
         let base = self.field_constant(FieldElement::from(2_u128));
         // we can safely cast to unsigned because overflow_checks prevent bit-shift with a negative value
         let rhs_unsigned = self.insert_cast(rhs, Type::unsigned(bit_size));