Merge 5c40387 into 199d129

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

@@ -1,4 +1,9 @@
-use crate::ssa::{ir::instruction::Instruction, ssa_gen::Ssa};
+use std::collections::HashMap;
+
+use crate::ssa::{
+    ir::instruction::{Instruction, InstructionId},
+    ssa_gen::Ssa,
+};
 
 impl Ssa {
     /// A simple SSA pass to go through each instruction and move every `Instruction::Constrain` to immediately
@@ -10,6 +15,12 @@ impl Ssa {
                 let instructions = function.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 = &function.dfg;
                 for instruction in instructions {
                     let (lhs, rhs) = match dfg[instruction] {
@@ -20,19 +31,38 @@ impl Ssa {
                         }
                     };
 
-                    let index = filtered_instructions
+                    let last_instruction_that_creates_inputs = filtered_instructions
                         .iter()
                         .rev()
-                        .position(|instruction_id| {
-                            let results = dfg.instruction_results(*instruction_id).to_vec();
+                        .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. Subtract from vector length to get correct index.
-                        .map(|reversed_index| filtered_instructions.len() - reversed_index)
-                        .unwrap_or(0);
+                        // 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(index, instruction);
+                    filtered_instructions.insert(
+                        insertion_index + *already_inserted_for_this_instruction,
+                        instruction,
+                    );
+
+                    *already_inserted_for_this_instruction += 1;
                 }
 
                 *function.dfg[block].instructions_mut() = filtered_instructions;
@@ -41,3 +71,83 @@ impl Ssa {
         self
     }
 }
+
+#[cfg(test)]
+mod test {
+    use crate::ssa::{
+        function_builder::FunctionBuilder,
+        ir::{
+            function::RuntimeType,
+            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, RuntimeType::Acir);
+        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()));
+        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()));
+        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())),
+            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())),
+            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]);
+        }
+    }
+}