Add proptest for breakdown-reveal aggregation

And some general proptest cleanup and improvements. Run longer versions
of the MPC proptests when EXEC_SLOW_TESTS is set.

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

@@ -943,7 +943,6 @@ mod tests {
     use proptest::{
         prelude::{Just, Strategy},
         prop_compose, prop_oneof, proptest,
-        test_runner::Config as ProptestConfig,
     };
     use rand::{distributions::Standard, prelude::Distribution};
 
@@ -973,7 +972,9 @@ mod tests {
         seq_join::{seq_join, SeqJoin},
         sharding::NotSharded,
         test_executor::run_random,
-        test_fixture::{join3v, Reconstruct, Runner, TestWorld, TestWorldConfig},
+        test_fixture::{
+            join3v, mpc_proptest_config, Reconstruct, Runner, TestWorld, TestWorldConfig,
+        },
     };
 
     async fn test_select_semi_honest<V>()
@@ -1329,34 +1330,26 @@ mod tests {
     }
 
     proptest! {
-        #![proptest_config(ProptestConfig::with_cases(20))]
+        #![proptest_config(mpc_proptest_config())]
         #[test]
-        fn batching_proptest((record_count, max_multiplications_per_gate) in batching()) {
-            println!("record_count {record_count} batch {max_multiplications_per_gate}");
-            // This condition is correct only for active_work = 16 and record size of 1 byte.
-            if max_multiplications_per_gate != 1 && max_multiplications_per_gate % 16 != 0 {
-                // TODO: #1300, read_size | batch_size.
-                // Note: for active work < 2048, read size matches active work.
-
-                // Besides read_size | batch_size, there is also a constraint
-                // something like active_work > read_size + batch_size - 1.
-                println!("skipping config due to read_size vs. batch_size constraints");
-            } else {
-                tokio::runtime::Runtime::new().unwrap().block_on(async {
-                    chained_multiplies_dzkp(record_count, max_multiplications_per_gate).await.unwrap();
-                    /*
-                    multi_select_malicious::<BA3>(record_count, max_multiplications_per_gate).await;
-                    multi_select_malicious::<BA8>(record_count, max_multiplications_per_gate).await;
-                    multi_select_malicious::<BA16>(record_count, max_multiplications_per_gate).await;
-                    */
-                    multi_select_malicious::<BA20>(record_count, max_multiplications_per_gate).await;
-                    /*
-                    multi_select_malicious::<BA32>(record_count, max_multiplications_per_gate).await;
-                    multi_select_malicious::<BA64>(record_count, max_multiplications_per_gate).await;
-                    multi_select_malicious::<BA256>(record_count, max_multiplications_per_gate).await;
-                    */
-                });
-            }
+        fn batching_proptest(
+            (record_count, max_multiplications_per_gate) in batching(),
+            protocol in 0..8,
+        ) {
+            println!("record_count {record_count} batch {max_multiplications_per_gate} protocol {protocol}");
+            tokio::runtime::Runtime::new().unwrap().block_on(async {
+                match protocol {
+                    0 => chained_multiplies_dzkp(record_count, max_multiplications_per_gate).await.unwrap(),
+                    1 => multi_select_malicious::<BA3>(record_count, max_multiplications_per_gate).await,
+                    2 => multi_select_malicious::<BA8>(record_count, max_multiplications_per_gate).await,
+                    3 => multi_select_malicious::<BA16>(record_count, max_multiplications_per_gate).await,
+                    4 => multi_select_malicious::<BA20>(record_count, max_multiplications_per_gate).await,
+                    5 => multi_select_malicious::<BA32>(record_count, max_multiplications_per_gate).await,
+                    6 => multi_select_malicious::<BA64>(record_count, max_multiplications_per_gate).await,
+                    7 => multi_select_malicious::<BA256>(record_count, max_multiplications_per_gate).await,
+                    _ => unreachable!(),
+                }
+            });
         }
     }
 

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

@@ -203,10 +203,19 @@ where
         intermediate_results = next_intermediate_results;
     }
 
-    Ok(intermediate_results
+    let mut result = intermediate_results
         .into_iter()
         .next()
-        .expect("aggregation input must not be empty"))
+        .expect("aggregation input must not be empty");
+
+    // If there were less than 2^(|ov| - |tv|) inputs, then we didn't add enough carries to produce
+    // a full-length output, so pad the output now.
+    result.resize(
+        usize::try_from(HV::BITS).unwrap(),
+        Replicated::<Boolean, B>::ZERO,
+    );
+
+    Ok(result)
 }
 
 /// Transforms the Breakdown key from a secret share into a revealed `usize`.
