diff --git a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs index c2bc1e32cd6..37128508086 100644 --- a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs +++ b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs @@ -938,7 +938,7 @@ impl<'block> BrilligBlock<'block> { if let Some((index_register, value_variable)) = opt_index_and_value { // Then set the value in the newly created array - self.store_variable_in_array( + self.brillig_context.codegen_store_variable_in_array( destination_pointer, SingleAddrVariable::new_usize(index_register), value_variable, @@ -949,47 +949,6 @@ impl<'block> BrilligBlock<'block> { source_size_as_register } - pub(crate) fn store_variable_in_array_with_ctx( - ctx: &mut BrilligContext, - destination_pointer: MemoryAddress, - index_register: SingleAddrVariable, - value_variable: BrilligVariable, - ) { - match value_variable { - BrilligVariable::SingleAddr(value_variable) => { - ctx.codegen_array_set(destination_pointer, index_register, value_variable.address); - } - BrilligVariable::BrilligArray(_) => { - let reference: MemoryAddress = ctx.allocate_register(); - ctx.codegen_allocate_array_reference(reference); - ctx.codegen_store_variable(reference, value_variable); - ctx.codegen_array_set(destination_pointer, index_register, reference); - ctx.deallocate_register(reference); - } - BrilligVariable::BrilligVector(_) => { - let reference = ctx.allocate_register(); - ctx.codegen_allocate_vector_reference(reference); - ctx.codegen_store_variable(reference, value_variable); - ctx.codegen_array_set(destination_pointer, index_register, reference); - ctx.deallocate_register(reference); - } - } - } - - pub(crate) fn store_variable_in_array( - &mut self, - destination_pointer: MemoryAddress, - index_variable: SingleAddrVariable, - value_variable: BrilligVariable, - ) { - Self::store_variable_in_array_with_ctx( - self.brillig_context, - destination_pointer, - index_variable, - value_variable, - ); - } - /// Convert the SSA slice operations to brillig slice operations fn convert_ssa_slice_intrinsic_call( &mut self, @@ -1748,10 +1707,19 @@ impl<'block> BrilligBlock<'block> { subitem_to_repeat_variables.push(self.convert_ssa_value(subitem_id, dfg)); } - let data_length_variable = self + // Initialize loop bound with the array length + let end_pointer_variable = self .brillig_context .make_usize_constant_instruction((item_count * item_types.len()).into()); + // Add the pointer to the array length + self.brillig_context.memory_op_instruction( + end_pointer_variable.address, + pointer, + end_pointer_variable.address, + BrilligBinaryOp::Add, + ); + // If this is an array with complex subitems, we need a custom step in the loop to write all the subitems while iterating. if item_types.len() > 1 { let step_variable = @@ -1760,33 +1728,54 @@ impl<'block> BrilligBlock<'block> { let subitem_pointer = SingleAddrVariable::new_usize(self.brillig_context.allocate_register()); - let initializer_fn = |ctx: &mut BrilligContext<_>, iterator: SingleAddrVariable| { - ctx.mov_instruction(subitem_pointer.address, iterator.address); - for subitem in subitem_to_repeat_variables.into_iter() { - Self::store_variable_in_array_with_ctx(ctx, pointer, subitem_pointer, subitem); - ctx.codegen_usize_op_in_place(subitem_pointer.address, BrilligBinaryOp::Add, 1); - } - }; + let one = self.brillig_context.make_usize_constant_instruction(1_usize.into()); + + // Initializes a single subitem + let initializer_fn = + |ctx: &mut BrilligContext<_>, subitem_start_pointer: SingleAddrVariable| { + ctx.mov_instruction(subitem_pointer.address, subitem_start_pointer.address); + for (subitem_index, subitem) in + subitem_to_repeat_variables.into_iter().enumerate() + { + ctx.codegen_store_variable_in_pointer(subitem_pointer.address, subitem); + if subitem_index != item_types.len() - 1 { + ctx.memory_op_instruction( + subitem_pointer.address, + one.address, + subitem_pointer.address, + BrilligBinaryOp::Add, + ); + } + } + }; - self.brillig_context.codegen_loop_with_bound_and_step( - data_length_variable.address, - step_variable.address, + // for (let subitem_start_pointer = pointer; subitem_start_pointer < pointer + data_length; subitem_start_pointer += step) { initializer_fn(iterator) } + self.brillig_context.codegen_for_loop( + Some(pointer), + end_pointer_variable.address, + Some(step_variable.address), initializer_fn, ); self.brillig_context.deallocate_single_addr(step_variable); self.brillig_context.deallocate_single_addr(subitem_pointer); + self.brillig_context.deallocate_single_addr(one); } else { let subitem = subitem_to_repeat_variables.into_iter().next().unwrap(); - let initializer_fn = |ctx: &mut _, iterator_register| { - Self::store_variable_in_array_with_ctx(ctx, pointer, iterator_register, subitem); + let initializer_fn = |ctx: &mut BrilligContext<_>, item_pointer: SingleAddrVariable| { + ctx.codegen_store_variable_in_pointer(item_pointer.address, subitem); }; - self.brillig_context.codegen_loop(data_length_variable.address, initializer_fn); + // for (let item_pointer = pointer; item_pointer < pointer + data_length; item_pointer += 1) { initializer_fn(iterator) } + self.brillig_context.codegen_for_loop( + Some(pointer), + end_pointer_variable.address, + None, + initializer_fn, + ); } - - self.brillig_context.deallocate_single_addr(data_length_variable); + self.brillig_context.deallocate_single_addr(end_pointer_variable); } fn initialize_constant_array_comptime( @@ -1796,22 +1785,29 @@ impl<'block> BrilligBlock<'block> { pointer: MemoryAddress, ) { // Allocate a register for the iterator - let iterator_register = - self.brillig_context.make_usize_constant_instruction(0_usize.into()); + let write_pointer_register = self.brillig_context.allocate_register(); + let one = self.brillig_context.make_usize_constant_instruction(1_usize.into()); - for element_id in data.iter() { + self.brillig_context.mov_instruction(write_pointer_register, pointer); + + for (element_idx, element_id) in data.iter().enumerate() { let element_variable = self.convert_ssa_value(*element_id, dfg); // Store the item in memory - self.store_variable_in_array(pointer, iterator_register, element_variable); - // Increment the iterator - self.brillig_context.codegen_usize_op_in_place( - iterator_register.address, - BrilligBinaryOp::Add, - 1, - ); + self.brillig_context + .codegen_store_variable_in_pointer(write_pointer_register, element_variable); + if element_idx != data.len() - 1 { + // Increment the write_pointer_register + self.brillig_context.memory_op_instruction( + write_pointer_register, + one.address, + write_pointer_register, + BrilligBinaryOp::Add, + ); + } } - self.brillig_context.deallocate_single_addr(iterator_register); + self.brillig_context.deallocate_register(write_pointer_register); + self.brillig_context.deallocate_single_addr(one); } /// Converts an SSA `ValueId` into a `MemoryAddress`. Initializes if necessary. @@ -1896,7 +1892,7 @@ impl<'block> BrilligBlock<'block> { let inner_array = self.allocate_nested_array(element_type, None); let idx = self.brillig_context.make_usize_constant_instruction(index.into()); - self.store_variable_in_array(array.pointer, idx, inner_array); + self.brillig_context.codegen_store_variable_in_array(array.pointer, idx, inner_array); } Type::Slice(_) => unreachable!("ICE: unsupported slice type in allocate_nested_array(), expects an array or a numeric type"), _ => (), diff --git a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs index d17b15a13b5..55679432b1f 100644 --- a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs +++ b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs @@ -38,7 +38,11 @@ impl<'block> BrilligBlock<'block> { target_index.address, BrilligBinaryOp::Add, ); - self.store_variable_in_array(target_vector.pointer, target_index, *variable); + self.brillig_context.codegen_store_variable_in_array( + target_vector.pointer, + target_index, + *variable, + ); self.brillig_context.deallocate_single_addr(target_index); } } @@ -79,7 +83,11 @@ impl<'block> BrilligBlock<'block> { // Then we write the items to insert at the start for (index, variable) in variables_to_insert.iter().enumerate() { let target_index = self.brillig_context.make_usize_constant_instruction(index.into()); - self.store_variable_in_array(target_vector.pointer, target_index, *variable); + self.brillig_context.codegen_store_variable_in_array( + target_vector.pointer, + target_index, + *variable, + ); self.brillig_context.deallocate_single_addr(target_index); } @@ -239,7 +247,11 @@ impl<'block> BrilligBlock<'block> { target_index.address, BrilligBinaryOp::Add, ); - self.store_variable_in_array(target_vector.pointer, target_index, *variable); + self.brillig_context.codegen_store_variable_in_array( + target_vector.pointer, + target_index, + *variable, + ); self.brillig_context.deallocate_single_addr(target_index); } diff --git a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_control_flow.rs b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_control_flow.rs index 5741089a497..8e52cf072b4 100644 --- a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_control_flow.rs +++ b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_control_flow.rs @@ -38,15 +38,28 @@ impl BrilligContext { self.stop_instruction(); } - /// This codegen will issue a loop do for (let iterator_register = 0; i < loop_bound; i += step) + /// This codegen will issue a loop for (let iterator_register = loop_start; i < loop_bound; i += step) /// The body of the loop should be issued by the caller in the on_iteration closure. - pub(crate) fn codegen_loop_with_bound_and_step( + pub(crate) fn codegen_for_loop( &mut self, + loop_start: Option, // Defaults to zero loop_bound: MemoryAddress, - step: MemoryAddress, + step: Option, // Defaults to 1 on_iteration: impl FnOnce(&mut BrilligContext, SingleAddrVariable), ) { - let iterator_register = self.make_usize_constant_instruction(0_u128.into()); + let iterator_register = if let Some(loop_start) = loop_start { + let iterator_register = SingleAddrVariable::new_usize(self.allocate_register()); + self.mov_instruction(iterator_register.address, loop_start); + iterator_register + } else { + self.make_usize_constant_instruction(0_usize.into()) + }; + + let step_register = if let Some(step) = step { + step + } else { + self.make_usize_constant_instruction(1_usize.into()).address + }; let (loop_section, loop_label) = self.reserve_next_section_label(); self.enter_section(loop_section); @@ -76,7 +89,7 @@ impl BrilligContext { // Add step to the iterator register self.memory_op_instruction( iterator_register.address, - step, + step_register, iterator_register.address, BrilligBinaryOp::Add, ); @@ -89,18 +102,20 @@ impl BrilligContext { // Deallocate our temporary registers self.deallocate_single_addr(iterator_less_than_iterations); self.deallocate_single_addr(iterator_register); + // Only deallocate step if we allocated it + if step.is_none() { + self.deallocate_register(step_register); + } } - /// This codegen will issue a loop that will iterate iteration_count times + /// This codegen will issue a loop that will iterate from 0 to iteration_count /// The body of the loop should be issued by the caller in the on_iteration closure. pub(crate) fn codegen_loop( &mut self, iteration_count: MemoryAddress, on_iteration: impl FnOnce(&mut BrilligContext, SingleAddrVariable), ) { - let step = self.make_usize_constant_instruction(1_u128.into()); - self.codegen_loop_with_bound_and_step(iteration_count, step.address, on_iteration); - self.deallocate_single_addr(step); + self.codegen_for_loop(None, iteration_count, None, on_iteration); } /// This codegen will issue an if-then branch that will check if the condition is true diff --git a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_memory.rs b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_memory.rs index 81c1c3847b1..d20f736ee6d 100644 --- a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_memory.rs +++ b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_memory.rs @@ -112,6 +112,50 @@ impl BrilligContext { self.deallocate_register(index_of_element_in_memory); } + pub(crate) fn codegen_store_variable_in_array( + &mut self, + array_pointer: MemoryAddress, + index: SingleAddrVariable, + value_variable: BrilligVariable, + ) { + assert!(index.bit_size == BRILLIG_MEMORY_ADDRESSING_BIT_SIZE); + let final_pointer_register = self.allocate_register(); + self.memory_op_instruction( + array_pointer, + index.address, + final_pointer_register, + BrilligBinaryOp::Add, + ); + self.codegen_store_variable_in_pointer(final_pointer_register, value_variable); + self.deallocate_register(final_pointer_register); + } + + pub(crate) fn codegen_store_variable_in_pointer( + &mut self, + destination_pointer: MemoryAddress, + value_variable: BrilligVariable, + ) { + match value_variable { + BrilligVariable::SingleAddr(value_variable) => { + self.store_instruction(destination_pointer, value_variable.address); + } + BrilligVariable::BrilligArray(_) => { + let reference: MemoryAddress = self.allocate_register(); + self.codegen_allocate_array_reference(reference); + self.codegen_store_variable(reference, value_variable); + self.store_instruction(destination_pointer, reference); + self.deallocate_register(reference); + } + BrilligVariable::BrilligVector(_) => { + let reference = self.allocate_register(); + self.codegen_allocate_vector_reference(reference); + self.codegen_store_variable(reference, value_variable); + self.store_instruction(destination_pointer, reference); + self.deallocate_register(reference); + } + } + } + /// Copies the values of an array pointed by source with length stored in `num_elements_register` /// Into the array pointed by destination pub(crate) fn codegen_copy_array(