Merge from main

ipa-core/Cargo.toml

@@ -136,7 +136,7 @@ sha2 = "0.10"
 shuttle-crate = { package = "shuttle", version = "0.6.1", optional = true }
 thiserror = "1.0"
 time = { version = "0.3", optional = true }
-tokio = { version = "1.35", features = ["rt", "rt-multi-thread", "macros"] }
+tokio = { version = "1.35", features = ["fs", "rt", "rt-multi-thread", "macros"] }
 # TODO: axum-server holds onto 0.24 and we can't upgrade until they do. Or we move away from axum-server
 tokio-rustls = { version = "0.24", optional = true }
 tokio-stream = "0.1.14"
@@ -159,7 +159,7 @@ cfg_aliases = "0.1.1"
 command-fds = "0.2.2"
 hex = "0.4"
 permutation = "0.4.1"
-proptest = "1"
+proptest = "1.4"
 rustls = { version = "0.21", features = ["dangerous_configuration"] }
 tempfile = "3"
 

ipa-core/src/ff/boolean.rs

@@ -16,6 +16,16 @@ impl Block for bool {
 #[derive(Clone, Copy, PartialEq, Debug, Eq)]
 pub struct Boolean(bool);
 
+impl Boolean {
+    pub const TRUE: Boolean = Self(true);
+    pub const FALSE: Boolean = Self(false);
+
+    #[must_use]
+    pub fn as_u128(&self) -> u128 {
+        u128::from(bool::from(*self))
+    }
+}
+
 impl ExtendableField for Boolean {
     type ExtendedField = Gf32Bit;
 
@@ -139,7 +149,7 @@ impl Field for Boolean {
     const ONE: Boolean = Boolean(true);
 
     fn as_u128(&self) -> u128 {
-        bool::from(*self).into()
+        Boolean::as_u128(self)
     }
 
     fn truncate_from<T: Into<u128>>(v: T) -> Self {

ipa-core/src/ff/galois_field.rs

@@ -687,5 +687,11 @@ bit_array_impl!(
                 v
             }
         }
+
+        impl From<Gf2> for bool {
+            fn from(value: Gf2) -> Self {
+                value != Gf2::ZERO
+            }
+        }
     }
 );

ipa-core/src/helpers/transport/stream/input.rs

@@ -727,7 +727,6 @@ mod test {
 
         proptest::proptest! {
             #[test]
-            #[allow(clippy::ignored_unit_patterns)] // https://github.com/proptest-rs/proptest/issues/371
             fn test_records_stream_works_with_any_chunks(
                 (expected_bytes, chunked_bytes, _seed) in arb_expected_and_chunked_body(30, 100)
             ) {
@@ -836,7 +835,6 @@ mod test {
 
         proptest::proptest! {
             #[test]
-            #[allow(clippy::ignored_unit_patterns)] // https://github.com/proptest-rs/proptest/issues/371
             fn test_delimited_stream_works_with_any_chunks(
                 (expected_items, chunked_bytes, _seed) in arb_expected_and_chunked_body(100)
             ) {

ipa-core/src/hpke/mod.rs

@@ -360,7 +360,6 @@ mod tests {
         proptest::proptest! {
             #![proptest_config(ProptestConfig::with_cases(50))]
             #[test]
-            #[allow(clippy::ignored_unit_patterns)] // https://github.com/proptest-rs/proptest/issues/371
             fn arbitrary_ct_corruption(bad_byte in 0..23_usize, bad_bit in 0..7_usize, seed: [u8; 32]) {
                 let rng = StdRng::from_seed(seed);
                 let mut suite = EncryptionSuite::new(1, rng);
@@ -374,7 +373,6 @@ mod tests {
         proptest::proptest! {
             #![proptest_config(ProptestConfig::with_cases(50))]
             #[test]
-            #[allow(clippy::ignored_unit_patterns)] // https://github.com/proptest-rs/proptest/issues/371
              fn arbitrary_enc_corruption(bad_byte in 0..32_usize, bad_bit in 0..7_usize, seed: [u8; 32]) {
                 let rng = StdRng::from_seed(seed);
                 let mut suite = EncryptionSuite::new(1, rng);
@@ -419,7 +417,6 @@ mod tests {
         proptest::proptest! {
             #![proptest_config(ProptestConfig::with_cases(50))]
             #[test]
-            #[allow(clippy::ignored_unit_patterns)] // https://github.com/proptest-rs/proptest/issues/371
             fn arbitrary_info_corruption(corrupted_info_field in 1..5,
                                          site_domain in "[a-z]{10}",
                                          helper_origin in "[a-z]{10}",

ipa-core/src/protocol/attribution/aggregate_credit.rs

@@ -8,9 +8,9 @@ use crate::{
     ff::{Gf2, PrimeField, Serializable},
     protocol::{
         context::{UpgradableContext, UpgradedContext, Validator},
+        ipa_prf::prf_sharding::bucket::move_single_value_to_bucket,
         modulus_conversion::convert_bits,
-        sort::{bitwise_to_onehot, generate_permutation::ShuffledPermutationWrapper},
-        step::BitOpStep,
+        sort::generate_permutation::ShuffledPermutationWrapper,
         BasicProtocols, RecordId,
     },
     secret_sharing::{
@@ -23,11 +23,6 @@ use crate::{
     seq_join::seq_join,
 };
 
-/// This is the number of breakdown keys above which it is more efficient to SORT by breakdown key.
-/// Below this number, it's more efficient to just do a ton of equality checks.
-/// This number was determined empirically on 27 Feb 2023
-const SIMPLE_AGGREGATION_BREAK_EVEN_POINT: u32 = 32;
-
 /// Aggregation step for Oblivious Attribution protocol.
 /// # Panics
 /// It probably won't
@@ -56,16 +51,9 @@ where
     ShuffledPermutationWrapper<S, C::UpgradedContext<F>>: DowngradeMalicious<Target = Vec<u32>>,
 {
     let m_ctx = validator.context();
-
-    if max_breakdown_key <= SIMPLE_AGGREGATION_BREAK_EVEN_POINT {
-        let res = simple_aggregate_credit(m_ctx, breakdown_keys, capped_credits, max_breakdown_key)
-            .await?;
-        Ok((validator, res))
-    } else {
-        Err(Error::Unsupported(
-            format!("query uses {max_breakdown_key} breakdown keys; only {SIMPLE_AGGREGATION_BREAK_EVEN_POINT} are supported")
-        ))
-    }
+    let res =
+        simple_aggregate_credit(m_ctx, breakdown_keys, capped_credits, max_breakdown_key).await?;
+    Ok((validator, res))
 }
 
 async fn simple_aggregate_credit<F, C, IC, IB, S>(
@@ -84,17 +72,10 @@ where
     S: LinearSecretSharing<F> + BasicProtocols<C, F> + Serializable + 'static,
 {
     let record_count = breakdown_keys.len();
-    // The number of records we compute is currently too high as the last row cannot have
-    // any credit associated with it.  TODO: don't compute that row when cap > 1.
-
-    let to_take = usize::try_from(max_breakdown_key).unwrap();
     let valid_bits_count = u32::BITS - (max_breakdown_key - 1).leading_zeros();
 
-    let equality_check_context = ctx
-        .narrow(&Step::ComputeEqualityChecks)
-        .set_total_records(record_count);
-    let check_times_credit_context = ctx
-        .narrow(&Step::CheckTimesCredit)
+    let move_value_to_bucket_context = ctx
+        .narrow(&Step::MoveValueToBucket)
         .set_total_records(record_count);
 
     let converted_bk = convert_bits(
@@ -110,23 +91,21 @@ where
             .zip(stream_iter(capped_credits))
             .enumerate()
             .map(|(i, (bk, cred))| {
-                let ceq = &equality_check_context;
-                let cmul = &check_times_credit_context;
+                let ctx = move_value_to_bucket_context.clone();
                 async move {
-                    let equality_checks = bitwise_to_onehot(ceq.clone(), i, &bk?).await?;
-                    ceq.try_join(equality_checks.into_iter().take(to_take).enumerate().map(
-                        |(check_idx, check)| {
-                            let step = BitOpStep::from(check_idx);
-                            let c = cmul.narrow(&step);
-                            let record_id = RecordId::from(i);
-                            let credit = &cred;
-                            async move { check.multiply(credit, c, record_id).await }
-                        },
-                    ))
+                    move_single_value_to_bucket(
+                        ctx,
+                        RecordId::from(i),
+                        bk.unwrap(),
+                        cred,
+                        usize::try_from(max_breakdown_key).unwrap(),
+                        true,
+                    )
                     .await
                 }
             }),
     );
+
     let aggregate = increments
         .try_fold(
             vec![S::ZERO; max_breakdown_key as usize],
@@ -143,8 +122,7 @@ where
 
 #[derive(Step)]
 pub(crate) enum Step {
-    ComputeEqualityChecks,
-    CheckTimesCredit,
+    MoveValueToBucket,
     ModConvBreakdownKeyBits,
 }
 

ipa-core/src/protocol/basics/share_known_value.rs

@@ -12,6 +12,10 @@ use crate::{
     },
 };
 
+/// Produce a share of some pre-determined constant.
+///
+/// The context is only used to determine the helper role. It is not used for communication or PRSS,
+/// and it is not necessary to use a uniquely narrowed context.
 pub trait ShareKnownValue<C: Context, V: SharedValue> {
     fn share_known_value(ctx: &C, value: V) -> Self;
 }

ipa-core/src/protocol/boolean/comparison.rs

@@ -558,7 +558,6 @@ mod tests {
 
     proptest! {
         #[test]
-        #[allow(clippy::ignored_unit_patterns)] // https://github.com/proptest-rs/proptest/issues/371
         fn gt_fp31_proptest(a in 0..Fp31::PRIME, c in 0..Fp31::PRIME) {
             type F = Fp31;
             let r = thread_rng().gen::<F>();
@@ -567,7 +566,6 @@ mod tests {
         }
 
         #[test]
-        #[allow(clippy::ignored_unit_patterns)] // https://github.com/proptest-rs/proptest/issues/371
         fn gt_fp_32bit_prime_proptest(a in 0..Fp32BitPrime::PRIME, c in 0..Fp32BitPrime::PRIME) {
             type F = Fp32BitPrime;
             let r = thread_rng().gen::<F>();

ipa-core/src/protocol/dp/insecure.rs

@@ -197,7 +197,6 @@ mod test {
     proptest! {
         #![proptest_config(ProptestConfig::with_cases(50))]
         #[test]
-        #[allow(clippy::ignored_unit_patterns)] // https://github.com/proptest-rs/proptest/issues/371
         fn output_differentially_private(
             rng_seed: u64,
             epsilon in 1..255_u8,

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

@@ -988,7 +988,6 @@ pub mod tests {
 
         proptest! {
             #[test]
-            #[allow(clippy::ignored_unit_patterns)] // https://github.com/proptest-rs/proptest/issues/371
             fn serde(timestamp in 0..u128::MAX, match_key in 0..u64::MAX, trigger_bit in 0..u128::MAX, breakdown_key in 0..u128::MAX, trigger_value in 0..u128::MAX, seed in 0..u128::MAX) {
                 serde_internal::<Fp31>(timestamp, match_key, trigger_bit, breakdown_key, trigger_value, seed);
                 serde_internal::<Fp32BitPrime>(timestamp, match_key, trigger_bit, breakdown_key, trigger_value, seed);
@@ -1129,9 +1128,9 @@ pub mod tests {
                 cap_one(),
                 SemiHonest,
                 PerfMetrics {
-                    records_sent: 14_421,
-                    bytes_sent: 47_100,
-                    indexed_prss: 19_137,
+                    records_sent: 14_397,
+                    bytes_sent: 47_004,
+                    indexed_prss: 19_113,
                     seq_prss: 1118,
                 },
             )
@@ -1144,9 +1143,9 @@ pub mod tests {
                 cap_three(),
                 SemiHonest,
                 PerfMetrics {
-                    records_sent: 21_756,
-                    bytes_sent: 76_440,
-                    indexed_prss: 28_146,
+                    records_sent: 21_732,
+                    bytes_sent: 76_344,
+                    indexed_prss: 28_122,
                     seq_prss: 1118,
                 },
             )
@@ -1159,9 +1158,9 @@ pub mod tests {
                 cap_one(),
                 Malicious,
                 PerfMetrics {
-                    records_sent: 35_163,
-                    bytes_sent: 130_068,
-                    indexed_prss: 72_447,
+                    records_sent: 35_115,
+                    bytes_sent: 129_876,
+                    indexed_prss: 72_375,
                     seq_prss: 1132,
                 },
             )
@@ -1174,9 +1173,9 @@ pub mod tests {
                 cap_three(),
                 Malicious,
                 PerfMetrics {
-                    records_sent: 53_865,
-                    bytes_sent: 204_876,
-                    indexed_prss: 109_734,
+                    records_sent: 53_817,
+                    bytes_sent: 204_684,
+                    indexed_prss: 109_662,
                     seq_prss: 1132,
                 },
             )