Use sharded shuffle in Hybrid

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

@@ -14,13 +14,11 @@ use crate::{
             dzkp_validator::DZKPValidator, Context, DZKPUpgraded, MaliciousProtocolSteps,
             ShardedContext, UpgradableContext,
         },
+        hybrid::step::AggregationStep as Step,
         ipa_prf::{
-            aggregation::{
-                aggregate_values, aggregate_values_proof_chunk, step::AggregationStep as Step,
-                AGGREGATE_DEPTH,
-            },
+            aggregation::{aggregate_values, aggregate_values_proof_chunk, AGGREGATE_DEPTH},
             oprf_padding::{apply_dp_padding, PaddingParameters},
-            shuffle::Shuffle,
+            shuffle::ShardedShuffle,
         },
         BooleanProtocols, RecordId,
     },
@@ -66,7 +64,7 @@ pub async fn breakdown_reveal_aggregation<C, BK, V, HV, const B: usize>(
     padding_params: &PaddingParameters,
 ) -> Result<BitDecomposed<Replicated<Boolean, B>>, Error>
 where
-    C: UpgradableContext + Shuffle + ShardedContext,
+    C: UpgradableContext + ShardedShuffle + ShardedContext,
     Boolean: FieldSimd<B>,
     Replicated<Boolean, B>: BooleanProtocols<DZKPUpgraded<C>, B>,
     BK: BooleanArray + U128Conversions,
@@ -94,7 +92,7 @@ where
 
     let attributions = ctx
         .narrow(&Step::Shuffle)
-        .shuffle(attributed_values_padded)
+        .sharded_shuffle(attributed_values_padded)
         .instrument(info_span!("shuffle_attribution_outputs"))
         .await?;
 

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

@@ -33,7 +33,7 @@ use crate::{
         ipa_prf::{
             oprf_padding::{apply_dp_padding, PaddingParameters},
             prf_eval::PrfSharing,
-            shuffle::Shuffle,
+            shuffle::ShardedShuffle,
         },
         prss::FromPrss,
         BooleanProtocols,
@@ -79,7 +79,7 @@ pub async fn hybrid_protocol<'ctx, C, BK, V, HV, const SS_BITS: usize, const B:
 where
     C: UpgradableContext
         + 'ctx
-        + Shuffle
+        + ShardedShuffle
         + ShardedContext
         + FinalizerContext<FinalizingContext = DZKPUpgraded<C>>,
     BK: BreakdownKey<B>,
@@ -121,7 +121,7 @@ where
 
     let shuffled_input_rows = ctx
         .narrow(&Step::InputShuffle)
-        .shuffle(padded_input_rows)
+        .sharded_shuffle(padded_input_rows)
         .instrument(info_span!("shuffle_inputs"))
         .await?;
 

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

@@ -181,9 +181,11 @@ where
 
     // reshard reports based on OPRF values. This ensures at the end of this function
     // reports with the same value end up on the same shard.
-    reshard_try_stream(ctx, report_stream, |ctx, _, report| {
-        report.match_key % ctx.shard_count()
-    })
+    reshard_try_stream(
+        ctx.narrow(&HybridStep::ReshardByPrf),
+        report_stream,
+        |ctx, _, report| report.match_key % ctx.shard_count(),
+    )
     .await
 }
 

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

