From 4786ddcca256d8786e3ba4c52d76be36375ddf06 Mon Sep 17 00:00:00 2001
From: sirasistant <sirasistant@gmail.com>
Date: Tue, 3 Sep 2024 12:58:02 +0000
Subject: [PATCH] fix: temporary register leaks

---
 .../brillig/brillig_gen/brillig_black_box.rs  | 67 ++++++++++++++----
 .../src/brillig/brillig_gen/brillig_block.rs  | 69 ++++++++++---------
 2 files changed, 91 insertions(+), 45 deletions(-)

diff --git a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_black_box.rs b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_black_box.rs
index 802d442885f..bd9190c1cfe 100644
--- a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_black_box.rs
+++ b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_black_box.rs
@@ -29,6 +29,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     message: message_vector.to_heap_vector(),
                     output: result_array.to_heap_array(),
                 });
+                deallocate_converted_vector(brillig_context, message, message_vector, bb_func);
             } else {
                 unreachable!("ICE: SHA256 expects one array argument and one array result")
             }
@@ -42,6 +43,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     message: message_vector.to_heap_vector(),
                     output: result_array.to_heap_array(),
                 });
+                deallocate_converted_vector(brillig_context, message, message_vector, bb_func);
             } else {
                 unreachable!("ICE: Blake2s expects one array argument and one array result")
             }
@@ -55,6 +57,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     message: message_vector.to_heap_vector(),
                     output: result_array.to_heap_array(),
                 });
+                deallocate_converted_vector(brillig_context, message, message_vector, bb_func);
             } else {
                 unreachable!("ICE: Blake3 expects one array argument and one array result")
             }
@@ -78,6 +81,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     output: result_array.to_heap_array(),
                 });
                 brillig_context.deallocate_single_addr(message_size_as_usize);
+                deallocate_converted_vector(brillig_context, message, message_vector, bb_func);
             } else {
                 unreachable!("ICE: Keccak256 expects message, message size and result array")
             }
@@ -92,6 +96,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     message: state_vector.to_heap_vector(),
                     output: result_array.to_heap_array(),
                 });
+                deallocate_converted_vector(brillig_context, message, state_vector, bb_func);
             } else {
                 unreachable!("ICE: Keccakf1600 expects one array argument and one array result")
             }
@@ -111,6 +116,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     signature: signature.to_heap_array(),
                     result: result_register.address,
                 });
+                deallocate_converted_vector(brillig_context, message, message_hash_vector, bb_func);
             } else {
                 unreachable!(
                     "ICE: EcdsaSecp256k1 expects four array arguments and one register result"
@@ -132,6 +138,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     signature: signature.to_heap_array(),
                     result: result_register.address,
                 });
+                deallocate_converted_vector(brillig_context, message, message_hash_vector, bb_func);
             } else {
                 unreachable!(
                     "ICE: EcdsaSecp256r1 expects four array arguments and one register result"
@@ -151,6 +158,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     domain_separator: domain_separator.address,
                     output: result_array.to_heap_array(),
                 });
+                deallocate_converted_vector(brillig_context, message, message_vector, bb_func);
             } else {
                 unreachable!("ICE: Pedersen expects one array argument, a register for the domain separator, and one array result")
             }
@@ -167,6 +175,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     domain_separator: domain_separator.address,
                     output: result.address,
                 });
+                deallocate_converted_vector(brillig_context, message, message_vector, bb_func);
             } else {
                 unreachable!("ICE: Pedersen hash expects one array argument, a register for the domain separator, and one register result")
             }
@@ -178,14 +187,16 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
             ) = (function_arguments, function_results)
             {
                 let message_hash = convert_array_or_vector(brillig_context, message, bb_func);
-                let signature = brillig_context.array_to_vector_instruction(signature);
+                let signature_vector = brillig_context.array_to_vector_instruction(signature);
                 brillig_context.black_box_op_instruction(BlackBoxOp::SchnorrVerify {
                     public_key_x: public_key_x.address,
                     public_key_y: public_key_y.address,
                     message: message_hash.to_heap_vector(),
-                    signature: signature.to_heap_vector(),
+                    signature: signature_vector.to_heap_vector(),
                     result: result_register.address,
                 });
+                deallocate_converted_vector(brillig_context, message, message_hash, bb_func);
+                brillig_context.deallocate_register(signature_vector.size);
             } else {
                 unreachable!("ICE: Schnorr verify expects two registers for the public key, an array for signature, an array for the message hash and one result register")
             }