@@ -302,29 +311,38 @@ where
 
 #[cfg(all(test, any(unit_test, feature = "shuttle")))]
 pub mod tests {
+    use std::cmp::min;
+
     use futures::TryFutureExt;
+    use proptest::{prelude::*, prop_compose};
     use rand::seq::SliceRandom;
 
-    #[cfg(not(feature = "shuttle"))]
-    use crate::{ff::boolean_array::BA16, test_executor::run};
     use crate::{
+        const_assert,
         ff::{
             boolean::Boolean,
-            boolean_array::{BA3, BA5, BA8},
+            boolean_array::{BA3, BA32, BA5, BA8},
             U128Conversions,
         },
         protocol::ipa_prf::{
             aggregation::breakdown_reveal::breakdown_reveal_aggregation,
             oprf_padding::PaddingParameters,
-            prf_sharding::{AttributionOutputsTestInput, SecretSharedAttributionOutputs},
+            prf_sharding::{
+                AttributionOutputs, AttributionOutputsTestInput, SecretSharedAttributionOutputs,
+            },
         },
         rand::Rng,
         secret_sharing::{
-            replicated::semi_honest::AdditiveShare as Replicated, BitDecomposed, TransposeFrom,
+            replicated::semi_honest::AdditiveShare as Replicated, BitDecomposed, IntoShares,
+            SharedValue, TransposeFrom,
         },
         test_executor::run_with,
-        test_fixture::{Reconstruct, Runner, TestWorld},
+        test_fixture::{
+            mpc_proptest_config_with_cases, Reconstruct, ReconstructArr, Runner, TestWorld,
+        },
     };
+    #[cfg(not(feature = "shuttle"))]
+    use crate::{ff::boolean_array::BA16, test_executor::run};
 
     fn input_row(bk: usize, tv: u128) -> AttributionOutputsTestInput<BA5, BA3> {
         let bk: u128 = bk.try_into().unwrap();
@@ -445,4 +463,119 @@ pub mod tests {
             assert_eq!(result, expectation);
         });
     }
+
+    type PropBreakdownKey = BA5;
+    type PropTriggerValue = BA3;
+    type PropHistogramValue = BA8;
+    type PropBucketsBitVec = BA32;
+    const PROP_MAX_INPUT_LEN: usize = 2500;
+    const PROP_BUCKETS: usize = PropBucketsBitVec::BITS as usize;
+
+    // We want to capture everything in this struct for visibility in the output of failing runs,
+    // even if it isn't used by the test.
+    #[allow(dead_code)]
+    #[derive(Debug)]
+    struct AggregatePropTestInputs {
+        inputs: Vec<AttributionOutputs<usize, u32>>,
+        expected: BitDecomposed<PropBucketsBitVec>,
+        len: usize,
+    }
+
+    const_assert!(
+        PropHistogramValue::BITS < u32::BITS,
+        "(1 << PropHistogramValue::BITS) must fit in u32",
+    );
+
+    const_assert!(
+        PROP_BUCKETS <= 1 << PropBreakdownKey::BITS,
+        "PROP_BUCKETS must fit in PropBreakdownKey",
+    );
+
+    impl<BK, TV> From<(BK, TV)> for AttributionOutputs<BK, TV> {
+        fn from(value: (BK, TV)) -> Self {
+            AttributionOutputs {
+                attributed_breakdown_key_bits: value.0,
+                capped_attributed_trigger_value: value.1,
+            }
+        }
+    }
+
+    impl IntoShares<SecretSharedAttributionOutputs<PropBreakdownKey, PropTriggerValue>>
+        for AttributionOutputs<usize, u32>
+    {
+        fn share_with<R: Rng>(
+            self,
+            rng: &mut R,
+        ) -> [SecretSharedAttributionOutputs<PropBreakdownKey, PropTriggerValue>; 3] {
+            let [bk_0, bk_1, bk_2] = PropBreakdownKey::truncate_from(
+                u128::try_from(self.attributed_breakdown_key_bits).unwrap(),
+            )
+            .share_with(rng);
+            let [tv_0, tv_1, tv_2] =
+                PropTriggerValue::truncate_from(u128::from(self.capped_attributed_trigger_value))
+                    .share_with(rng);
+            [(bk_0, tv_0), (bk_1, tv_1), (bk_2, tv_2)].map(Into::into)
+        }
+    }
+
+    prop_compose! {
+        fn inputs(max_len: usize)
+                 (
+                     len in 1..=max_len,
+                 )
+                 (
+                     len in Just(len),
+                     inputs in prop::collection::vec((0..PROP_BUCKETS, 0u32..1 << PropTriggerValue::BITS).prop_map(Into::into), len),
+                 )
+        -> AggregatePropTestInputs {
+            let mut expected = [0; PROP_BUCKETS];
+            for input in &inputs {
+                let AttributionOutputs {
+                    attributed_breakdown_key_bits: bk,
+                    capped_attributed_trigger_value: tv,
+                } = *input;
+                expected[bk] = min(expected[bk] + tv, (1 << PropHistogramValue::BITS) - 1);
+            }
+
+            let expected = BitDecomposed::decompose(PropHistogramValue::BITS, |i| {
+                expected.iter().map(|v| Boolean::from((v >> i) & 1 == 1)).collect()
+            });
+
+            AggregatePropTestInputs {
+                inputs,
+                expected,
+                len,
+            }
+        }
+    }
+
+    proptest! {
+        #![proptest_config(mpc_proptest_config_with_cases(100))]
+        #[test]
+        fn breakdown_reveal_proptest(
+            input_struct in inputs(PROP_MAX_INPUT_LEN),
+            seed in any::<u64>(),
+        ) {
+            tokio::runtime::Runtime::new().unwrap().block_on(async {
+                let AggregatePropTestInputs {
+                    inputs,
+                    expected,
+                    ..
+                } = input_struct;
+                let result = TestWorld::with_seed(seed)
+                    .malicious(inputs.into_iter(), |ctx, inputs| async move {
+                        breakdown_reveal_aggregation::<_, _, _, PropHistogramValue, {PropBucketsBitVec::BITS as usize}>(
+                            ctx,
+                            inputs,
+                            &PaddingParameters::no_padding(),
+                        ).await
+                    })
+                    .await
+                    .map(Result::unwrap)
+                    .reconstruct_arr();
+
+                assert_eq!(result, expected);
+            });
+        }
+    }
 }

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

