From 7164402b75cac5ea93f7a17bd4c2370d87e5bfa7 Mon Sep 17 00:00:00 2001
From: Andy Leiserson <aleiserson@mozilla.com>
Date: Fri, 13 Dec 2024 11:15:08 -0800
Subject: [PATCH] Fix chunk calculation for hybrid aggregation

---
 ipa-core/src/protocol/hybrid/agg.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/ipa-core/src/protocol/hybrid/agg.rs b/ipa-core/src/protocol/hybrid/agg.rs
index cfaef9fe1..d03461058 100644
--- a/ipa-core/src/protocol/hybrid/agg.rs
+++ b/ipa-core/src/protocol/hybrid/agg.rs
@@ -18,6 +18,7 @@ use crate::{
     },
     report::hybrid::{AggregateableHybridReport, PrfHybridReport},
     secret_sharing::replicated::semi_honest::AdditiveShare as Replicated,
+    utils::non_zero_prev_power_of_two,
 };
 
 enum MatchEntry<BK, V>
@@ -112,7 +113,8 @@ where
 {
     let report_pairs = group_report_pairs_ordered(reports);
 
-    let chunk_size: usize = TARGET_PROOF_SIZE / (BK::BITS as usize + V::BITS as usize);
+    let chunk_size =
+        non_zero_prev_power_of_two(TARGET_PROOF_SIZE / (BK::BITS as usize + V::BITS as usize));
 
     let ctx = ctx.set_total_records(TotalRecords::specified(report_pairs.len())?);
 
@@ -121,7 +123,7 @@ where
             protocol: &HybridStep::GroupBySum,
             validate: &HybridStep::GroupBySumValidate,
         },
-        chunk_size.next_power_of_two(),
+        chunk_size,
     );
 
     let agg_ctx = dzkp_validator.context();