@@ -194,13 +205,15 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
             if let ([points, scalars], [BrilligVariable::BrilligArray(outputs)]) =
                 (function_arguments, function_results)
             {
-                let points = convert_array_or_vector(brillig_context, points, bb_func);
-                let scalars = convert_array_or_vector(brillig_context, scalars, bb_func);
+                let points_vector = convert_array_or_vector(brillig_context, points, bb_func);
+                let scalars_vector = convert_array_or_vector(brillig_context, scalars, bb_func);
                 brillig_context.black_box_op_instruction(BlackBoxOp::MultiScalarMul {
-                    points: points.to_heap_vector(),
-                    scalars: scalars.to_heap_vector(),
+                    points: points_vector.to_heap_vector(),
+                    scalars: scalars_vector.to_heap_vector(),
                     outputs: outputs.to_heap_array(),
                 });
+                deallocate_converted_vector(brillig_context, points, points_vector, bb_func);
+                deallocate_converted_vector(brillig_context, scalars, scalars_vector, bb_func);
             } else {
                 unreachable!(
                     "ICE: MultiScalarMul expects two register arguments and one array result"
@@ -319,6 +332,8 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     modulus: modulus_vector.to_heap_vector(),
                     output: output.address,
                 });
+                deallocate_converted_vector(brillig_context, inputs, inputs_vector, bb_func);
+                deallocate_converted_vector(brillig_context, modulus, modulus_vector, bb_func);
             } else {
                 unreachable!(
                     "ICE: BigIntFromLeBytes expects a register and an array  as arguments and two result registers"
@@ -336,6 +351,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     input: input.address,
                     output: output.to_heap_vector(),
                 });
+                deallocate_converted_vector(brillig_context, result_array, output, bb_func);
             } else {
                 unreachable!(
                     "ICE: BigIntToLeBytes expects two register arguments and one array result"
@@ -354,6 +370,7 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     output: result_array.to_heap_array(),
                     len: state_len.address,
                 });
+                deallocate_converted_vector(brillig_context, message, message_vector, bb_func);
             } else {
                 unreachable!("ICE: Poseidon2Permutation expects one array argument, a length and one array result")
             }
@@ -369,6 +386,8 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                     hash_values: hash_vector.to_heap_vector(),
                     output: result_array.to_heap_array(),
                 });
+                deallocate_converted_vector(brillig_context, message, message_vector, bb_func);
+                deallocate_converted_vector(brillig_context, hash_values, hash_vector, bb_func);
             } else {
                 unreachable!("ICE: Sha256Compression expects two array argument, one array result")
             }
@@ -379,18 +398,19 @@ pub(crate) fn convert_black_box_call<F: AcirField + DebugToString, Registers: Re
                 [BrilligVariable::SingleAddr(out_len), outputs],
             ) = (function_arguments, function_results)
             {
-                let inputs = convert_array_or_vector(brillig_context, inputs, bb_func);
-                let outputs = convert_array_or_vector(brillig_context, outputs, bb_func);
-                let output_vec = outputs.to_heap_vector();
+                let inputs_vector = convert_array_or_vector(brillig_context, inputs, bb_func);
+                let outputs_vector = convert_array_or_vector(brillig_context, outputs, bb_func);
                 brillig_context.black_box_op_instruction(BlackBoxOp::AES128Encrypt {
-                    inputs: inputs.to_heap_vector(),
+                    inputs: inputs_vector.to_heap_vector(),
                     iv: iv.to_heap_array(),
                     key: key.to_heap_array(),
-                    outputs: output_vec,
+                    outputs: outputs_vector.to_heap_vector(),
                 });
-                brillig_context.mov_instruction(out_len.address, output_vec.size);
+                brillig_context.mov_instruction(out_len.address, outputs_vector.size);
                 // Returns slice, so we need to allocate memory for it after the fact
-                brillig_context.increase_free_memory_pointer_instruction(output_vec.size);
+                brillig_context.increase_free_memory_pointer_instruction(outputs_vector.size);
+                deallocate_converted_vector(brillig_context, inputs, inputs_vector, bb_func);
+                deallocate_converted_vector(brillig_context, outputs, outputs_vector, bb_func);
             } else {
                 unreachable!("ICE: AES128Encrypt expects three array arguments, one array result")
             }
@@ -413,3 +433,24 @@ fn convert_array_or_vector<F: AcirField + DebugToString, Registers: RegisterAllo
         ),
     }
 }
