chore(experimental): Function::simple_optimization for SSA optimizations

compiler/noirc_evaluator/src/ssa/opt/as_slice_length.rs

@@ -3,7 +3,7 @@ use crate::ssa::{
         function::Function,
         instruction::{Instruction, Intrinsic},
         types::{NumericType, Type},
-        value::{Value, ValueMapping},
+        value::Value,
     },
     ssa_gen::Ssa,
 };
@@ -27,46 +27,30 @@ impl Ssa {
 
 impl Function {
     pub(crate) fn as_slice_optimization(&mut self) {
-        let mut values_to_replace = ValueMapping::default();
+        self.simple_reachable_blocks_optimization(|context| {
+            let instruction_id = context.instruction_id;
+            let instruction = context.instruction();
 
-        for block in self.reachable_blocks() {
-            let instruction_ids = self.dfg[block].take_instructions();
-            for instruction_id in &instruction_ids {
-                let (target_func, first_argument) = {
-                    let instruction = &mut self.dfg[*instruction_id];
-                    instruction.replace_values(&values_to_replace);
+            let (target_func, arguments) = match &instruction {
+                Instruction::Call { func, arguments } => (func, arguments),
+                _ => return,
+            };
 
-                    let (target_func, arguments) = match &instruction {
-                        Instruction::Call { func, arguments } => (func, arguments),
-                        _ => continue,
-                    };
+            let Value::Intrinsic(Intrinsic::AsSlice) = context.dfg[*target_func] else {
+                return;
+            };
 
-                    (*target_func, arguments.first().copied())
-                };
+            let first_argument = arguments.first().unwrap();
+            let array_typ = context.dfg.type_of_value(*first_argument);
+            let Type::Array(_, length) = array_typ else {
+                unreachable!("AsSlice called with non-array {}", array_typ);
+            };
 
-                match &self.dfg[target_func] {
-                    Value::Intrinsic(Intrinsic::AsSlice) => {
-                        let first_argument = first_argument.unwrap();
-                        let array_typ = self.dfg.type_of_value(first_argument);
-                        if let Type::Array(_, length) = array_typ {
-                            let call_returns = self.dfg.instruction_results(*instruction_id);
-                            let original_slice_length = call_returns[0];
-                            let known_length =
-                                self.dfg.make_constant(length.into(), NumericType::length_type());
-                            values_to_replace.insert(original_slice_length, known_length);
-                        } else {
-                            unreachable!("AsSlice called with non-array {}", array_typ);
-                        }
-                    }
-                    _ => continue,
-                };
-            }
-
-            *self.dfg[block].instructions_mut() = instruction_ids;
-            self.dfg.replace_values_in_block_terminator(block, &values_to_replace);
-        }
-
-        self.dfg.data_bus.replace_values(&values_to_replace);
+            let call_returns = context.dfg.instruction_results(instruction_id);
+            let original_slice_length = call_returns[0];
+            let known_length = context.dfg.make_constant(length.into(), NumericType::length_type());
+            context.replace_value(original_slice_length, known_length);
+        });
     }
 }
 

compiler/noirc_evaluator/src/ssa/opt/brillig_array_gets.rs

@@ -10,8 +10,6 @@
 //! is unnecessary as we already know the index. This pass looks for such array operations
 //! with constant indices and replaces their index with the appropriate offset.
 
-use fxhash::FxHashMap as HashMap;
-
 use crate::{
     brillig::brillig_ir::BRILLIG_MEMORY_ADDRESSING_BIT_SIZE,
     ssa::{
@@ -39,46 +37,34 @@ impl Ssa {
 
 impl Function {
     pub(super) fn brillig_array_gets(&mut self) {
-        let reachable_blocks = self.reachable_blocks();
-
-        let mut instructions_to_update = HashMap::default();
-        for block_id in reachable_blocks.into_iter() {
-            for instruction_id in self.dfg[block_id].instructions() {
-                if let Instruction::ArrayGet { array, index } = self.dfg[*instruction_id] {
-                    if self.dfg.is_constant(index) {
-                        instructions_to_update.insert(
-                            *instruction_id,
-                            (Instruction::ArrayGet { array, index }, block_id),
-                        );
-                    }
-                }
+        self.simple_reachable_blocks_optimization(|context| {
+            let instruction = context.instruction();
+            let Instruction::ArrayGet { array, index } = instruction else {
+                return;
+            };
+
+            let array = *array;
+            let index = *index;
+            if !context.dfg.is_constant(index) {
+                return;
             }
-        }
 
-        for (instruction_id, _) in instructions_to_update {
-            let new_instruction = match self.dfg[instruction_id] {
-                Instruction::ArrayGet { array, index } => {
-                    let index_constant =
-                        self.dfg.get_numeric_constant(index).expect("ICE: Expected constant index");
-                    let offset = if matches!(self.dfg.type_of_value(array), Type::Array(..)) {
-                        // Brillig arrays are [RC, ...items]
-                        1u128
-                    } else {
-                        // Brillig vectors are [RC, Size, Capacity, ...items]
-                        3u128
-                    };
-                    let index = self.dfg.make_constant(
-                        index_constant + offset.into(),
-                        NumericType::unsigned(BRILLIG_MEMORY_ADDRESSING_BIT_SIZE),
-                    );
-                    Instruction::ArrayGet { array, index }
-                }
-                _ => {
-                    continue;
-                }
+            let index_constant =
+                context.dfg.get_numeric_constant(index).expect("ICE: Expected constant index");
+            let offset = if matches!(context.dfg.type_of_value(array), Type::Array(..)) {
+                // Brillig arrays are [RC, ...items]
+                1u128
+            } else {
+                // Brillig vectors are [RC, Size, Capacity, ...items]
+                3u128
             };
-            self.dfg[instruction_id] = new_instruction;
-        }
+            let index = context.dfg.make_constant(
+                index_constant + offset.into(),
+                NumericType::unsigned(BRILLIG_MEMORY_ADDRESSING_BIT_SIZE),
+            );
+            let new_instruction = Instruction::ArrayGet { array, index };
+            context.replace_current_instruction_with(new_instruction);
+        });
     }
 }
 

compiler/noirc_evaluator/src/ssa/opt/check_u128_mul_overflow.rs

@@ -4,7 +4,6 @@ use crate::ssa::{
     ir::{
         basic_block::BasicBlockId,
         call_stack::CallStackId,
-        dfg::DataFlowGraph,
         function::Function,
         instruction::{Binary, BinaryOp, ConstrainError, Instruction},
         types::NumericType,
@@ -13,6 +12,8 @@ use crate::ssa::{
     ssa_gen::Ssa,
 };
 
+use super::simple_optimization::SimpleOptimizationContext;
+
 impl Ssa {
     /// An SSA pass that checks that multiplying two u128 doesn't overflow because
     /// both operands are greater or equal than 2^64.
@@ -34,40 +35,40 @@ impl Function {
             return;
         }
 
-        for block in self.reachable_blocks() {
-            let instructions = self.dfg[block].take_instructions();
-
-            for instruction in instructions {
-                self.dfg[block].insert_instruction(instruction);
-
-                let Instruction::Binary(Binary {
-                    lhs,
-                    rhs,
-                    operator: BinaryOp::Mul { unchecked: false },
-                }) = &self.dfg[instruction]
-                else {
-                    continue;
-                };
-
-                let binary_type = self.dfg.type_of_value(*lhs).unwrap_numeric();
-                let NumericType::Unsigned { bit_size: 128 } = binary_type else {
-                    continue;
-                };
-
-                let call_stack = self.dfg.get_instruction_call_stack_id(instruction);
-                check_u128_mul_overflow(*lhs, *rhs, block, &mut self.dfg, call_stack);
-            }
-        }
+        self.simple_reachable_blocks_optimization(|context| {
+            context.insert_current_instruction();
+
+            let block_id = context.block_id;
+            let instruction_id = context.instruction_id;
+            let instruction = context.instruction();
+            let Instruction::Binary(Binary {
+                lhs,
+                rhs,
+                operator: BinaryOp::Mul { unchecked: false },
+            }) = instruction
+            else {
+                return;
+            };
+
+            let binary_type = context.dfg.type_of_value(*lhs).unwrap_numeric();
+            let NumericType::Unsigned { bit_size: 128 } = binary_type else {
+                return;
+            };
+
+            let call_stack = context.dfg.get_instruction_call_stack_id(instruction_id);
+            check_u128_mul_overflow(*lhs, *rhs, block_id, context, call_stack);
+        });
     }
 }
 
 fn check_u128_mul_overflow(
     lhs: ValueId,
     rhs: ValueId,
     block: BasicBlockId,
-    dfg: &mut DataFlowGraph,
+    context: &mut SimpleOptimizationContext<'_, '_>,
     call_stack: CallStackId,
 ) {
+    let dfg = &mut context.dfg;
     let lhs_value = dfg.get_numeric_constant(lhs);
     let rhs_value = dfg.get_numeric_constant(rhs);
 

compiler/noirc_evaluator/src/ssa/opt/make_constrain_not_equal.rs

@@ -30,41 +30,56 @@ impl Function {
             return;
         }
 
-        for block in self.reachable_blocks() {
-            let instructions = self.dfg[block].instructions().to_vec();
+        self.simple_reachable_blocks_optimization(|context| {
+            let instruction = context.instruction();
 
-            for instruction in instructions {
-                let constrain_ne: Instruction = match &self.dfg[instruction] {
-                    Instruction::Constrain(lhs, rhs, msg) => {
-                        if self
-                            .dfg
-                            .get_numeric_constant(*rhs)
-                            .is_some_and(|constant| constant.is_zero())
-                        {
-                            if let Value::Instruction { instruction, .. } = &self.dfg[*lhs] {
-                                if let Instruction::Binary(Binary {
-                                    lhs,
-                                    rhs,
-                                    operator: BinaryOp::Eq,
-                                    ..
-                                }) = self.dfg[*instruction]
-                                {
-                                    Instruction::ConstrainNotEqual(lhs, rhs, msg.clone())
-                                } else {
-                                    continue;
-                                }
-                            } else {
-                                continue;
-                            }
-                        } else {
-                            continue;
-                        }
-                    }
-                    _ => continue,
-                };
+            let Instruction::Constrain(lhs, rhs, msg) = instruction else {
+                return;
+            };
 
-                self.dfg[instruction] = constrain_ne;
+            if !context.dfg.get_numeric_constant(*rhs).is_some_and(|constant| constant.is_zero()) {
+                return;
             }
+
+            let Value::Instruction { instruction, .. } = &context.dfg[*lhs] else {
+                return;
+            };
+
+            let Instruction::Binary(Binary { lhs, rhs, operator: BinaryOp::Eq, .. }) =
+                context.dfg[*instruction]
+            else {
+                return;
+            };
+
+            let new_instruction = Instruction::ConstrainNotEqual(lhs, rhs, msg.clone());
+            context.replace_current_instruction_with(new_instruction);
+        });
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{assert_ssa_snapshot, ssa::ssa_gen::Ssa};
+
+    #[test]
+    fn test_make_constrain_not_equals() {
+        let src = "
+        acir(inline) fn main f1 {
+          b0(v0: Field, v1: Field):
+            v2 = eq v0, v1
+            constrain v2 == u1 0
+            return
+        }
+        ";
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.make_constrain_not_equal_instructions();
+        assert_ssa_snapshot!(ssa, @r"
+        acir(inline) fn main f0 {
+          b0(v0: Field, v1: Field):
+            v2 = eq v0, v1
+            constrain v0 != v1
+            return
         }
+        ");
     }
 }

compiler/noirc_evaluator/src/ssa/opt/mod.rs

@@ -29,6 +29,7 @@ mod remove_enable_side_effects;
 mod remove_if_else;
 mod remove_truncate_after_range_check;
 mod remove_unreachable;
+mod simple_optimization;
 mod simplify_cfg;
 mod unrolling;