Review follow-up from private-attribution#1323

ipa-core/src/protocol/context/dzkp_malicious.rs

@@ -43,6 +43,10 @@ impl<'a, B: ShardBinding> DZKPUpgraded<'a, B> {
             // in tests; there shouldn't be a risk of deadlocks with one record per
             // batch; and UnorderedReceiver capacity (which is set from active_work)
             // must be at least two.
+            //
+            // Also rely on the protocol to ensure an appropriate active_work if
+            // records_per_batch is `usize::MAX` (unlimited batch size). Allocating
+            // storage for `usize::MAX` active records won't work.
             base_ctx.active_work()
         } else {
             // Adjust active_work to match records_per_batch. If it is less, we will

ipa-core/src/protocol/ipa_prf/aggregation/breakdown_reveal.rs

@@ -1,4 +1,4 @@
-use std::{convert::Infallible, mem, pin::pin};
+use std::{convert::Infallible, pin::pin};
 
 use futures::stream;
 use futures_util::{StreamExt, TryStreamExt};
@@ -34,7 +34,6 @@ use crate::{
         TransposeFrom, Vectorizable,
     },
     seq_join::seq_join,
-    utils::vec_chunks::vec_chunks,
 };
 
 /// Improved Aggregation a.k.a Aggregation revealing breakdown.
@@ -100,7 +99,8 @@ where
 
     while intermediate_results.len() > 1 {
         let mut record_ids = [RecordId::FIRST; AGGREGATE_DEPTH];
-        for chunk in vec_chunks(mem::take(&mut intermediate_results), agg_proof_chunk) {
+        let mut next_intermediate_results = Vec::new();
+        for chunk in intermediate_results.chunks(agg_proof_chunk) {
             let chunk_len = chunk.len();
             let validator = ctx.clone().dzkp_validator(
                 MaliciousProtocolSteps {
@@ -109,21 +109,22 @@ where
                 },
                 // We have to specify usize::MAX here because the procession through
                 // record IDs is different at each step of the reduction. The batch
-                // size is limited by `vec_chunks`, above.
+                // size is limited by `intermediate_results.chunks()`, above.
                 usize::MAX,
             );
             let result = aggregate_values::<_, HV, B>(
                 validator.context(),
-                stream::iter(chunk).map(Ok).boxed(),
+                stream::iter(chunk).map(|v| Ok(v.clone())).boxed(),
                 chunk_len,
                 Some(&mut record_ids),
             )
             .await?;
             validator.validate().await?;
             chunk_counter += 1;
-            intermediate_results.push(result);
+            next_intermediate_results.push(result);
         }
         depth += 1;
+        intermediate_results = next_intermediate_results;
     }
 
     Ok(intermediate_results

ipa-core/src/utils/mod.rs

@@ -2,7 +2,6 @@ pub mod array;
 pub mod arraychunks;
 #[cfg(target_pointer_width = "64")]
 mod power_of_two;
-pub mod vec_chunks;
 
 #[cfg(target_pointer_width = "64")]
 pub use power_of_two::NonZeroU32PowerOfTwo;