Add proptest for breakdown-reveal aggregation (#1470)

* 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.

* Include slow unit tests in CI

.github/workflows/check.yml

@@ -172,6 +172,9 @@ jobs:
             target/
           key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.toml') }}
 
+      - name: Slow Unit Tests
+        run: cargo test -p ipa-core --lib -- mpc_proptest semi_honest_with_dp_slow gen_binomial_noise_16_breakdowns
+
       - name: Integration Tests - Compact Gate
         run: cargo test --release --test "compact_gate" --no-default-features --features "cli web-app real-world-infra test-fixture compact-gate"
 

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_mpc_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/hybrid/breakdown_reveal.rs

@@ -1,4 +1,4 @@
-use std::{convert::Infallible, pin::pin};
+use std::{convert::Infallible, iter::repeat, pin::pin};
 
 use futures::stream;
 use futures_util::{StreamExt, TryStreamExt};
@@ -76,6 +76,14 @@ where
     BitDecomposed<Replicated<Boolean, B>>:
         for<'a> TransposeFrom<&'a [Replicated<V>; B], Error = Infallible>,
 {
+    // This was checked early in the protocol, but we need to check again here, in case
+    // there were no matching pairs of reports.
+    if attributed_values.is_empty() {
+        return Ok(BitDecomposed::new(
+            repeat(Replicated::<Boolean, B>::ZERO).take(usize::try_from(HV::BITS).unwrap()),
+        ));
+    }
+
     // Apply DP padding for Breakdown Reveal Aggregation
     let attributed_values_padded = apply_dp_padding::<_, AggregateableHybridReport<BK, V>, B>(
         ctx.narrow(&Step::PaddingDp),
@@ -290,6 +298,78 @@ pub mod tests {
         (inputs, expectation)
     }
 
+    #[test]
+    fn empty() {
+        run_with::<_, _, 3>(|| async {
+            let world = TestWorld::<WithShards<1>>::with_shards(TestWorldConfig::default());
+            let expectation = vec![0; 32];
+            let inputs: Vec<TestAggregateableHybridReport> = vec![];
+            let result: Vec<_> = world
+                .semi_honest(inputs.into_iter(), |ctx, input_rows| async move {
+                    let r: Vec<Replicated<BA8>> =
+                        breakdown_reveal_aggregation::<_, BA5, BA3, BA8, 32>(
+                            ctx,
+                            input_rows,
+                            &PaddingParameters::no_padding(),
+                        )
+                        .map_ok(|d: BitDecomposed<Replicated<Boolean, 32>>| {
+                            Vec::transposed_from(&d).unwrap()
+                        })
+                        .await
+                        .unwrap();
+                    r
+                })
+                .await
+                .reconstruct();
+            let result = result
+                .first()
+                .unwrap()
+                .iter()
+                .map(|&v| v.as_u128())
+                .collect::<Vec<_>>();
+            assert_eq!(32, result.len());
+            assert_eq!(result, expectation);
+        });
+    }
+
+    #[test]
+    fn single() {
+        // Test that the output is padded to the full size, when there are not enough inputs
+        // for the computation to naturally grow to the full size.
+        run_with::<_, _, 3>(|| async {
+            let world = TestWorld::<WithShards<1>>::with_shards(TestWorldConfig::default());
+            let mut expectation = vec![0; 32];
+            expectation[0] = 7;
+            let expectation = expectation; // no more mutability for safety
+            let inputs = vec![input_row(0, 7)];
+            let result: Vec<_> = world
+                .semi_honest(inputs.into_iter(), |ctx, input_rows| async move {
+                    let r: Vec<Replicated<BA8>> =
+                        breakdown_reveal_aggregation::<_, BA5, BA3, BA8, 32>(
+                            ctx,
+                            input_rows,
+                            &PaddingParameters::no_padding(),
+                        )
+                        .map_ok(|d: BitDecomposed<Replicated<Boolean, 32>>| {
+                            Vec::transposed_from(&d).unwrap()
+                        })
+                        .await
+                        .unwrap();
+                    r
+                })
+                .await
+                .reconstruct();
+            let result = result
+                .first()
+                .unwrap()
+                .iter()
+                .map(|&v| v.as_u128())
+                .collect::<Vec<_>>();
+            assert_eq!(32, result.len());
+            assert_eq!(result, expectation);
+        });
+    }
+
     #[test]
     fn breakdown_reveal_semi_honest_happy_path() {
         // if shuttle executor is enabled, run this test only once.
@@ -433,3 +513,164 @@ pub mod tests {
         });
     }
 }
+
+#[cfg(all(test, unit_test))]
+mod proptests {
+    use std::{cmp::min, time::Duration};
+
+    use futures::TryFutureExt;
+    use proptest::{prelude::*, prop_compose};
+
+    use crate::{
+        const_assert,
+        ff::{
+            boolean::Boolean,
+            boolean_array::{BA3, BA32, BA5, BA8},
+            U128Conversions,
+        },
+        protocol::{
+            hybrid::breakdown_reveal::breakdown_reveal_aggregation,
+            ipa_prf::oprf_padding::PaddingParameters,
+        },
+        secret_sharing::{
+            replicated::semi_honest::AdditiveShare as Replicated, BitDecomposed, SharedValue,
+            TransposeFrom,
+        },
+        test_fixture::{
+            hybrid::{TestAggregateableHybridReport, TestIndistinguishableHybridReport},
+            mpc_proptest_config_with_cases, Reconstruct, Runner, TestWorld, TestWorldConfig,
+            WithShards,
+        },
+    };
+
+    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;
+    const PROP_SHARDS: usize = 2;
+
+    // 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<TestAggregateableHybridReport>,
+        expected: Vec<u32>,
+        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 From<(u32, u32)> for TestAggregateableHybridReport {
+        fn from(value: (u32, u32)) -> Self {
+            TestAggregateableHybridReport {
+                match_key: (),
+                breakdown_key: value.0,
+                value: value.1,
+            }
+        }
+    }
+
+    prop_compose! {
+        fn inputs(max_len: usize)
+                 (
+                     len in 0..=max_len,
+                 )
+                 (
+                     len in Just(len),
+                     inputs in prop::collection::vec((
+                        0u32..u32::try_from(PROP_BUCKETS).unwrap(),
+                        0u32..1 << PropTriggerValue::BITS,
+                    ).prop_map(Into::into), len),
+                 )
+        -> AggregatePropTestInputs {
+            let mut expected = vec![0; PROP_BUCKETS];
+            for input in &inputs {
+                let TestIndistinguishableHybridReport {
+                    match_key: (),
+                    breakdown_key: bk,
+                    value: tv,
+                } = *input;
+                let bk = usize::try_from(bk).unwrap();
+                expected[bk] = min(expected[bk] + tv, (1 << PropHistogramValue::BITS) - 1);
+            }
+
+            AggregatePropTestInputs {
+                inputs,
+                expected,
+                len,
+            }
+        }
+    }
+
+    proptest! {
+        #![proptest_config(mpc_proptest_config_with_cases(100))]
+        #[test]
+        fn breakdown_reveal_mpc_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 config = TestWorldConfig {
+                    seed,
+                    timeout: Some(Duration::from_secs(20)),
+                    ..Default::default()
+                };
+                let result = TestWorld::<WithShards<PROP_SHARDS>>::with_config(&config)
+                    .malicious(inputs.into_iter(), |ctx, inputs| async move {
+                        breakdown_reveal_aggregation::<
+                            _,
+                            PropBreakdownKey,
+                            PropTriggerValue,
+                            PropHistogramValue,
+                            {PropBucketsBitVec::BITS as usize},
+                        >(
+                            ctx,
+                            inputs,
+                            &PaddingParameters::no_padding(),
+                        )
+                        .map_ok(|d: BitDecomposed<Replicated<Boolean, PROP_BUCKETS>>| {
+                            Vec::transposed_from(&d).unwrap()
+                        })
+                        .await
+                        .unwrap()
+                    })
+                    .await
+                    .reconstruct();
+
+                let initial = vec![0; PROP_BUCKETS];
+                let result = result
+                    .iter()
+                    .fold(initial, |mut acc, vec: &Vec<PropHistogramValue>| {
+                        acc.iter_mut()
+                            .zip(vec)
+                            .for_each(|(a, &b)| {
+                                *a = min(
+                                    *a + u32::try_from(b.as_u128()).unwrap(),
+                                    (1 << PropHistogramValue::BITS) - 1,
+                                );
+                            });
+                        acc
+                    })
+                    .into_iter()
+                    .collect::<Vec<_>>();
+                assert_eq!(result, expected);
+            });
+        }
+    }
+}