diff --git a/crates/noirc_evaluator/src/brillig/brillig_gen.rs b/crates/noirc_evaluator/src/brillig/brillig_gen.rs index 39a012f2d3c..00d26f6c7fd 100644 --- a/crates/noirc_evaluator/src/brillig/brillig_gen.rs +++ b/crates/noirc_evaluator/src/brillig/brillig_gen.rs @@ -22,7 +22,7 @@ pub(crate) fn convert_ssa_function(func: &Function) -> BrilligArtifact { let mut function_context = FunctionContext { function_id: func.id(), ssa_value_to_register: HashMap::new() }; - let mut brillig_context = BrilligContext::default(); + let mut brillig_context = BrilligContext::new(); for block in reverse_post_order { BrilligBlock::compile(&mut function_context, &mut brillig_context, block, &func.dfg); @@ -33,7 +33,7 @@ pub(crate) fn convert_ssa_function(func: &Function) -> BrilligArtifact { /// Creates an entry point artifact, that will be linked with the brillig functions being called pub(crate) fn create_entry_point_function(num_arguments: usize) -> BrilligArtifact { - let mut brillig_context = BrilligContext::default(); + let mut brillig_context = BrilligContext::new(); brillig_context.entry_point_instruction(num_arguments); brillig_context.artifact() } diff --git a/crates/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/crates/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs index af25f2945ad..8d74b11332e 100644 --- a/crates/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs +++ b/crates/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs @@ -263,7 +263,7 @@ impl<'block> BrilligBlock<'block> { } Value::Array { array, element_type } => { // Allocate a register for the iterator - let iterator_register = self.brillig_context.create_register(); + let iterator_register = self.brillig_context.allocate_register(); // Set the iterator to the address of the array self.brillig_context.mov_instruction(iterator_register, address_register); @@ -368,7 +368,7 @@ impl<'block> BrilligBlock<'block> { register_index } Value::Array { .. } => { - let address_register = self.brillig_context.create_register(); + let address_register = self.brillig_context.allocate_register(); self.brillig_context.allocate_fixed_length_array( address_register, compute_size_of_type(&dfg.type_of_value(value_id)), diff --git a/crates/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs b/crates/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs index 76b660d1c5b..60de8743373 100644 --- a/crates/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs +++ b/crates/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs @@ -26,7 +26,7 @@ impl FunctionContext { return *register_index; } - let register = brillig_context.create_register(); + let register = brillig_context.allocate_register(); // Cache the `ValueId` so that if we call it again, it will // return the register that has just been created. diff --git a/crates/noirc_evaluator/src/brillig/brillig_ir.rs b/crates/noirc_evaluator/src/brillig/brillig_ir.rs index 9886981d49a..b4e30b701f5 100644 --- a/crates/noirc_evaluator/src/brillig/brillig_ir.rs +++ b/crates/noirc_evaluator/src/brillig/brillig_ir.rs @@ -5,8 +5,12 @@ //! ssa types and types in this module. //! A similar paradigm can be seen with the `acir_ir` module. pub(crate) mod artifact; +pub(crate) mod registers; -use self::artifact::{BrilligArtifact, UnresolvedJumpLocation}; +use self::{ + artifact::{BrilligArtifact, UnresolvedJumpLocation}, + registers::BrilligRegistersContext, +}; use acvm::{ acir::brillig_vm::{ BinaryFieldOp, BinaryIntOp, Opcode as BrilligOpcode, RegisterIndex, RegisterValueOrArray, @@ -50,11 +54,10 @@ impl ReservedRegisters { /// Brillig context object that is used while constructing the /// Brillig bytecode. -#[derive(Default)] pub(crate) struct BrilligContext { obj: BrilligArtifact, - /// A usize indicating the latest un-used register. - latest_register: usize, + /// Tracks register allocations + registers: BrilligRegistersContext, /// Context label, must be unique with respect to the function /// being linked. context_label: String, @@ -63,11 +66,21 @@ pub(crate) struct BrilligContext { } impl BrilligContext { + /// Initial context state + pub(crate) fn new() -> BrilligContext { + BrilligContext { + obj: BrilligArtifact::default(), + registers: BrilligRegistersContext::new(), + context_label: String::default(), + section_label: 0, + } + } + /// Adds the instructions needed to handle entry point parameters /// And sets the starting value of the reserved registers pub(crate) fn entry_point_instruction(&mut self, num_arguments: usize) { // Translate the inputs by the reserved registers offset - for i in (0..num_arguments).into_iter().rev() { + for i in (0..num_arguments).rev() { self.mov_instruction(self.user_register_index(i), RegisterIndex::from(i)); } // Set the initial value of the stack pointer register @@ -123,7 +136,7 @@ impl BrilligContext { result: RegisterIndex, ) { // Computes array_ptr + index, ie array[index] - let index_of_element_in_memory = self.create_register(); + let index_of_element_in_memory = self.allocate_register(); self.binary_instruction( array_ptr, index, @@ -132,6 +145,8 @@ impl BrilligContext { ); self.load_instruction(result, index_of_element_in_memory); + // Free up temporary register + self.deallocate_register(index_of_element_in_memory); } /// Sets the item in the array at index `index` to `value` @@ -142,7 +157,7 @@ impl BrilligContext { value: RegisterIndex, ) { // Computes array_ptr + index, ie array[index] - let index_of_element_in_memory = self.create_register(); + let index_of_element_in_memory = self.allocate_register(); self.binary_instruction( array_ptr, index, @@ -151,6 +166,8 @@ impl BrilligContext { ); self.store_instruction(index_of_element_in_memory, value); + // Free up temporary register + self.deallocate_register(index_of_element_in_memory); } /// Copies the values of an array pointed by source with length stored in `num_elements_register` @@ -161,7 +178,7 @@ impl BrilligContext { destination: RegisterIndex, num_elements_register: RegisterIndex, ) { - let index = self.make_constant(0_u128.into()); + let index_register = self.make_constant(0_u128.into()); let loop_label = self.next_section_label(); self.enter_next_section(); @@ -169,9 +186,9 @@ impl BrilligContext { // Loop body // Check if index < num_elements - let index_less_than_array_len = self.create_register(); + let index_less_than_array_len = self.allocate_register(); self.binary_instruction( - index, + index_register, num_elements_register, index_less_than_array_len, BrilligBinaryOp::Integer { @@ -186,16 +203,16 @@ impl BrilligContext { self.jump_if_instruction(index_less_than_array_len, exit_loop_label); // Copy the element from source to destination - let value_register = self.create_register(); - self.array_get(source, index, value_register); - self.array_set(destination, index, value_register); + let value_register = self.allocate_register(); + self.array_get(source, index_register, value_register); + self.array_set(destination, index_register, value_register); // Increment the index register - let one = self.make_constant(1u128.into()); + let one_register = self.make_constant(1u128.into()); self.binary_instruction( - index, - one, - index, + index_register, + one_register, + index_register, BrilligBinaryOp::Integer { op: BinaryIntOp::Add, bit_size: BRILLIG_MEMORY_ADDRESSING_BIT_SIZE, @@ -206,6 +223,11 @@ impl BrilligContext { // Exit the loop self.enter_next_section(); + // Deallocate our temporary registers + self.deallocate_register(index_less_than_array_len); + self.deallocate_register(value_register); + self.deallocate_register(one_register); + self.deallocate_register(index_register); } /// Adds a label to the next opcode @@ -272,11 +294,16 @@ impl BrilligContext { RegisterIndex::from(index + ReservedRegisters::len()) } - /// Creates a new register. - pub(crate) fn create_register(&mut self) -> RegisterIndex { - let register = self.user_register_index(self.latest_register); - self.latest_register += 1; - register + /// Allocates an unused register. + pub(crate) fn allocate_register(&mut self) -> RegisterIndex { + self.registers.allocate_register() + } + + /// Push a register to the deallocation list, ready for reuse. + /// TODO(AD): currently, register deallocation is only done with immediate values. + /// TODO(AD): See https://github.com/noir-lang/noir/issues/1720 + pub(crate) fn deallocate_register(&mut self, register_index: RegisterIndex) { + self.registers.deallocate_register(register_index); } } @@ -303,13 +330,8 @@ impl BrilligContext { /// the VM. pub(crate) fn return_instruction(&mut self, return_registers: &[RegisterIndex]) { for (destination_index, return_register) in return_registers.iter().enumerate() { - // If the destination register index is more than the latest register, - // we update the latest register to be the destination register because the - // brillig vm will expand the number of registers internally, when it encounters - // a register that has not been initialized. - if destination_index > self.latest_register { - self.latest_register = destination_index; - } + // In case we have fewer return registers than indices to write to, ensure we've allocated this register + self.registers.ensure_register_is_allocated(destination_index.into()); self.mov_instruction(destination_index.into(), *return_register); } self.stop_instruction(); @@ -434,7 +456,7 @@ impl BrilligContext { /// Returns a register which holds the value of a constant pub(crate) fn make_constant(&mut self, constant: Value) -> RegisterIndex { - let register = self.create_register(); + let register = self.allocate_register(); self.const_instruction(register, constant); register } @@ -456,8 +478,8 @@ impl BrilligContext { bit_size: u32, signed: bool, ) { - let scratch_register_i = self.create_register(); - let scratch_register_j = self.create_register(); + let scratch_register_i = self.allocate_register(); + let scratch_register_j = self.allocate_register(); // i = left / right self.push_opcode(BrilligOpcode::BinaryIntOp { @@ -488,6 +510,9 @@ impl BrilligContext { lhs: left, rhs: scratch_register_j, }); + // Free scratch registers + self.deallocate_register(scratch_register_i); + self.deallocate_register(scratch_register_j); } /// Emits a modulo instruction against 2**target_bit_size @@ -509,13 +534,14 @@ impl BrilligContext { // The brillig VM performs all arithmetic operations modulo 2**bit_size // So to cast any value to a target bit size we can just issue a no-op arithmetic operation // With bit size equal to target_bit_size - let zero = self.make_constant(Value::from(FieldElement::zero())); + let zero_register = self.make_constant(Value::from(FieldElement::zero())); self.binary_instruction( source, - zero, + zero_register, destination, BrilligBinaryOp::Integer { op: BinaryIntOp::Add, bit_size: target_bit_size }, ); + self.deallocate_register(zero_register); } } diff --git a/crates/noirc_evaluator/src/brillig/brillig_ir/registers.rs b/crates/noirc_evaluator/src/brillig/brillig_ir/registers.rs new file mode 100644 index 00000000000..cad5670d76b --- /dev/null +++ b/crates/noirc_evaluator/src/brillig/brillig_ir/registers.rs @@ -0,0 +1,59 @@ +use acvm::acir::brillig_vm::RegisterIndex; + +use super::ReservedRegisters; + +/// Every brillig stack frame/call context has its own view of register space. +/// This is maintained by copying these registers to the stack during calls and reading them back. +/// +/// Each has a stack base pointer from which all stack allocations can be offset. +pub(crate) struct BrilligRegistersContext { + /// A free-list of registers that have been deallocated and can be used again. + /// TODO(AD): currently, register deallocation is only done with immediate values. + /// TODO(AD): See https://github.com/noir-lang/noir/issues/1720 + deallocated_registers: Vec, + /// A usize indicating the next un-used register. + next_free_register_index: usize, +} + +impl BrilligRegistersContext { + /// Initial register allocation + pub(crate) fn new() -> BrilligRegistersContext { + BrilligRegistersContext { + deallocated_registers: Vec::new(), + next_free_register_index: ReservedRegisters::len(), + } + } + + /// Ensures a register is allocated. + pub(crate) fn ensure_register_is_allocated(&mut self, register: RegisterIndex) { + let index = register.to_usize(); + if index < self.next_free_register_index { + // If it could be allocated, check if it's in the deallocated list and remove it from there + self.deallocated_registers.retain(|&r| r != register); + } else { + // If it couldn't yet be, expand the register space. + self.next_free_register_index = index + 1; + } + } + + /// Creates a new register. + pub(crate) fn allocate_register(&mut self) -> RegisterIndex { + // If we have a register in our free list of deallocated registers, + // consume it first. This prioritizes reuse. + if let Some(register) = self.deallocated_registers.pop() { + return register; + } + // Otherwise, move to our latest register. + let register = RegisterIndex::from(self.next_free_register_index); + self.next_free_register_index += 1; + register + } + + /// Push a register to the deallocation list, ready for reuse. + /// TODO(AD): currently, register deallocation is only done with immediate values. + /// TODO(AD): See https://github.com/noir-lang/noir/issues/1720 + pub(crate) fn deallocate_register(&mut self, register_index: RegisterIndex) { + assert!(!self.deallocated_registers.contains(®ister_index)); + self.deallocated_registers.push(register_index); + } +}