feat: add CapsuleArray::for_each

noir-projects/aztec-nr/aztec/src/capsules/mod.nr

@@ -76,6 +76,42 @@ impl<T> CapsuleArray<T> {
         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<Env, let N: u32>(self, f: unconstrained fn[Env](u32, T) -> ())
+    where
+        T: Deserialize<N>,
+    {
+        // 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);
+    // }
 }

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<Env>(
         [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<Env>(
             // 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(

noir-projects/aztec-nr/aztec/src/discovery/private_logs.nr

@@ -37,23 +37,13 @@ pub unconstrained fn fetch_and_process_private_tagged_logs<Env>(
     // 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::<PendingTaggedLog>::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.