@@ -5,7 +5,7 @@ pub(crate) enum HybridStep {
     ReshardByTag,
     #[step(child = crate::protocol::ipa_prf::oprf_padding::step::PaddingDpStep, name="report_padding_dp")]
     PaddingDp,
-    #[step(child = crate::protocol::ipa_prf::shuffle::step::OPRFShuffleStep)]
+    #[step(child = crate::protocol::ipa_prf::shuffle::step::ShardedShuffleStep)]
     InputShuffle,
     #[step(child = crate::protocol::ipa_prf::boolean_ops::step::Fp25519ConversionStep)]
     ConvertFp25519,
@@ -19,7 +19,7 @@ pub(crate) enum HybridStep {
     GroupBySum,
     #[step(child = crate::protocol::context::step::DzkpValidationProtocolStep)]
     GroupBySumValidate,
-    #[step(child = crate::protocol::ipa_prf::aggregation::step::AggregationStep)]
+    #[step(child = AggregationStep)]
     Aggregate,
     #[step(child = FinalizeSteps)]
     Finalize,
@@ -40,3 +40,18 @@ pub(crate) enum FinalizeSteps {
     #[step(child = crate::protocol::context::step::DzkpValidationProtocolStep)]
     Validate,
 }
+
+#[derive(CompactStep)]
+pub(crate) enum AggregationStep {
+    #[step(child = crate::protocol::ipa_prf::oprf_padding::step::PaddingDpStep, name="padding_dp")]
+    PaddingDp,
+    #[step(child = crate::protocol::ipa_prf::shuffle::step::ShardedShuffleStep)]
+    Shuffle,
+    Reveal,
+    #[step(child = crate::protocol::context::step::DzkpValidationProtocolStep)]
+    RevealValidate, // only partly used -- see code
+    #[step(count = 4, child = crate::protocol::ipa_prf::aggregation::step::AggregateChunkStep, name = "chunks")]
+    Aggregate(usize),
+    #[step(count = 4, child = crate::protocol::context::step::DzkpValidationProtocolStep)]
+    AggregateValidate(usize),
+}

ipa-core/src/protocol/ipa_prf/shuffle/malicious.rs

@@ -30,7 +30,7 @@ use crate::{
                 h1_shuffle_for_shard, h2_shuffle_for_shard, h3_shuffle_for_shard,
                 MaliciousShuffleable, ShuffleShare, Shuffleable,
             },
-            step::{OPRFShuffleStep, VerifyShuffleStep},
+            step::{OPRFShuffleStep, ShardedShuffleStep, VerifyShuffleStep},
             IntermediateShuffleMessages,
         },
         prss::SharedRandomness,
@@ -179,11 +179,11 @@ where
 
     // prepare keys
     let amount_of_keys: usize = (usize::try_from(S::Share::BITS).unwrap() + 31) / 32;
-    let keys = setup_keys(ctx.narrow(&OPRFShuffleStep::SetupKeys), amount_of_keys).await?;
+    let keys = setup_keys(ctx.narrow(&ShardedShuffleStep::SetupKeys), amount_of_keys).await?;
 
     // compute and append tags to rows
     let shares_and_tags: Vec<Pair<S::ShareAndTag>> =
-        compute_and_add_tags(ctx.narrow(&OPRFShuffleStep::GenerateTags), &keys, shares).await?;
+        compute_and_add_tags(ctx.narrow(&ShardedShuffleStep::GenerateTags), &keys, shares).await?;
 
     let (shuffled_shares, messages) = match ctx.role() {
         Role::H1 => h1_shuffle_for_shard(ctx.clone(), shares_and_tags).await,
@@ -193,7 +193,7 @@ where
 
     // verify the shuffle
     verify_shuffle::<_, S>(
-        ctx.narrow(&OPRFShuffleStep::VerifyShuffle),
+        ctx.narrow(&ShardedShuffleStep::VerifyShuffle),
         &keys,
         &shuffled_shares,
         messages,

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

@@ -9,7 +9,7 @@ use crate::{
         context::{Context, MaliciousContext, SemiHonestContext},
         ipa_prf::shuffle::sharded::ShuffleContext,
     },
-    sharding::{ShardBinding, Sharded},
+    sharding::Sharded,
 };
 
 mod base;
@@ -22,6 +22,8 @@ use malicious::{malicious_sharded_shuffle, malicious_shuffle};
 use sharded::shuffle as sharded_shuffle;
 pub use sharded::{MaliciousShuffleable, Shuffleable};
 
+use crate::sharding::NotSharded;
+
 /// This struct stores some intermediate messages during the shuffle.
 /// In a maliciously secure shuffle,
 /// these messages need to be checked for consistency across helpers.
@@ -63,7 +65,7 @@ pub trait Shuffle: Context {
         S: MaliciousShuffleable;
 }
 
-impl<T: ShardBinding> Shuffle for SemiHonestContext<'_, T> {
+impl Shuffle for SemiHonestContext<'_, NotSharded> {
     fn shuffle<S>(self, shares: Vec<S>) -> impl Future<Output = Result<Vec<S>, Error>> + Send
     where
         S: MaliciousShuffleable,
@@ -73,7 +75,7 @@ impl<T: ShardBinding> Shuffle for SemiHonestContext<'_, T> {
     }
 }
 
-impl<T: ShardBinding> Shuffle for MaliciousContext<'_, T> {
+impl Shuffle for MaliciousContext<'_, NotSharded> {
     fn shuffle<S>(self, shares: Vec<S>) -> impl Future<Output = Result<Vec<S>, Error>> + Send
     where
         S: MaliciousShuffleable,
@@ -84,7 +86,6 @@ impl<T: ShardBinding> Shuffle for MaliciousContext<'_, T> {
 
 /// Trait used by protocols to invoke either semi-honest or malicious sharded shuffle,
 /// depending on the type of context being used.
-#[allow(dead_code)]
 pub trait ShardedShuffle: ShuffleContext {
     fn sharded_shuffle<S>(
         self,

ipa-core/src/protocol/ipa_prf/shuffle/step.rs

@@ -29,6 +29,9 @@ pub(crate) enum VerifyShuffleStep {
 
 #[derive(CompactStep)]
 pub(crate) enum ShardedShuffleStep {
+    SetupKeys,
+    #[step(child = crate::protocol::boolean::step::EightBitStep)]
+    GenerateTags,
     /// Depending on the helper position inside the MPC ring, generate Ã, B̃ or both.
     PseudoRandomTable,
     /// Permute the input according to the PRSS shared between H1 and H2.
@@ -46,6 +49,8 @@ pub(crate) enum ShardedShuffleStep {
     TransferXY,
     /// H2 and H3 interaction - Exchange `C_1` and `C_2`.
     TransferC,
+    #[step(child = crate::protocol::ipa_prf::shuffle::step::VerifyShuffleStep)]
+    VerifyShuffle,
 }
 
 #[derive(CompactStep)]

ipa-core/src/query/runner/hybrid.rs

@@ -35,7 +35,7 @@ use crate::{
             oprf::{CONV_CHUNK, PRF_CHUNK},
             step::HybridStep,
         },
-        ipa_prf::{oprf_padding::PaddingParameters, prf_eval::PrfSharing, shuffle::Shuffle},
+        ipa_prf::{oprf_padding::PaddingParameters, prf_eval::PrfSharing, shuffle::ShardedShuffle},
         prss::{Endpoint, FromPrss},
         step::ProtocolStep::Hybrid,
         Gate,
@@ -73,7 +73,7 @@ impl<C, HV, R: PrivateKeyRegistry> Query<C, HV, R> {
 impl<C, HV, R> Query<C, HV, R>
 where
     C: UpgradableContext
-        + Shuffle
+        + ShardedShuffle
         + ShardedContext
         + FinalizerContext<FinalizingContext = DZKPUpgraded<C>>,
     HV: BooleanArray + U128Conversions,