test shuffle hybrid

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

@@ -3,7 +3,7 @@ pub(crate) mod step;
 
 use step::HybridStep as Step;
 
-use self::shuffle::shuffle_inputs;
+use self::shuffle::shuffle_hybrid_inputs;
 use crate::{
     error::Error,
     ff::{
@@ -91,7 +91,7 @@ where
     )
     .await?;
 
-    let _shuffled = shuffle_inputs(ctx.narrow(&Step::Shuffle), padded_input_rows).await?;
+    let _shuffled = shuffle_hybrid_inputs(ctx.narrow(&Step::Shuffle), padded_input_rows).await?;
 
     unimplemented!("protocol::hybrid::hybrid_protocol is not fully implemented")
 }

ipa-core/src/protocol/hybrid/shuffle.rs

@@ -15,8 +15,11 @@ use crate::{
     },
 };
 
+/// Shuffles a Vec of IndistinguishableHybridReport
+/// # Errors
+/// Propogates errors from ctx.shuffle
 #[tracing::instrument(name = "shuffle_inputs", skip_all)]
-pub async fn shuffle_inputs<C, BK, V>(
+pub async fn shuffle_hybrid_inputs<C, BK, V>(
     ctx: C,
     input: Vec<IndistinguishableHybridReport<BK, V>>,
 ) -> Result<Vec<IndistinguishableHybridReport<BK, V>>, Error>
@@ -91,14 +94,16 @@ where
 pub mod tests {
     use rand::Rng;
 
-    use super::{hybrid_report_to_shuffle_input, shuffled_to_hybrid_report};
+    use super::{hybrid_report_to_shuffle_input, shuffle_hybrid_inputs, shuffled_to_hybrid_report};
     use crate::{
         ff::boolean_array::{BA112, BA3, BA8},
         report::hybrid::IndistinguishableHybridReport,
         secret_sharing::{
             replicated::{semi_honest::AdditiveShare, ReplicatedSecretSharing},
             SharedValue,
         },
+        test_executor::run,
+        test_fixture::{hybrid::TestIndistinguishableHybridReport, Reconstruct, Runner, TestWorld},
     };
 
     #[test]
@@ -117,4 +122,38 @@ pub mod tests {
         let report_copy = shuffled_to_hybrid_report::<BA112, BA8, BA3>(&additive_share);
         assert_eq!(report, report_copy);
     }
+
+    #[test]
+    fn test_shuffle_hybrid_inputs() {
+        const BATCHSIZE: usize = 50;
+        run(|| async {
+            let world = TestWorld::default();
+
+            let mut rng = rand::thread_rng();
+            let mut records = Vec::new();
+
+            for _ in 0..BATCHSIZE {
+                records.push({
+                    TestIndistinguishableHybridReport {
+                        match_key: rng.gen::<u64>(),
+                        breakdown_key: rng.gen_range(0u32..1 << 8),
+                        value: rng.gen_range(0u32..1 << 3),
+                    }
+                });
+            }
+
+            let mut result: Vec<TestIndistinguishableHybridReport> = world
+                .semi_honest(records.clone().into_iter(), |ctx, input_rows| async move {
+                    shuffle_hybrid_inputs::<_, BA8, BA3>(ctx, input_rows)
+                        .await
+                        .unwrap()
+                })
+                .await
+                .reconstruct();
+            assert_ne!(result, records);
+            records.sort();
+            result.sort();
+            assert_eq!(result, records);
+        });
+    }
 }

ipa-core/src/test_fixture/hybrid.rs

@@ -1,7 +1,15 @@
-use std::collections::{HashMap, HashSet};
+use std::{
+    collections::{HashMap, HashSet},
+    iter::zip,
+};
+
+use rand::Rng;
 
 use crate::{
-    ff::{boolean_array::BooleanArray, U128Conversions},
+    ff::{
+        boolean_array::{BooleanArray, BA64},
+        U128Conversions,
+    },
     report::hybrid::IndistinguishableHybridReport,
     secret_sharing::{replicated::semi_honest::AdditiveShare as Replicated, IntoShares},
     test_fixture::sharing::Reconstruct,
@@ -13,7 +21,7 @@ pub enum TestHybridRecord {
     TestConversion { match_key: u64, value: u32 },
 }
 
-#[derive(PartialEq, Eq)]
+#[derive(Debug, Clone, Ord, PartialEq, PartialOrd, Eq)]
 pub struct TestIndistinguishableHybridReport {
     pub match_key: u64,
     pub value: u32,
@@ -51,6 +59,34 @@ where
     }
 }
 
+impl<BK, V> IntoShares<IndistinguishableHybridReport<BK, V>> for TestIndistinguishableHybridReport
+where
+    BK: BooleanArray + U128Conversions + IntoShares<Replicated<BK>>,
+    V: BooleanArray + U128Conversions + IntoShares<Replicated<V>>,
+{
+    fn share_with<R: Rng>(self, rng: &mut R) -> [IndistinguishableHybridReport<BK, V>; 3] {
+        let match_key = BA64::try_from(u128::from(self.match_key))
+            .unwrap()
+            .share_with(rng);
+        let breakdown_key = BK::try_from(self.breakdown_key.into())
+            .unwrap()
+            .share_with(rng);
+        let value = V::try_from(self.value.into()).unwrap().share_with(rng);
+
+        zip(zip(match_key, breakdown_key), value)
+            .map(
+                |((match_key_share, bk_share), value_share)| IndistinguishableHybridReport {
+                    match_key: match_key_share,
+                    breakdown_key: bk_share,
+                    value: value_share,
+                },
+            )
+            .collect::<Vec<_>>()
+            .try_into()
+            .unwrap()
+    }
+}
+
 struct HashmapEntry {
     breakdown_key: u32,
     total_value: u32,