From 322e3772d1ab39e7fe24a25017def0c4a32ce1ed Mon Sep 17 00:00:00 2001
From: Alex Koshelev <koshelev@meta.com>
Date: Sun, 11 Aug 2024 22:54:49 -0700
Subject: [PATCH] Feedback

---
 ipa-core/src/protocol/context/malicious.rs |  5 ++++-
 ipa-core/src/protocol/context/validator.rs | 16 +++++++++++-----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/ipa-core/src/protocol/context/malicious.rs b/ipa-core/src/protocol/context/malicious.rs
index 685961b45..564b7e1d6 100644
--- a/ipa-core/src/protocol/context/malicious.rs
+++ b/ipa-core/src/protocol/context/malicious.rs
@@ -194,14 +194,17 @@ impl<'a, F: ExtendableField> Upgraded<'a, F> {
         });
     }
 
+    /// `TestWorld` malicious methods require access to r share to perform validation.
+    /// This method allows such access only in non-prod code.
     #[cfg(any(test, feature = "test-fixture"))]
     #[must_use]
     pub fn r(&self, record_id: RecordId) -> Replicated<F::ExtendedField> {
         self.r_share(record_id)
     }
 
+    /// It is intentionally not public, allows access to it only from within
+    /// this module
     fn r_share(&self, record_id: RecordId) -> Replicated<F::ExtendedField> {
-        // its unfortunate, but carrying references across mutex boundaries is not possible
         self.with_batch(record_id, |v| v.r_share().clone())
     }
 
diff --git a/ipa-core/src/protocol/context/validator.rs b/ipa-core/src/protocol/context/validator.rs
index 10405f21b..6e5a4de6f 100644
--- a/ipa-core/src/protocol/context/validator.rs
+++ b/ipa-core/src/protocol/context/validator.rs
@@ -211,11 +211,12 @@ impl<'a, F: ExtendableField> BatchValidator<'a, F> {
     /// If total records is not set.
     #[must_use]
     pub fn new(ctx: MaliciousContext<'a>) -> Self {
-        assert!(
-            ctx.total_records().is_specified(),
-            "Total records must be specified before creating the validator"
-        );
-        let total_records = ctx.total_records().count().unwrap();
+        let Some(total_records) = ctx.total_records().count() else {
+            panic!("Total records must be specified before creating the validator");
+        };
+
+        // TODO: Right now we set the batch work to be equal to active_work,
+        // but it does not need to be. We can make this configurable if needed.
         let records_per_batch = ctx.active_work().get().min(total_records);
 
         Self {
@@ -254,6 +255,11 @@ impl<F: ExtendableField> Malicious<'_, F> {
         let narrow_ctx = self
             .validate_ctx
             .narrow(&ValidateStep::RevealR)
+            // TODO: propagate_u_and_w, RevealR and CheckZero all use indeterminate record count
+            // to communicate data right away. We could make it better if we had support from
+            // compact gate infrastructure to override batch size per step. All of the steps
+            // above require batch size to be set to 1, but we know the total number of records
+            // sent through these channels (total_records / batch_size)
             .set_total_records(TotalRecords::Indeterminate);
         let r = <F as ExtendableField>::ExtendedField::from_array(
             &malicious_reveal(