fix!: Only decrement the counter of an array if its address has not changed

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs

@@ -20,6 +20,7 @@ use crate::ssa::ir::{
     value::{Value, ValueId},
 };
 use acvm::acir::brillig::{MemoryAddress, ValueOrArray};
+use acvm::brillig_vm::MEMORY_ADDRESSING_BIT_SIZE;
 use acvm::{acir::AcirField, FieldElement};
 use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
 use iter_extended::vecmap;
@@ -835,20 +836,65 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
                     .store_instruction(array_or_vector.extract_register(), rc_register);
                 self.brillig_context.deallocate_register(rc_register);
             }
-            Instruction::DecrementRc { value } => {
+            Instruction::DecrementRc { value, original } => {
                 let array_or_vector = self.convert_ssa_value(*value, dfg);
-                let rc_register = self.brillig_context.allocate_register();
+                let orig_array_or_vector = self.convert_ssa_value(*original, dfg);
 
-                self.brillig_context
-                    .load_instruction(rc_register, array_or_vector.extract_register());
-                self.brillig_context.codegen_usize_op_in_place(
-                    rc_register,
-                    BrilligBinaryOp::Sub,
-                    1,
-                );
-                self.brillig_context
-                    .store_instruction(array_or_vector.extract_register(), rc_register);
-                self.brillig_context.deallocate_register(rc_register);
+                let array_register = array_or_vector.extract_register();
+                let orig_array_register = orig_array_or_vector.extract_register();
+
+                let codegen_dec_rc = |ctx: &mut BrilligContext<FieldElement, Registers>| {
+                    let rc_register = ctx.allocate_register();
+
+                    ctx.load_instruction(rc_register, array_register);
+
+                    ctx.codegen_usize_op_in_place(rc_register, BrilligBinaryOp::Sub, 1);
+
+                    // Check that the refcount didn't go below 1. If we allow it to underflow
+                    // and become 0 or usize::MAX, and then return to 1, then it will indicate
+                    // an array as mutable when it probably shouldn't be.
+                    {
+                        let min_addr = ReservedRegisters::usize_one();
+                        let condition = SingleAddrVariable::new(ctx.allocate_register(), 1);
+                        ctx.memory_op_instruction(
+                            min_addr,
+                            rc_register,
+                            condition.address,
+                            BrilligBinaryOp::LessThanEquals,
+                        );
+                        ctx.codegen_constrain(
+                            condition,
+                            Some("array ref-count went to zero".to_owned()),
+                        );
+                        ctx.deallocate_single_addr(condition);
+                    }
+
+                    ctx.store_instruction(array_register, rc_register);
+                    ctx.deallocate_register(rc_register);
+                };
+
+                if array_register == orig_array_register {
+                    codegen_dec_rc(&mut self.brillig_context);
+                } else {
+                    // Check if the current and original register point at the same memory.
+                    // If they don't, it means that an array-copy procedure has created a
+                    // new array, which includes decrementing the RC of the original,
+                    // which means we don't have to do the decrement here.
+                    let condition =
+                        SingleAddrVariable::new(self.brillig_context.allocate_register(), 1);
+                    // We will use a binary op to interpret the pointer as a number.
+                    let bit_size: u32 = MEMORY_ADDRESSING_BIT_SIZE.into();
+                    let array_var = SingleAddrVariable::new(array_register, bit_size);
+                    let orig_array_var = SingleAddrVariable::new(orig_array_register, bit_size);
+                    self.brillig_context.binary_instruction(
+                        array_var,
+                        orig_array_var,
+                        condition,
+                        BrilligBinaryOp::Equals,
+                    );
+                    self.brillig_context.codegen_if(condition.address, codegen_dec_rc);
+                    self.brillig_context.deallocate_single_addr(condition);
+                }
             }
             Instruction::EnableSideEffectsIf { .. } => {
                 unreachable!("enable_side_effects not supported by brillig")

compiler/noirc_evaluator/src/ssa/function_builder/mod.rs

@@ -390,8 +390,8 @@ impl FunctionBuilder {
 
     /// Insert an instruction to decrement an array's reference count. This only has an effect
     /// in unconstrained code where arrays are reference counted and copy on write.
-    pub(crate) fn insert_dec_rc(&mut self, value: ValueId) {
-        self.insert_instruction(Instruction::DecrementRc { value }, None);
+    pub(crate) fn insert_dec_rc(&mut self, value: ValueId, original: ValueId) {
+        self.insert_instruction(Instruction::DecrementRc { value, original }, None);
     }
 
     /// Insert an enable_side_effects_if instruction. These are normally only automatically
@@ -491,53 +491,63 @@ impl FunctionBuilder {
     /// within the given value. If the given value is not an array and does not contain
     /// any arrays, this does nothing.
     ///
-    /// Returns whether a reference count instruction was issued.
-    pub(crate) fn increment_array_reference_count(&mut self, value: ValueId) -> bool {
-        self.update_array_reference_count(value, true)
+    /// Returns the ID of the array that was affected, if any.
+    pub(crate) fn increment_array_reference_count(&mut self, value: ValueId) -> Option<ValueId> {
+        self.update_array_reference_count(value, None)
     }
 
     /// Insert instructions to decrement the reference count of any array(s) stored
     /// within the given value. If the given value is not an array and does not contain
     /// any arrays, this does nothing.
     ///
-    /// Returns whether a reference count instruction was issued.
-    pub(crate) fn decrement_array_reference_count(&mut self, value: ValueId) -> bool {
-        self.update_array_reference_count(value, false)
+    /// Returns the ID of the array that was affected, if any.
+    pub(crate) fn decrement_array_reference_count(
+        &mut self,
+        value: ValueId,
+        original: ValueId,
+    ) -> Option<ValueId> {
+        self.update_array_reference_count(value, Some(original))
     }
 
     /// Increment or decrement the given value's reference count if it is an array.
     /// If it is not an array, this does nothing. Note that inc_rc and dec_rc instructions
     /// are ignored outside of unconstrained code.
     ///
-    /// Returns whether a reference count instruction was issued.
-    fn update_array_reference_count(&mut self, value: ValueId, increment: bool) -> bool {
-        if self.current_function.runtime().is_brillig() {
-            match self.type_of_value(value) {
-                Type::Numeric(_) => false,
-                Type::Function => false,
-                Type::Reference(element) => {
-                    if element.contains_an_array() {
-                        let reference = value;
-                        let value = self.insert_load(reference, element.as_ref().clone());
-                        self.update_array_reference_count(value, increment);
-                        true
-                    } else {
-                        false
-                    }
+    /// If there is an `original` value it indicates that we're now decrementing it,
+    /// but that should only happen if it hasn't been changed by an `array_set` in
+    /// the meantime, which in itself would make sure its RC is decremented.
+    ///
+    /// Returns the ID of the array that was affected, if any.
+    fn update_array_reference_count(
+        &mut self,
+        value: ValueId,
+        original: Option<ValueId>,
+    ) -> Option<ValueId> {
+        if self.current_function.runtime().is_acir() {
+            return None;
+        }
+        match self.type_of_value(value) {
+            Type::Numeric(_) | Type::Function => None,
+            Type::Reference(element) => {
+                if element.contains_an_array() {
+                    let reference = value;
+                    let value = self.insert_load(reference, element.as_ref().clone());
+                    self.update_array_reference_count(value, original);
+                    Some(value)
+                } else {
+                    None
                 }
-                Type::Array(..) | Type::Slice(..) => {
-                    // If there are nested arrays or slices, we wait until ArrayGet
-                    // is issued to increment the count of that array.
-                    if increment {
-                        self.insert_inc_rc(value);
-                    } else {
-                        self.insert_dec_rc(value);
-                    }
-                    true
+            }
+            Type::Array(..) | Type::Slice(..) => {
+                // If there are nested arrays or slices, we wait until ArrayGet
+                // is issued to increment the count of that array.
+                if let Some(original) = original {
+                    self.insert_dec_rc(value, original);
+                } else {
+                    self.insert_inc_rc(value);
                 }
+                Some(value)
             }
-        } else {
-            false
         }
     }
 

compiler/noirc_evaluator/src/ssa/ir/instruction.rs

@@ -324,7 +324,9 @@
     /// This currently only has an effect in Brillig code where array sharing and copy on write is
     /// implemented via reference counting. In ACIR code this is done with im::Vector and these
     /// DecrementRc instructions are ignored.
-    DecrementRc { value: ValueId },
+    ///
+    /// The `original` contains the value of the array which was incremented by the pair of this decrement.
+    DecrementRc { value: ValueId, original: ValueId },
 
     /// Merge two values returned from opposite branches of a conditional into one.
     ///
@@ -423,7 +425,7 @@
             // These can fail.
             Constrain(..) | ConstrainNotEqual(..) | RangeCheck { .. } => true,
 
             // This should never be side-effectful
             MakeArray { .. } | Noop => false,
 
             // Some binary math can overflow or underflow
@@ -717,7 +719,9 @@
                 mutable: *mutable,
             },
             Instruction::IncrementRc { value } => Instruction::IncrementRc { value: f(*value) },
-            Instruction::DecrementRc { value } => Instruction::DecrementRc { value: f(*value) },
+            Instruction::DecrementRc { value, original } => {
+                Instruction::DecrementRc { value: f(*value), original: f(*original) }
+            }
             Instruction::RangeCheck { value, max_bit_size, assert_message } => {
                 Instruction::RangeCheck {
                     value: f(*value),
@@ -788,7 +792,10 @@
                 *value = f(*value);
             }
             Instruction::IncrementRc { value } => *value = f(*value),
-            Instruction::DecrementRc { value } => *value = f(*value),
+            Instruction::DecrementRc { value, original } => {
+                *value = f(*value);
+                *original = f(*original);
+            }
             Instruction::RangeCheck { value, max_bit_size: _, assert_message: _ } => {
                 *value = f(*value);
             }
@@ -854,10 +861,12 @@
             Instruction::EnableSideEffectsIf { condition } => {
                 f(*condition);
             }
-            Instruction::IncrementRc { value }
-            | Instruction::DecrementRc { value }
-            | Instruction::RangeCheck { value, .. } => {
+            Instruction::IncrementRc { value } | Instruction::RangeCheck { value, .. } => {
+                f(*value);
+            }
+            Instruction::DecrementRc { value, original } => {
                 f(*value);
+                f(*original);
             }
             Instruction::IfElse { then_condition, then_value, else_condition, else_value } => {
                 f(*then_condition);

compiler/noirc_evaluator/src/ssa/ir/printer.rs

@@ -250,8 +250,8 @@ fn display_instruction_inner(
         Instruction::IncrementRc { value } => {
             writeln!(f, "inc_rc {}", show(*value))
         }
-        Instruction::DecrementRc { value } => {
-            writeln!(f, "dec_rc {}", show(*value))
+        Instruction::DecrementRc { value, original } => {
+            writeln!(f, "dec_rc {} {}", show(*value), show(*original))
         }
         Instruction::RangeCheck { value, max_bit_size, .. } => {
             writeln!(f, "range_check {} to {} bits", show(*value), *max_bit_size,)

compiler/noirc_evaluator/src/ssa/opt/die.rs

@@ -284,7 +284,7 @@ impl Context {
             self.rc_instructions.iter().fold(HashMap::default(), |mut acc, (rc, block)| {
                 let value = match &dfg[*rc] {
                     Instruction::IncrementRc { value } => *value,
-                    Instruction::DecrementRc { value } => *value,
+                    Instruction::DecrementRc { value, .. } => *value,
                     other => {
                         unreachable!(
                             "Expected IncrementRc or DecrementRc instruction, found {other:?}"
@@ -434,7 +434,7 @@ impl Context {
         dfg: &DataFlowGraph,
     ) -> bool {
         use Instruction::*;
-        if let IncrementRc { value } | DecrementRc { value } = instruction {
+        if let IncrementRc { value } | DecrementRc { value, .. } = instruction {
             let Some(instruction) = dfg.get_local_or_global_instruction(*value) else {
                 return false;
             };
@@ -685,7 +685,7 @@ impl<'a> RcTracker<'a> {
                 // Remember that this array was RC'd by this instruction.
                 self.inc_rcs.entry(*value).or_default().insert(instruction_id);
             }
-            Instruction::DecrementRc { value } => {
+            Instruction::DecrementRc { value, .. } => {
                 let typ = function.dfg.type_of_value(*value);
 
                 // We assume arrays aren't mutated until we find an array_set

compiler/noirc_evaluator/src/ssa/opt/rc.rs

@@ -113,7 +113,7 @@ impl Context {
         let mut to_remove = HashSet::default();
 
         for instruction in function.dfg[last_block].instructions() {
-            if let Instruction::DecrementRc { value } = &function.dfg[*instruction] {
+            if let Instruction::DecrementRc { value, .. } = &function.dfg[*instruction] {
                 if let Some(inc_rc) = pop_rc_for(*value, function, &mut self.inc_rcs) {
                     if !inc_rc.possibly_mutated {
                         to_remove.insert(inc_rc.id);
@@ -213,7 +213,7 @@ mod test {
 
         builder.insert_inc_rc(v0);
         builder.insert_inc_rc(v0);
-        builder.insert_dec_rc(v0);
+        builder.insert_dec_rc(v0, v0);
 
         let outer_array_type = Type::Array(Arc::new(vec![inner_array_type]), 1);
         let v1 = builder.insert_make_array(vec![v0].into(), outer_array_type);
@@ -261,7 +261,7 @@ mod test {
         let v7 = builder.insert_array_set(v2, zero, five);
 
         builder.insert_store(v1, v7);
-        builder.insert_dec_rc(v0);
+        builder.insert_dec_rc(v0, v0);
         builder.terminate_with_return(vec![]);
 
         let ssa = builder.finish().remove_paired_rc();
@@ -314,7 +314,7 @@ mod test {
 
         builder.insert_store(v0, v7);
         let v8 = builder.insert_load(v0, array_type);
-        builder.insert_dec_rc(v8);
+        builder.insert_dec_rc(v8, v1);
         builder.insert_store(v0, v8);
         builder.terminate_with_return(vec![]);
 

compiler/noirc_evaluator/src/ssa/ssa_gen/context.rs

@@ -964,11 +964,13 @@ impl<'a> FunctionContext<'a> {
     ///
     /// This is done on parameters rather than call arguments so that we can optimize out
     /// paired inc/dec instructions within brillig functions more easily.
-    pub(crate) fn increment_parameter_rcs(&mut self) -> HashSet<ValueId> {
+    ///
+    /// Returns the list of parameters incremented, together with the value ID of the arrays they refer to.
+    pub(crate) fn increment_parameter_rcs(&mut self) -> Vec<(ValueId, ValueId)> {
         let entry = self.builder.current_function.entry_block();
         let parameters = self.builder.current_function.dfg.block_parameters(entry).to_vec();
 
-        let mut incremented = HashSet::default();
+        let mut incremented = Vec::default();
         let mut seen_array_types = HashSet::default();
 
         for parameter in parameters {
@@ -979,10 +981,11 @@ impl<'a> FunctionContext<'a> {
                 if element.contains_an_array() {
                     // If we haven't already seen this array type, the value may be possibly
                     // aliased, so issue an inc_rc for it.
-                    if !seen_array_types.insert(element.get_contained_array().clone())
-                        && self.builder.increment_array_reference_count(parameter)
-                    {
-                        incremented.insert(parameter);
+                    if seen_array_types.insert(element.get_contained_array().clone()) {
+                        continue;
+                    }
+                    if let Some(id) = self.builder.increment_array_reference_count(parameter) {
+                        incremented.push((parameter, id));
                     }
                 }
             }
@@ -997,14 +1000,14 @@ impl<'a> FunctionContext<'a> {
     /// ignored.
     pub(crate) fn end_scope(
         &mut self,
-        mut incremented_params: HashSet<ValueId>,
+        mut incremented_params: Vec<(ValueId, ValueId)>,
         terminator_args: &[ValueId],
     ) {
-        incremented_params.retain(|parameter| !terminator_args.contains(parameter));
+        incremented_params.retain(|(parameter, _)| !terminator_args.contains(parameter));
 
-        for parameter in incremented_params {
+        for (parameter, original) in incremented_params {
             if self.builder.current_function.dfg.value_is_reference(parameter) {
-                self.builder.decrement_array_reference_count(parameter);
+                self.builder.decrement_array_reference_count(parameter, original);
             }
         }
     }