@@ -257,7 +257,7 @@ pub mod tests {
         helpers::Role,
         secret_sharing::{BitDecomposed, SharedValue},
         test_executor::run,
-        test_fixture::{ReconstructArr, Runner, TestWorld},
+        test_fixture::{mpc_proptest_config, ReconstructArr, Runner, TestWorld},
     };
 
     fn input_row<const B: usize>(tv_bits: usize, values: &[u32]) -> BitDecomposed<[Boolean; B]> {
@@ -482,7 +482,7 @@ pub mod tests {
     // Any of the supported aggregation configs can be used here (search for "aggregation output" in
     // transpose.rs). This small config keeps CI runtime within reason, however, it does not exercise
     // saturated addition at the output.
-    const PROP_MAX_INPUT_LEN: usize = 10;
+    const PROP_MAX_INPUT_LEN: usize = 100;
     const PROP_MAX_TV_BITS: usize = 3; // Limit: (1 << TV_BITS) must fit in u32
     const PROP_BUCKETS: usize = 8;
     type PropHistogramValue = BA8;
@@ -535,6 +535,7 @@ pub mod tests {
     }
 
     proptest! {
+        #![proptest_config(mpc_proptest_config())]
         #[test]
         fn aggregate_values_proptest(
             input_struct in arb_aggregate_values_inputs(PROP_MAX_INPUT_LEN),

ipa-core/src/test_fixture/mod.rs

@@ -197,3 +197,31 @@ pub fn bits_to_value<F: Field + U128Conversions>(x: &[F]) -> u128 {
 pub fn bits_to_field<F: Field + U128Conversions>(x: &[F]) -> F {
     F::try_from(bits_to_value(x)).unwrap()
 }
+
+/// Useful proptest configuration when testing MPC protocols.
+#[cfg(test)]
+#[must_use]
+pub fn mpc_proptest_config() -> proptest::prelude::ProptestConfig {
+    mpc_proptest_config_with_cases(proptest::prelude::ProptestConfig::default().cases)
+}
+
+#[cfg(test)]
+#[must_use]
+pub fn mpc_proptest_config_with_cases(cases: u32) -> proptest::prelude::ProptestConfig {
+    use std::cmp::max;
+
+    // TestWorld imposes a per-iteration timeout. In addition to that, we want to limit
+    // the shrinking time, since most of the shrinking attempts for a timeout case, are
+    // likely to also time out. Proptest also has a timeout mechanism, but it interacts
+    // badly with max_shrink_time, and is mostly redundant with the TestWorld timeout.
+    let mut config = proptest::prelude::ProptestConfig {
+        max_shrink_time: 60_000,
+        cases,
+        ..Default::default()
+    };
+    if std::env::var("EXEC_SLOW_TESTS").is_err() {
+        // Reduce the number of cases when not in slow-tests mode.
+        config.cases = max(cases / 20, 1);
+    }
+    config
+}