implement Shuffleable for IndistinguishableHybridReport (#1435)

* implement Shuffleable for IndistinguishableHybridReport

* pr feedback

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

@@ -1,6 +1,8 @@
 pub(crate) mod oprf;
 pub(crate) mod step;
 
+use tracing::{info_span, Instrument};
+
 use crate::{
     error::Error,
     ff::{
@@ -81,8 +83,11 @@ where
     )
     .await?;
 
-    // TODO shuffle input rows
-    let shuffled_input_rows = padded_input_rows;
+    let shuffled_input_rows = ctx
+        .narrow(&Step::InputShuffle)
+        .shuffle(padded_input_rows)
+        .instrument(info_span!("shuffle_inputs"))
+        .await?;
 
     let _sharded_reports = compute_prf_and_reshard(ctx.clone(), shuffled_input_rows).await?;
 

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

@@ -5,6 +5,8 @@ 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)]
+    InputShuffle,
     #[step(child = crate::protocol::ipa_prf::boolean_ops::step::Fp25519ConversionStep)]
     ConvertFp25519,
     #[step(child = crate::protocol::context::step::DzkpValidationProtocolStep)]

ipa-core/src/report/hybrid.rs

@@ -39,15 +39,21 @@ use crate::{
     const_assert_eq,
     error::{BoxError, Error},
     ff::{
-        boolean_array::{BA3, BA64, BA8},
+        boolean_array::{
+            BooleanArray, BooleanArrayReader, BooleanArrayWriter, BA112, BA3, BA64, BA8,
+        },
         Serializable,
     },
     hpke::{
         open_in_place, seal_in_place, CryptError, EncapsulationSize, PrivateKeyRegistry,
         PublicKeyRegistry, TagSize,
     },
+    protocol::ipa_prf::shuffle::Shuffleable,
     report::hybrid_info::{HybridConversionInfo, HybridImpressionInfo, HybridInfo},
-    secret_sharing::{replicated::semi_honest::AdditiveShare as Replicated, SharedValue},
+    secret_sharing::{
+        replicated::{semi_honest::AdditiveShare as Replicated, ReplicatedSecretSharing},
+        SharedValue,
+    },
     sharding::ShardIndex,
 };
 
@@ -682,8 +688,8 @@ pub type PrfHybridReport<BK, V> = IndistinguishableHybridReport<BK, V, u64>;
 #[derive(Clone, Debug, Eq, PartialEq, Default)]
 pub struct IndistinguishableHybridReport<BK, V, MK = Replicated<BA64>>
 where
-    BK: SharedValue,
-    V: SharedValue,
+    BK: BooleanArray,
+    V: BooleanArray,
 {
     pub match_key: MK,
     pub value: Replicated<V>,
@@ -692,20 +698,39 @@ where
 
 impl<BK, V> IndistinguishableHybridReport<BK, V>
 where
-    BK: SharedValue,
-    V: SharedValue,
+    BK: BooleanArray,
+    V: BooleanArray,
 {
     pub const ZERO: Self = Self {
         match_key: Replicated::<BA64>::ZERO,
         value: Replicated::<V>::ZERO,
         breakdown_key: Replicated::<BK>::ZERO,
     };
+
+    fn join_fields(match_key: BA64, value: V, breakdown_key: BK) -> <Self as Shuffleable>::Share {
+        let mut share = <Self as Shuffleable>::Share::ZERO;
+
+        BooleanArrayWriter::new(&mut share)
+            .write(&match_key)
+            .write(&value)
+            .write(&breakdown_key);
+
+        share
+    }
+
+    fn split_fields(share: &<Self as Shuffleable>::Share) -> (BA64, V, BK) {
+        let bits = BooleanArrayReader::new(share);
+        let (match_key, bits) = bits.read();
+        let (value, bits) = bits.read();
+        let (breakdown_key, _) = bits.read();
+        (match_key, value, breakdown_key)
+    }
 }
 
 impl<BK, V> From<Replicated<BA64>> for IndistinguishableHybridReport<BK, V>
 where
-    BK: SharedValue,
-    V: SharedValue,
+    BK: BooleanArray,
+    V: BooleanArray,
 {
     fn from(match_key: Replicated<BA64>) -> Self {
         Self {
@@ -718,8 +743,8 @@ where
 
 impl<BK, V> From<HybridReport<BK, V>> for IndistinguishableHybridReport<BK, V>
 where
-    BK: SharedValue,
-    V: SharedValue,
+    BK: BooleanArray,
+    V: BooleanArray,
 {
     fn from(report: HybridReport<BK, V>) -> Self {
         match report {
@@ -731,8 +756,8 @@ where
 
 impl<BK, V> From<HybridImpressionReport<BK>> for IndistinguishableHybridReport<BK, V>
 where
-    BK: SharedValue,
-    V: SharedValue,
+    BK: BooleanArray,
+    V: BooleanArray,
 {
     fn from(impression_report: HybridImpressionReport<BK>) -> Self {
         Self {
@@ -745,8 +770,8 @@ where
 
 impl<BK, V> From<HybridConversionReport<V>> for IndistinguishableHybridReport<BK, V>
 where
-    BK: SharedValue,
-    V: SharedValue,
+    BK: BooleanArray,
+    V: BooleanArray,
 {
     fn from(conversion_report: HybridConversionReport<V>) -> Self {
         Self {
@@ -757,6 +782,54 @@ where
     }
 }
 
+impl<BK, V> Shuffleable for IndistinguishableHybridReport<BK, V>
+where
+    BK: BooleanArray,
+    V: BooleanArray,
+{
+    // this requires BK:BAXX + V:BAYY  such that XX + YY <= 48
+    // this is checked in a debud_assert call in ::new below
+    // PERF OPPORTUNITY
+    // note that BA96 would likely be a better fit here. however, that would require a `BA128`
+    // in order to use `impl_malicious_shuffle_share!` and `BA128` cannot currently be
+    // implemented with `the boolean_array_impl!` macro as the trait `secret_sharing::Block`
+    // is not implemented for `bitvec::array::BitArray<[u8; 16]>`
+    type Share = BA112;
+
+    fn left(&self) -> Self::Share {
+        Self::join_fields(
+            ReplicatedSecretSharing::left(&self.match_key),
+            self.value.left(),
+            self.breakdown_key.left(),
+        )
+    }
+
+    fn right(&self) -> Self::Share {
+        Self::join_fields(
+            ReplicatedSecretSharing::right(&self.match_key),
+            self.value.right(),
+            self.breakdown_key.right(),
+        )
+    }
+
+    fn new(l: Self::Share, r: Self::Share) -> Self {
+        debug_assert!(
+            BA64::BITS + BK::BITS + V::BITS <= Self::Share::BITS,
+            "share type {} is too small",
+            std::any::type_name::<Self::Share>(),
+        );
+
+        let left = Self::split_fields(&l);
+        let right = Self::split_fields(&r);
+
+        Self {
+            match_key: ReplicatedSecretSharing::new(left.0, right.0),
+            value: ReplicatedSecretSharing::new(left.1, right.1),
+            breakdown_key: ReplicatedSecretSharing::new(left.2, right.2),
+        }
+    }
+}
+
 impl PrfHybridReport<BA8, BA3> {
     const PRF_MK_SZ: usize = 8;
     const V_SZ: usize = <Replicated<BA3> as Serializable>::Size::USIZE;