Fix compact gate steps for sharded shuffle

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

@@ -14,11 +14,9 @@ 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::ShardedShuffle,
         },

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::{Sharded},
+    sharding::Sharded,
 };
 
 mod base;
@@ -21,6 +21,7 @@ use base::shuffle_protocol as base_shuffle;
 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.

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)]