Hybrid shuffle

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

@@ -1,7 +1,9 @@
+pub(crate) mod shuffle;
 pub(crate) mod step;
 
 use step::HybridStep as Step;
 
+use self::shuffle::shuffle_hybrid_inputs;
 use crate::{
     error::Error,
     ff::{
@@ -19,6 +21,7 @@ use crate::{
     report::hybrid::IndistinguishableHybridReport,
     secret_sharing::replicated::semi_honest::AdditiveShare as Replicated,
 };
+
 // In theory, we could support (runtime-configured breakdown count) ≤ (compile-time breakdown count)
 // ≤ 2^|bk|, with all three values distinct, but at present, there is no runtime configuration and
 // the latter two must be equal. The implementation of `move_single_value_to_bucket` does support a
@@ -81,12 +84,14 @@ where
     }
 
     // Apply DP padding for OPRF
-    let _padded_input_rows = apply_dp_padding::<_, IndistinguishableHybridReport<BK, V>, B>(
+    let padded_input_rows = apply_dp_padding::<_, IndistinguishableHybridReport<BK, V>, B>(
         ctx.narrow(&Step::PaddingDp),
         input_rows,
         &dp_padding_params,
     )
     .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

@@ -0,0 +1,159 @@
+use crate::{
+    error::Error,
+    ff::boolean_array::{BooleanArray, BA112, BA144, BA64},
+    protocol::{
+        context::Context,
+        ipa_prf::{
+            boolean_ops::{expand_shared_array_in_place, extract_from_shared_array},
+            shuffle::Shuffle,
+        },
+    },
+    report::hybrid::IndistinguishableHybridReport,
+    secret_sharing::{
+        replicated::{semi_honest::AdditiveShare, ReplicatedSecretSharing},
+        SharedValue,
+    },
+};
+
+/// Shuffles a Vec of IndistinguishableHybridReport
+/// # Errors
+/// Propogates errors from ctx.shuffle
+#[tracing::instrument(name = "shuffle_inputs", skip_all)]
+pub async fn shuffle_hybrid_inputs<C, BK, V>(
+    ctx: C,
+    input: Vec<IndistinguishableHybridReport<BK, V>>,
+) -> Result<Vec<IndistinguishableHybridReport<BK, V>>, Error>
+where
+    C: Context + Shuffle,
+    BK: BooleanArray,
+    V: BooleanArray,
+{
+    let shuffle_input: Vec<AdditiveShare<BA112>> = input
+        .into_iter()
+        .map(|item| hybrid_report_to_shuffle_input::<BA112, BK, V>(&item))
+        .collect::<Vec<_>>();
+
+    let shuffled = ctx.shuffle::<BA112, BA144, _>(shuffle_input).await?;
+
+    Ok(shuffled
+        .into_iter()
+        .map(|item| shuffled_to_hybrid_report(&item))
+        .collect::<Vec<_>>())
+}
+
+/// Converts `IndistinguishableHybridReport` into
+/// an `AdditiveShare` needed for shuffle protocol.
+pub fn hybrid_report_to_shuffle_input<YS, BK, V>(
+    input: &IndistinguishableHybridReport<BK, V>,
+) -> AdditiveShare<YS>
+where
+    YS: BooleanArray,
+    BK: BooleanArray,
+    V: BooleanArray,
+{
+    let mut y = AdditiveShare::new(YS::ZERO, YS::ZERO);
+    expand_shared_array_in_place(&mut y, &input.match_key, 0);
+
+    let mut offset = BA64::BITS as usize;
+
+    expand_shared_array_in_place(&mut y, &input.breakdown_key, offset);
+
+    offset += BK::BITS as usize;
+    expand_shared_array_in_place(&mut y, &input.value, offset);
+
+    y
+}
+
+/// Converts an `AdditiveShare` obtained from shuffle protocol
+/// into an `IndistinguishableHybridReport`.
+pub fn shuffled_to_hybrid_report<YS, BK, V>(
+    input: &AdditiveShare<YS>,
+) -> IndistinguishableHybridReport<BK, V>
+where
+    YS: BooleanArray,
+    BK: BooleanArray,
+    V: BooleanArray,
+{
+    let match_key = extract_from_shared_array::<YS, BA64>(input, 0);
+
+    let mut offset = BA64::BITS as usize;
+
+    let breakdown_key = extract_from_shared_array::<YS, BK>(input, offset);
+
+    offset += BK::BITS as usize;
+    let value = extract_from_shared_array::<YS, V>(input, offset);
+
+    IndistinguishableHybridReport {
+        match_key,
+        value,
+        breakdown_key,
+    }
+}
+
+#[cfg(all(test, unit_test))]
+pub mod tests {
+    use rand::Rng;
+
+    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]
+    fn hybrid_shuffle_conversion() {
+        let mut rng = rand::thread_rng();
+        let report = IndistinguishableHybridReport::<BA8, BA3> {
+            match_key: AdditiveShare::new(rng.gen(), rng.gen()),
+            breakdown_key: AdditiveShare::new(rng.gen(), rng.gen()),
+            value: AdditiveShare::new(rng.gen(), rng.gen()),
+        };
+
+        let additive_share = hybrid_report_to_shuffle_input::<BA112, BA8, BA3>(&report);
+        assert_ne!(additive_share.left(), BA112::ZERO);
+        assert_ne!(additive_share.right(), BA112::ZERO);
+
+        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/protocol/hybrid/step.rs

@@ -5,4 +5,6 @@ pub(crate) enum HybridStep {
     ReshardByTag,
     #[step(child = crate::protocol::ipa_prf::oprf_padding::step::PaddingDpStep, name="padding_dp")]
     PaddingDp,
+    #[step(child = crate::protocol::ipa_prf::shuffle::step::OPRFShuffleStep)]
+    Shuffle,
 }

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,