chore(brillig): Add a register abstraction/optimization (#1737)

* chore: brillig register abstraction

* chore: small refactor

* fix: latest_register starting value

* fix: comments

* chore: remove unused method

* fix: reintroduce a safety register alloc check

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()
 }

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)),

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.

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,17 +178,17 @@ 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();
 
         // 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);
     }
 }
 

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<RegisterIndex>,
+    /// 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(&register_index));
+        self.deallocated_registers.push(register_index);
+    }
+}