added argument for initial rc

noir-lang · Nov 21, 2024 · 1e12188 · 1e12188
1 parent 204de4a
commit 1e12188
Show file tree

Hide file tree

Showing 8 changed files with 24 additions and 13 deletions.
diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
@@ -492,6 +492,7 @@ impl<'block> BrilligBlock<'block> {
                         destination_vector,
                         source_size_register,
                         None,
+                        None,
                     );
 
                     // Items
@@ -772,13 +773,14 @@ impl<'block> BrilligBlock<'block> {
                     // Initialize the variable
                     match new_variable {
                         BrilligVariable::BrilligArray(brillig_array) => {
-                            self.brillig_context.codegen_initialize_array(brillig_array);
+                            self.brillig_context.codegen_initialize_array(brillig_array, None);
                         }
                         BrilligVariable::BrilligVector(vector) => {
                             let size = self
                                 .brillig_context
                                 .make_usize_constant_instruction(array.len().into());
-                            self.brillig_context.codegen_initialize_vector(vector, size, None);
+                            self.brillig_context
+                                .codegen_initialize_vector(vector, size, None, None);
                             self.brillig_context.deallocate_single_addr(size);
                         }
                         _ => unreachable!(
@@ -1773,7 +1775,7 @@ impl<'block> BrilligBlock<'block> {
             unreachable!("ICE: allocate_foreign_call_array() expects an array, got {typ:?}")
         };
 
-        self.brillig_context.codegen_initialize_array(array);
+        self.brillig_context.codegen_initialize_array(array, None);
 
         let mut index = 0_usize;
         for _ in 0..*size {

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_intrinsic.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_intrinsic.rs
@@ -75,7 +75,7 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
         assert!(source_field.bit_size == F::max_num_bits());
         assert!(radix.bit_size == 32);
 
-        self.codegen_initialize_array(target_array);
+        self.codegen_initialize_array(target_array, None);
 
         let heap_array = self.codegen_brillig_array_to_heap_array(target_array);
 

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_memory.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_memory.rs
@@ -421,12 +421,16 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
     }
 
     /// Initializes an array, allocating memory to store its representation and initializing the reference counter.
-    pub(crate) fn codegen_initialize_array(&mut self, array: BrilligArray) {
+    pub(crate) fn codegen_initialize_array(
+        &mut self,
+        array: BrilligArray,
+        initial_rc: Option<usize>,
+    ) {
         self.codegen_allocate_immediate_mem(array.pointer, array.size + 1);
         self.indirect_const_instruction(
             array.pointer,
             BRILLIG_MEMORY_ADDRESSING_BIT_SIZE,
-            INITIAL_ARRAY_REF_COUNT.into(),
+            initial_rc.unwrap_or(INITIAL_ARRAY_REF_COUNT).into(),
         );
     }
 
@@ -435,12 +439,13 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
         vector: BrilligVector,
         size: SingleAddrVariable,
         capacity: Option<SingleAddrVariable>,
+        initial_rc: Option<usize>,
     ) {
         // Write RC
         self.indirect_const_instruction(
             vector.pointer,
             BRILLIG_MEMORY_ADDRESSING_BIT_SIZE,
-            INITIAL_ARRAY_REF_COUNT.into(),
+            initial_rc.unwrap_or(INITIAL_ARRAY_REF_COUNT).into(),
         );
 
         // Write size
@@ -461,6 +466,7 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
         vector: BrilligVector,
         size: SingleAddrVariable,
         capacity: Option<SingleAddrVariable>, // Defaults to size if None
+        initial_rc: Option<usize>,            // Defaults to INITIAL_ARRAY_REF_COUNT if none
     ) {
         let allocation_size = self.allocate_register();
         // Allocation size = capacity + 3 (rc, size, capacity)
@@ -473,7 +479,7 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
         self.codegen_allocate_mem(vector.pointer, allocation_size);
         self.deallocate_register(allocation_size);
 
-        self.codegen_initialize_vector_metadata(vector, size, capacity);
+        self.codegen_initialize_vector_metadata(vector, size, capacity, initial_rc);
     }
 
     /// We don't know the length of a vector returned externally before the call

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/entry_point.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/entry_point.rs
@@ -199,15 +199,15 @@ impl<F: AcirField + DebugToString> BrilligContext<F, Stack> {
         let deflattened_items_pointer = if is_vector {
             let vector = BrilligVector { pointer: deflattened_array_pointer };
 
-            self.codegen_initialize_vector(vector, deflattened_size_variable, None);
+            self.codegen_initialize_vector(vector, deflattened_size_variable, None, Some(1));
 
             self.codegen_make_vector_items_pointer(vector)
         } else {
             let arr = BrilligArray {
                 pointer: deflattened_array_pointer,
                 size: item_count * item_type.len(),
             };
-            self.codegen_initialize_array(arr);
+            self.codegen_initialize_array(arr, Some(1));
             self.codegen_make_array_items_pointer(arr)
         };
 

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/prepare_vector_push.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/prepare_vector_push.rs
@@ -185,6 +185,7 @@ pub(crate) fn reallocate_vector_for_insertion<
                             target_vector,
                             target_size,
                             Some(source_capacity),
+                            Some(1),
                         );
                     }
                 });
@@ -201,6 +202,7 @@ pub(crate) fn reallocate_vector_for_insertion<
                     target_vector,
                     target_size,
                     Some(double_size),
+                    Some(1),
                 );
                 brillig_context.deallocate_single_addr(double_size);
             }

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/vector_pop_back.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/vector_pop_back.rs
@@ -78,7 +78,7 @@ pub(super) fn compile_vector_pop_back_procedure<F: AcirField + DebugToString>(
             brillig_context.codegen_update_vector_length(target_vector, target_size);
         } else {
             // We need to clone the source vector
-            brillig_context.codegen_initialize_vector(target_vector, target_size, None);
+            brillig_context.codegen_initialize_vector(target_vector, target_size, None, Some(1));
 
             let target_vector_items_pointer =
                 brillig_context.codegen_make_vector_items_pointer(target_vector);

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/vector_pop_front.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/vector_pop_front.rs
@@ -95,9 +95,10 @@ pub(super) fn compile_vector_pop_front_procedure<F: AcirField + DebugToString>(
                 target_vector,
                 target_size,
                 Some(source_capacity),
+                Some(1),
             );
         } else {
-            brillig_context.codegen_initialize_vector(target_vector, target_size, None);
+            brillig_context.codegen_initialize_vector(target_vector, target_size, None, Some(1));
 
             let target_vector_items_pointer =
                 brillig_context.codegen_make_vector_items_pointer(target_vector);

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/vector_remove.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/vector_remove.rs
@@ -81,7 +81,7 @@ pub(super) fn compile_vector_remove_procedure<F: AcirField + DebugToString>(
             brillig_context
                 .codegen_vector_items_pointer(target_vector, target_vector_items_pointer);
         } else {
-            brillig_context.codegen_initialize_vector(target_vector, target_size, None);
+            brillig_context.codegen_initialize_vector(target_vector, target_size, None, Some(1));
 
             // Copy the elements to the left of the index
             brillig_context