fix(ssa_fuzzer): nested conditions in loops

tooling/ssa_fuzzer/fuzzer/src/fuzz_lib/base_context.rs

@@ -4,8 +4,8 @@
 use super::options::{ProgramContextOptions, SsaBlockOptions};
 use acvm::FieldElement;
 use acvm::acir::native_types::Witness;
 use libfuzzer_sys::arbitrary;
 use libfuzzer_sys::arbitrary::Arbitrary;
 use noir_ssa_fuzzer::{
     builder::{FuzzerBuilder, FuzzerBuilderError},
     typed_value::{TypedValue, ValueType},
@@ -238,59 +238,6 @@
         let block_else_instruction_block =
             self.instruction_blocks[block_else_idx % self.instruction_blocks.len()].clone();
 
-        // TODO(sn): add support of nested jmp_if in cycles
-        // If the current block is a loop body
-        if self.cycle_bodies_to_iters_ids.contains_key(&self.current_block.block_id) {
-            let CycleInfo { block_iter_id, block_end_id } =
-                self.cycle_bodies_to_iters_ids[&self.current_block.block_id];
-            let current_block = self.stored_blocks[&block_end_id].clone();
-
-            // init then branch
-            let block_then_id = self.insert_ssa_block();
-
-            self.switch_to_block(block_then_id);
-            // context of then branch is shared with predecessor (loop body)
-            current_block.context.clone().insert_instructions(
-                &mut self.acir_builder,
-                &mut self.brillig_builder,
-                &block_then_instruction_block.instructions,
-            );
-            // jump to the loop iter block
-            self.insert_jmp_instruction(block_iter_id, vec![]);
-
-            // init else branch
-            let block_else_id = self.insert_ssa_block();
-            self.switch_to_block(block_else_id);
-            // context of else branch is shared with predecessor (loop body)
-            current_block.context.clone().insert_instructions(
-                &mut self.acir_builder,
-                &mut self.brillig_builder,
-                &block_else_instruction_block.instructions,
-            );
-            // jump to the loop iter block
-            self.insert_jmp_instruction(block_iter_id, vec![]);
-
-            // finalize the loop body with jmp_if
-            self.switch_to_block(self.current_block.block_id);
-            self.current_block.context.finalize_block_with_jmp_if(
-                &mut self.acir_builder,
-                &mut self.brillig_builder,
-                block_then_id,
-                block_else_id,
-            );
-            // remove children blocks not to merge them (they already jumped to iter block)
-            self.current_block.context.children_blocks.pop();
-            self.current_block.context.children_blocks.pop();
-
-            // switch context to the loop end block
-            self.current_block =
-                StoredBlock { context: current_block.context, block_id: block_end_id };
-            self.switch_to_block(self.current_block.block_id);
-            // remove loop body from the map, we don't need it anymore
-            self.cycle_bodies_to_iters_ids.remove(&self.current_block.block_id);
-            return;
-        }
-
         self.store_variables();
 
         if block_then_instruction_block.instructions.len()
@@ -359,7 +306,25 @@
             StoredBlock { context: block_then_context, block_id: block_then_id };
         let else_stored_block =
             StoredBlock { context: block_else_context, block_id: block_else_id };
-        self.not_terminated_blocks.push_back(else_stored_block);
+
+        // if current context is cycle body we define then and else branch as new bodies
+        if self.cycle_bodies_to_iters_ids.contains_key(&self.current_block.block_id) {
+            let CycleInfo { block_iter_id, block_end_id } =
+                self.cycle_bodies_to_iters_ids[&self.current_block.block_id];
+            // block cannot have more than two predecessors
+            // so we create join block that terminated with jmp to iter block
+            // and then terminate then and else blocks with jmp join block in Self::finalize_cycles
+            let block_join_id = self.insert_ssa_block();
+            self.switch_to_block(block_join_id);
+            self.insert_jmp_instruction(block_iter_id, vec![]);
+            self.cycle_bodies_to_iters_ids
+                .insert(block_then_id, CycleInfo { block_iter_id: block_join_id, block_end_id });
+            self.cycle_bodies_to_iters_ids
+                .insert(block_else_id, CycleInfo { block_iter_id: block_join_id, block_end_id });
+            self.cycle_bodies_to_iters_ids.remove(&self.current_block.block_id);
+        } else {
+            self.not_terminated_blocks.push_back(else_stored_block);
+        }
         self.switch_to_block(then_stored_block.block_id);
         self.current_block = then_stored_block.clone();
     }
