diff --git a/noir-projects/aztec-nr/aztec/src/capsules/mod.nr b/noir-projects/aztec-nr/aztec/src/capsules/mod.nr index 5b0a8bb95929..c3d2ab7202de 100644 --- a/noir-projects/aztec-nr/aztec/src/capsules/mod.nr +++ b/noir-projects/aztec-nr/aztec/src/capsules/mod.nr @@ -76,6 +76,42 @@ impl CapsuleArray { capsules::store(self.contract_address, self.base_slot, current_length - 1); } + /// Iterates over the entire array, calling the callback with all values and their array index. The order in which + /// values are processed is arbitrary. + /// + /// It is safe to delete the current element (and only the current element) from inside the callback via `remove`: + /// ```noir + /// array.for_each(|index, value| { + /// if some_condition(value) { + /// array.remove(index); // safe only for this index + /// } + /// } + /// ``` + /// + /// If all elements in the array need to iterated over and then removed, then using `for_each` results in optimal + /// efficiency. + /// + /// It is **not** safe to push new elements into the array from inside the callback. + pub unconstrained fn for_each(self, f: unconstrained fn[Env](u32, T) -> ()) + where + T: Deserialize, + { + // Iterating over all elements is simple, but we want to do it in such a way that a) deleting the current + // element is safe to do, and b) deleting *all* elements is optimally efficient. This is because CapsuleArrays + // are typically used to hold pending tasks, so iterating them while clearing completed tasks (sometimes + // unconditionally, resulting in a full clear) is a very common access pattern. + // + // The way we achieve this is by iterating backwards: each element can always be deleted since it won't change + // any preceding (lower) indices, and if every element is deleted then every element will (in turn) be the last + // element. This results in an optimal full clear since `remove` will be able to skip the `capsules::copy` call + // to shift any elements past the deleted one (because there will be none). + let mut i = self.len(); + while i > 0 { + i -= 1; + f(i, self.get(i)); + } + } + unconstrained fn slot_at(self, index: u32) -> Field { // Elements are stored immediately after the base slot, so we add 1 to it to compute the slot for the first // element. @@ -181,4 +217,78 @@ mod test { assert_eq(array.len(), 0); } + + #[test] + unconstrained fn for_each_called_with_all_elements() { + let contract_address = setup(); + let array = CapsuleArray::at(contract_address, SLOT); + + array.push(4); + array.push(5); + array.push(6); + + // We store all values that we were called with and check that all (value, index) tuples are present. Note that + // we do not care about the order in which each tuple was passed to the closure. + let called_with = &mut BoundedVec::<(u32, Field), 3>::new(); + array.for_each(|index, value| { called_with.push((index, value)); }); + + assert_eq(called_with.len(), 3); + assert(called_with.any(|(index, value)| (index == 0) & (value == 4))); + assert(called_with.any(|(index, value)| (index == 1) & (value == 5))); + assert(called_with.any(|(index, value)| (index == 2) & (value == 6))); + } + + #[test] + unconstrained fn for_each_remove_some() { + let contract_address = setup(); + let array = CapsuleArray::at(contract_address, SLOT); + + array.push(4); + array.push(5); + array.push(6); + + array.for_each(|index, _| { + if index == 1 { + array.remove(index); + } + }); + + assert_eq(array.len(), 2); + assert_eq(array.get(0), 4); + assert_eq(array.get(1), 6); + } + + #[test] + unconstrained fn for_each_remove_all() { + let contract_address = setup(); + let array = CapsuleArray::at(contract_address, SLOT); + + array.push(4); + array.push(5); + array.push(6); + + array.for_each(|index, _| { array.remove(index); }); + + assert_eq(array.len(), 0); + } + + // TODO: uncomment this test once OracleMock::count is implemented in the stdlib. + // #[test] + // unconstrained fn for_each_remove_all_no_copy() { + // let contract_address = setup(); + // let array = CapsuleArray::at(contract_address, SLOT); + + // array.push(4); + // array.push(5); + // array.push(6); + + // // We test that the copyCapsule was never called, which is the expensive operation we want to avoid. + // let mock = OracleMock::mock("copyCapsule"); + + // array.for_each(|index, _| { + // array.remove(index); + // }); + + // assert_eq(mock.count(), 0); + // } } diff --git a/noir-projects/aztec-nr/aztec/src/discovery/partial_notes.nr b/noir-projects/aztec-nr/aztec/src/discovery/partial_notes.nr index 212492ea45e8..63e11becf18f 100644 --- a/noir-projects/aztec-nr/aztec/src/discovery/partial_notes.nr +++ b/noir-projects/aztec-nr/aztec/src/discovery/partial_notes.nr @@ -93,20 +93,17 @@ pub unconstrained fn fetch_and_process_public_partial_note_completion_logs( [pending_partial_notes.len() as Field], ); - let mut i = 0; - while i < pending_partial_notes.len() { - let pending_partial_note: DeliveredPendingPartialNote = pending_partial_notes.get(i); - + pending_partial_notes.for_each(|i, pending_partial_note: DeliveredPendingPartialNote| { let maybe_log = get_log_by_tag(pending_partial_note.note_completion_log_tag); if maybe_log.is_none() { debug_log_format( "Found no completion logs for partial note with tag {}", [pending_partial_note.note_completion_log_tag], ); - i += 1 as u32; - // Note that we're not removing the pending partial note from the PXE DB, so we will continue searching - // for this tagged log when performing message discovery in the future until we either find it or the - // entry is somehow removed from the PXE DB. + + // Note that we're not removing the pending partial note from the capsule array, so we will continue + // searching for this tagged log when performing message discovery in the future until we either find it or + // the entry is somehow removed from the array. } else { debug_log_format( "Completion log found for partial note with tag {}", @@ -174,11 +171,8 @@ pub unconstrained fn fetch_and_process_public_partial_note_completion_logs( // TODO(#11627): only remove the pending entry if we actually process a log that results in the note // being completed. pending_partial_notes.remove(i); - - // We don't increment `i` here, because CapsuleArray is contiguous and its `remove(...)` function - // shifts the elements to the left if the removed element is not the last element. } - } + }); } fn decode_partial_note_private_msg( diff --git a/noir-projects/aztec-nr/aztec/src/discovery/private_logs.nr b/noir-projects/aztec-nr/aztec/src/discovery/private_logs.nr index 5331bcb85b72..58f0a0f0d023 100644 --- a/noir-projects/aztec-nr/aztec/src/discovery/private_logs.nr +++ b/noir-projects/aztec-nr/aztec/src/discovery/private_logs.nr @@ -37,23 +37,13 @@ pub unconstrained fn fetch_and_process_private_tagged_logs( // retrieved and processed here. sync_notes(PENDING_TAGGED_LOG_ARRAY_BASE_SLOT); - // Get logs from capsules + // Get the logs from the capsule array and process them one by one let logs = CapsuleArray::::at(contract_address, PENDING_TAGGED_LOG_ARRAY_BASE_SLOT); - let len = logs.len(); - - // We iterate over the logs in reverse order to avoid shifting elements. - let mut i = len; - while i > 0 { - i -= 1; - - // Get and process each log - let log = logs.get(i); + logs.for_each(|i, log: PendingTaggedLog| { process_log(contract_address, compute_note_hash_and_nullifier, log); - - // Remove the log from the capsule array logs.remove(i); - } + }); } /// Processes a log's ciphertext by decrypting it and then searching the plaintext for private notes or partial notes.