+
+/// Deallocates any new register allocated by the function above.
+/// Concretely, the only allocated register between array and vector is the size register if the array was converted to a vector.
+fn deallocate_converted_vector<F: AcirField + DebugToString, Registers: RegisterAllocator>(
+    brillig_context: &mut BrilligContext<F, Registers>,
+    original_array_or_vector: &BrilligVariable,
+    converted_vector: BrilligVector,
+    bb_func: &BlackBoxFunc,
+) {
+    match original_array_or_vector {
+        BrilligVariable::BrilligArray(_) => {
+            brillig_context.deallocate_register(converted_vector.size);
+        }
+        BrilligVariable::BrilligVector(_) => {}
+        _ => unreachable!(
+            "ICE: {} expected an array or a vector, but got {:?}",
+            bb_func.name(),
+            original_array_or_vector
+        ),
+    }
+}
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 a7cb1571e34..4f678303542 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
@@ -810,6 +810,12 @@ impl<'block> BrilligBlock<'block> {
         // puts the returns into the returned_registers and restores saved_registers
         self.brillig_context
             .codegen_post_call_prep_returns_load_registers(&returned_registers, &saved_registers);
+
+        // Reset the register state to the one needed to hold the current available variables
+        let variables = self.variables.get_available_variables(self.function_context);
+        let registers =
+            variables.into_iter().flat_map(|variable| variable.extract_registers()).collect();
+        self.brillig_context.set_allocated_registers(registers);
     }
 
     fn validate_array_index(
@@ -1794,7 +1800,7 @@ impl<'block> BrilligBlock<'block> {
                     dfg,
                 );
                 let array = variable.extract_array();
-                self.allocate_nested_array(typ, Some(array));
+                self.allocate_foreign_call_result_array(typ, array);
 
                 variable
             }
@@ -1821,40 +1827,39 @@ impl<'block> BrilligBlock<'block> {
         }
     }
 
-    fn allocate_nested_array(
-        &mut self,
-        typ: &Type,
-        array: Option<BrilligArray>,
-    ) -> BrilligVariable {
-        match typ {
-            Type::Array(types, size) => {
-                let array = array.unwrap_or(BrilligArray {
-                    pointer: self.brillig_context.allocate_register(),
-                    size: *size,
-                    rc: self.brillig_context.allocate_register(),
-                });
-                self.brillig_context.codegen_allocate_fixed_length_array(array.pointer, array.size);
-                self.brillig_context.usize_const_instruction(array.rc, 1_usize.into());
-
-                let mut index = 0_usize;
-                for _ in 0..*size {
-                    for element_type in types.iter() {
-                        match element_type {
-                            Type::Array(_, _) => {
-                                let inner_array = self.allocate_nested_array(element_type, None);
-                                let idx =
-                                    self.brillig_context.make_usize_constant_instruction(index.into());
-                                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"),
-                            _ => (),
-                        }
-                        index += 1;
+    fn allocate_foreign_call_result_array(&mut self, typ: &Type, array: BrilligArray) {
+        let Type::Array(types, size) = typ else {
+            unreachable!("ICE: allocate_foreign_call_array() expects an array, got {typ:?}")
+        };
+
+        self.brillig_context.codegen_allocate_fixed_length_array(array.pointer, array.size);
+        self.brillig_context.usize_const_instruction(array.rc, 1_usize.into());
+
+        let mut index = 0_usize;
+        for _ in 0..*size {
+            for element_type in types.iter() {
+                match element_type {
+                    Type::Array(_, nested_size) => {
+                        let inner_array = BrilligArray {
+                            pointer: self.brillig_context.allocate_register(),
+                            rc: self.brillig_context.allocate_register(),
+                            size: *nested_size,
+                        };
+                        self.allocate_foreign_call_result_array(element_type, inner_array);
+
+                        let idx =
+                            self.brillig_context.make_usize_constant_instruction(index.into());
+                        self.brillig_context.codegen_store_variable_in_array(array.pointer, idx, BrilligVariable::BrilligArray(inner_array));
+
+                        self.brillig_context.deallocate_single_addr(idx);
+                        self.brillig_context.deallocate_register(inner_array.pointer);
+                        self.brillig_context.deallocate_register(inner_array.rc);
                     }
+                    Type::Slice(_) => unreachable!("ICE: unsupported slice type in allocate_nested_array(), expects an array or a numeric type"),
+                    _ => (),
                 }
-                BrilligVariable::BrilligArray(array)
+                index += 1;
             }
-            _ => unreachable!("ICE: allocate_nested_array() expects an array, got {typ:?}"),
         }
     }