@@ -569,10 +534,19 @@
             &block_body.instructions,
         );
 
+        let end_context =
+            if self.cycle_bodies_to_iters_ids.contains_key(&self.current_block.block_id) {
+                self.stored_blocks
+                    [&self.cycle_bodies_to_iters_ids[&self.current_block.block_id].block_end_id]
+                    .context
+                    .clone()
+            } else {
+                self.current_block.context.clone()
+            };
         // end block does not share variables with body block, so we copy them from the current block
         let block_end_context = BlockContext::new(
-            self.current_block.context.stored_values.clone(),
-            self.current_block.context.memory_addresses.clone(),
+            end_context.stored_values.clone(),
+            end_context.memory_addresses.clone(),
             block_body_context.parent_blocks_history.clone(),
             SsaBlockOptions::from(self.context_options.clone()),
         );

tooling/ssa_fuzzer/fuzzer/src/fuzz_lib/fuzz_target_lib.rs

@@ -5,11 +5,11 @@
 use super::options::FuzzerOptions;
 use acvm::FieldElement;
 use acvm::acir::native_types::{Witness, WitnessMap};
 use libfuzzer_sys::arbitrary;
 use libfuzzer_sys::arbitrary::Arbitrary;
 use noir_ssa_fuzzer::typed_value::ValueType;
 
 /// Field modulus has 254 bits, and FieldElement::from supports u128, so we use two unsigneds to represent a field element
 /// field = low + high * 2^128
 #[derive(Debug, Clone, Hash, Arbitrary)]
 pub(crate) struct FieldRepresentation {
@@ -428,8 +428,8 @@
     #[test]
     fn test_jmp_if_in_cycle() {
         let arg_2_field = Argument { index: 2, value_type: ValueType::Field };
-        let arg_5_field = Argument { index: 5, value_type: ValueType::Field };
         let arg_6_field = Argument { index: 6, value_type: ValueType::Field };
+        let arg_7_field = Argument { index: 7, value_type: ValueType::Field };
         // v9 = allocate -> &mut Field
         // store v2 at v9
         let add_to_memory_block =
@@ -446,15 +446,15 @@
         let load_mul_set_block = InstructionBlock {
             instructions: vec![
                 Instruction::LoadFromMemory { memory_addr: typed_memory_0 }, // v15 = load v9 -> Field (loaded value is 5th defined field in then block)
-                Instruction::MulChecked { lhs: arg_5_field, rhs: arg_2_field }, // v16 = mul v15, v2 (v16 -- 6th defined field in then block)
-                Instruction::SetToMemory { memory_addr_index: 0, value: arg_6_field }, // store v16 at v9
+                Instruction::MulChecked { lhs: arg_6_field, rhs: arg_2_field }, // v16 = mul v15, v2 (v16 -- 6th defined field in then block)
+                Instruction::SetToMemory { memory_addr_index: 0, value: arg_7_field }, // store v16 at v9
             ],
         };
         let load_add_set_block = InstructionBlock {
             instructions: vec![
                 Instruction::LoadFromMemory { memory_addr: typed_memory_0 }, // v17 = load v9 -> Field (loaded value is 5th defined field in else block)
-                Instruction::AddChecked { lhs: arg_5_field, rhs: arg_2_field }, // v18 = add v17, v2 (v18 -- 6th defined field in else block)
-                Instruction::SetToMemory { memory_addr_index: 0, value: arg_6_field }, // store v18 at v9
+                Instruction::AddChecked { lhs: arg_6_field, rhs: arg_2_field }, // v18 = add v17, v2 (v18 -- 6th defined field in else block)
+                Instruction::SetToMemory { memory_addr_index: 0, value: arg_7_field }, // store v18 at v9
             ],
         };
         let commands = vec![
@@ -509,4 +509,85 @@
             None => panic!("Program failed to execute"),
         }
     }
+
+    /// fn main(x: Field, condition: bool) -> pub Field {
+    /// let mut y = x;
+    /// for i in 0..4 {
+    ///     if condition {
+    ///         for _ in 0..4 {
+    ///             y *= x;
+    ///         }
+    ///     }
+    ///    }
+    /// y
+    /// }
+    #[test]
+    fn test_cycle_if_then_cycle() {
+        // this argument doesn't exists
+        // fuzzer doesn't support empty instruction blocks
+        let dummy_arg = Argument { index: 0, value_type: ValueType::I64 };
+        let arg_2_field = Argument { index: 2, value_type: ValueType::Field };
+        let arg_6_field = Argument { index: 5, value_type: ValueType::Field };
+        let arg_7_field = Argument { index: 6, value_type: ValueType::Field };
+        // v9 = allocate -> &mut Field
+        // store v2 at v9
+        let add_to_memory_block =
+            InstructionBlock { instructions: vec![Instruction::AddToMemory { lhs: arg_2_field }] };
+        let typed_memory_0 = Argument { index: 0, value_type: ValueType::Field };
+        let load_mul_set_block = InstructionBlock {
+            instructions: vec![
+                Instruction::LoadFromMemory { memory_addr: typed_memory_0 }, // v15 = load v9 -> Field (loaded value is 5th defined field in then block)
+                Instruction::MulChecked { lhs: arg_6_field, rhs: arg_2_field }, // v16 = mul v15, v2 (v16 -- 6th defined field in then block)
+                Instruction::SetToMemory { memory_addr_index: 0, value: arg_7_field }, // store v16 at v9
+            ],
+        };
+        let load_block = InstructionBlock {
+            instructions: vec![Instruction::LoadFromMemory { memory_addr: typed_memory_0 }],
+        };
+        let dummy_block = InstructionBlock {
+            instructions: vec![Instruction::AddChecked { lhs: dummy_arg, rhs: dummy_arg }],
+        };
+        let arg_0_boolean = Argument { index: 0, value_type: ValueType::Boolean };
+        let arg_1_boolean = Argument { index: 1, value_type: ValueType::Boolean };
+        let add_boolean_block = InstructionBlock {
+            instructions: vec![Instruction::Or { lhs: arg_0_boolean, rhs: arg_1_boolean }], // jmpif uses last defined boolean variable
+                                                                                            // [initialize_witness_map] func inserts two boolean variables itself, first is true, last is false
+                                                                                            // so by inserting new boolean = first | last, we will get last variable = true
+        };
+        let mut commands = vec![
+            FuzzerCommand::InsertSimpleInstructionBlock { instruction_block_idx: 0 }, // add v2 to memory
+            FuzzerCommand::InsertCycle { block_body_idx: 1, start_iter: 0, end_iter: 4 }, // for i in 0..4 do if ...
+            FuzzerCommand::InsertJmpIfBlock { block_then_idx: 1, block_else_idx: 1 }, // insert if and switch to then branch
+            FuzzerCommand::InsertCycle { block_body_idx: 2, start_iter: 0, end_iter: 4 }, // for i in 0..4 do load_mul_set
+        ];
+        let mut blocks = vec![add_to_memory_block, dummy_block, load_mul_set_block, load_block];
+        let data = FuzzerData {
+            blocks: blocks.to_vec(),
+            commands: commands.clone(),
+            initial_witness: default_witness(),
+            return_instruction_block_idx: 3, // load stored field and return it
+        };
+        // last boolean is false, so this cycle does not execute
+        let result = fuzz_target(data, FuzzerOptions::default());
+        match result {
+            Some(result) => assert_eq!(result, FieldElement::from(2_u32)),
+            None => panic!("Program failed to execute"),
+        }
+
+        // same program but with condition=true
+        commands
+            .insert(0, FuzzerCommand::InsertSimpleInstructionBlock { instruction_block_idx: 4 }); // set last boolean as true
+        blocks.push(add_boolean_block);
+        let data = FuzzerData {
+            blocks: blocks.to_vec(),
+            commands: commands.clone(),
+            initial_witness: default_witness(),
+            return_instruction_block_idx: 3, // load stored field and return it
+        };
+        let result = fuzz_target(data, FuzzerOptions::default());
+        match result {
+            Some(result) => assert_eq!(result, FieldElement::from(2_u32.pow(17))),
+            None => panic!("Program failed to execute"),
+        }
+    }
 }