feat: refactor SSA passes to run on individual functions

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

@@ -2,6 +2,7 @@ use crate::ssa::{
     ir::{
         basic_block::BasicBlockId,
         dfg::DataFlowGraph,
+        function::Function,
         instruction::{Instruction, InstructionId, TerminatorInstruction},
         types::Type::{Array, Slice},
         value::ValueId,
@@ -17,32 +18,38 @@ impl Ssa {
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn array_set_optimization(mut self) -> Self {
         for func in self.functions.values_mut() {
-            let reachable_blocks = func.reachable_blocks();
-
-            if !func.runtime().is_entry_point() {
-                assert_eq!(reachable_blocks.len(), 1, "Expected there to be 1 block remaining in Acir function for array_set optimization");
-            }
-            let mut array_to_last_use = HashMap::default();
-            let mut instructions_to_update = HashSet::default();
-            let mut arrays_from_load = HashSet::default();
-
-            for block in reachable_blocks.iter() {
-                analyze_last_uses(
-                    &func.dfg,
-                    *block,
-                    &mut array_to_last_use,
-                    &mut instructions_to_update,
-                    &mut arrays_from_load,
-                );
-            }
-            for block in reachable_blocks {
-                make_mutable(&mut func.dfg, block, &instructions_to_update);
-            }
+            func.array_set_optimization();
         }
         self
     }
 }
 
+impl Function {
+    pub(crate) fn array_set_optimization(&mut self) {
+        let reachable_blocks = self.reachable_blocks();
+
+        if !self.runtime().is_entry_point() {
+            assert_eq!(reachable_blocks.len(), 1, "Expected there to be 1 block remaining in Acir function for array_set optimization");
+        }
+        let mut array_to_last_use = HashMap::default();
+        let mut instructions_to_update = HashSet::default();
+        let mut arrays_from_load = HashSet::default();
+
+        for block in reachable_blocks.iter() {
+            analyze_last_uses(
+                &self.dfg,
+                *block,
+                &mut array_to_last_use,
+                &mut instructions_to_update,
+                &mut arrays_from_load,
+            );
+        }
+        for block in reachable_blocks {
+            make_mutable(&mut self.dfg, block, &instructions_to_update);
+        }
+    }
+}
+
 /// Builds the set of ArraySet instructions that can be made mutable
 /// because their input value is unused elsewhere afterward.
 fn analyze_last_uses(

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

@@ -20,13 +20,19 @@ impl Ssa {
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn as_slice_optimization(mut self) -> Self {
         for func in self.functions.values_mut() {
-            let known_slice_lengths = known_slice_lengths(func);
-            replace_known_slice_lengths(func, known_slice_lengths);
+            func.as_slice_optimization();
         }
         self
     }
 }
 
+impl Function {
+    pub(crate) fn as_slice_optimization(&mut self) {
+        let known_slice_lengths = known_slice_lengths(self);
+        replace_known_slice_lengths(self, known_slice_lengths);
+    }
+}
+
 fn known_slice_lengths(func: &Function) -> HashMap<InstructionId, usize> {
     let mut known_slice_lengths = HashMap::default();
     for block_id in func.reachable_blocks() {

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

@@ -26,22 +26,31 @@ impl Ssa {
         mut self,
     ) -> Result<Ssa, RuntimeError> {
         for function in self.functions.values_mut() {
-            for block in function.reachable_blocks() {
-                // Unfortunately we can't just use instructions.retain(...) here since
-                // check_instruction can also return an error
-                let instructions = function.dfg[block].take_instructions();
-                let mut filtered_instructions = Vec::with_capacity(instructions.len());
+            function.evaluate_static_assert_and_assert_constant()?;
+        }
+        Ok(self)
+    }
+}
 
-                for instruction in instructions {
-                    if check_instruction(function, instruction)? {
-                        filtered_instructions.push(instruction);
-                    }
-                }
+impl Function {
+    pub(crate) fn evaluate_static_assert_and_assert_constant(
+        &mut self,
+    ) -> Result<(), RuntimeError> {
+        for block in self.reachable_blocks() {
+            // Unfortunately we can't just use instructions.retain(...) here since
+            // check_instruction can also return an error
+            let instructions = self.dfg[block].take_instructions();
+            let mut filtered_instructions = Vec::with_capacity(instructions.len());
 
-                *function.dfg[block].instructions_mut() = filtered_instructions;
+            for instruction in instructions {
+                if check_instruction(self, instruction)? {
+                    filtered_instructions.push(instruction);
+                }
             }
+
+            *self.dfg[block].instructions_mut() = filtered_instructions;
         }
-        Ok(self)
+        Ok(())
     }
 }
 

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

@@ -0,0 +1,159 @@
+use std::collections::HashMap;
+
+use crate::ssa::{
+    ir::{
+        function::Function,
+        instruction::{Instruction, InstructionId},
+    },
+    ssa_gen::Ssa,
+};
+
+impl Ssa {
+    /// A simple SSA pass to go through each instruction and move every `Instruction::Constrain` to immediately
+    /// after when all of its inputs are available.
+    #[tracing::instrument(level = "trace", skip(self))]
+    pub(crate) fn bubble_up_constrains(mut self) -> Ssa {
+        for function in self.functions.values_mut() {
+            function.bubble_up_constrains();
+        }
+        self
+    }
+}
+
+impl Function {
+    pub(crate) fn bubble_up_constrains(&mut self) {
+        for block in self.reachable_blocks() {
+            let instructions = self.dfg[block].take_instructions();
+            let mut filtered_instructions = Vec::with_capacity(instructions.len());
+
+            // Multiple constrains can bubble up to sit under a single instruction. We want to maintain the ordering of these constraints,
+            // so we need to keep track of how many constraints are attached to a given instruction.
+            // Some assertions don't operate on instruction results, so we use Option so we also track the None case
+            let mut inserted_at_instruction: HashMap<Option<InstructionId>, usize> =
+                HashMap::with_capacity(instructions.len());
+
+            let dfg = &self.dfg;
+            for instruction in instructions {
+                let (lhs, rhs) = match dfg[instruction] {
+                    Instruction::Constrain(lhs, rhs, ..) => (lhs, rhs),
+                    _ => {
+                        filtered_instructions.push(instruction);
+                        continue;
+                    }
+                };
+
+                let last_instruction_that_creates_inputs = filtered_instructions
+                    .iter()
+                    .rev()
+                    .position(|&instruction_id| {
+                        let results = dfg.instruction_results(instruction_id).to_vec();
+                        results.contains(&lhs) || results.contains(&rhs)
+                    })
+                    // We iterate through the previous instructions in reverse order so the index is from the
+                    // back of the vector
+                    .map(|reversed_index| filtered_instructions.len() - reversed_index - 1);
+
+                let insertion_index = last_instruction_that_creates_inputs
+                    .map(|index| {
+                        // We want to insert just after the last instruction that creates the inputs
+                        index + 1
+                    })
+                    // If it doesn't depend from the previous instructions, then we insert at the start
+                    .unwrap_or_default();
+
+                let already_inserted_for_this_instruction = inserted_at_instruction
+                    .entry(
+                        last_instruction_that_creates_inputs
+                            .map(|index| filtered_instructions[index]),
+                    )
+                    .or_default();
+
+                filtered_instructions
+                    .insert(insertion_index + *already_inserted_for_this_instruction, instruction);
+
+                *already_inserted_for_this_instruction += 1;
+            }
+
+            *self.dfg[block].instructions_mut() = filtered_instructions;
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::ssa::{
+        function_builder::FunctionBuilder,
+        ir::{
+            instruction::{Binary, BinaryOp, Instruction},
+            map::Id,
+            types::Type,
+        },
+    };
+
+    #[test]
+    fn check_bubble_up_constrains() {
+        // fn main f0 {
+        //   b0(v0: Field):
+        //     v1 = add v0, Field 1
+        //     v2 = add v1, Field 1
+        //     constrain v0 == Field 1 'With message'
+        //     constrain v2 == Field 3
+        //     constrain v0 == Field 1
+        //     constrain v1 == Field 2
+        //     constrain v1 == Field 2 'With message'
+        // }
+        //
+        let main_id = Id::test_new(0);
+
+        // Compiling main
+        let mut builder = FunctionBuilder::new("main".into(), main_id);
+        let v0 = builder.add_parameter(Type::field());
+
+        let one = builder.field_constant(1u128);
+        let two = builder.field_constant(2u128);
+        let three = builder.field_constant(3u128);
+
+        let v1 = builder.insert_binary(v0, BinaryOp::Add, one);
+        let v2 = builder.insert_binary(v1, BinaryOp::Add, one);
+        builder.insert_constrain(v0, one, Some("With message".to_string().into()));
+        builder.insert_constrain(v2, three, None);
+        builder.insert_constrain(v0, one, None);
+        builder.insert_constrain(v1, two, None);
+        builder.insert_constrain(v1, two, Some("With message".to_string().into()));
+        builder.terminate_with_return(vec![]);
+
+        let ssa = builder.finish();
+
+        // Expected output:
+        //
+        // fn main f0 {
+        //   b0(v0: Field):
+        //     constrain v0 == Field 1 'With message'
+        //     constrain v0 == Field 1
+        //     v1 = add v0, Field 1
+        //     constrain v1 == Field 2
+        //     constrain v1 == Field 2 'With message'
+        //     v2 = add v1, Field 1
+        //     constrain v2 == Field 3
+        // }
+        //
+        let ssa = ssa.bubble_up_constrains();
+        let main = ssa.main();
+        let block = &main.dfg[main.entry_block()];
+        assert_eq!(block.instructions().len(), 7);
+
+        let expected_instructions = vec![
+            Instruction::Constrain(v0, one, Some("With message".to_string().into())),
+            Instruction::Constrain(v0, one, None),
+            Instruction::Binary(Binary { lhs: v0, rhs: one, operator: BinaryOp::Add }),
+            Instruction::Constrain(v1, two, None),
+            Instruction::Constrain(v1, two, Some("With message".to_string().into())),
+            Instruction::Binary(Binary { lhs: v1, rhs: one, operator: BinaryOp::Add }),
+            Instruction::Constrain(v2, three, None),
+        ];
+
+        for (index, instruction) in block.instructions().iter().enumerate() {
+            assert_eq!(&main.dfg[*instruction], &expected_instructions[index]);
+        }
+    }
+}

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

@@ -44,7 +44,7 @@ impl Ssa {
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn fold_constants(mut self) -> Ssa {
         for function in self.functions.values_mut() {
-            constant_fold(function, false);
+            function.constant_fold(false);
         }
         self
     }
@@ -57,25 +57,27 @@ impl Ssa {
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn fold_constants_using_constraints(mut self) -> Ssa {
         for function in self.functions.values_mut() {
-            constant_fold(function, true);
+            function.constant_fold(true);
         }
         self
     }
 }
 
-/// The structure of this pass is simple:
-/// Go through each block and re-insert all instructions.
-fn constant_fold(function: &mut Function, use_constraint_info: bool) {
-    let mut context = Context { use_constraint_info, ..Default::default() };
-    context.block_queue.push(function.entry_block());
+impl Function {
+    /// The structure of this pass is simple:
+    /// Go through each block and re-insert all instructions.
+    pub(crate) fn constant_fold(&mut self, use_constraint_info: bool) {
+        let mut context = Context { use_constraint_info, ..Default::default() };
+        context.block_queue.push(self.entry_block());
 
-    while let Some(block) = context.block_queue.pop() {
-        if context.visited_blocks.contains(&block) {
-            continue;
-        }
+        while let Some(block) = context.block_queue.pop() {
+            if context.visited_blocks.contains(&block) {
+                continue;
+            }
 
-        context.visited_blocks.insert(block);
-        context.fold_constants_in_block(function, block);
+            context.visited_blocks.insert(block);
+            context.fold_constants_in_block(self, block);
+        }
     }
 }