Fix compact gate steps for sharded shuffle

akoshelev · Dec 26, 2024 · f4f4304 · f4f4304
1 parent bf79364
commit f4f4304
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/ipa-core/src/protocol/hybrid/step.rs b/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,

diff --git a/ipa-core/src/protocol/ipa_prf/aggregation/step.rs b/ipa-core/src/protocol/ipa_prf/aggregation/step.rs
@@ -7,7 +7,7 @@ pub(crate) enum AggregationStep {
     /// should be deleted
     #[step(child = crate::protocol::ipa_prf::oprf_padding::step::PaddingDpStep, name="padding_dp")]
     PaddingDp,
-    #[step(child = crate::protocol::ipa_prf::shuffle::step::OPRFShuffleStep)]
+    #[step(child = crate::protocol::ipa_prf::shuffle::step::ShardedShuffleStep)]
     Shuffle,
     Reveal,
     #[step(child = crate::protocol::context::step::DzkpValidationProtocolStep)]

diff --git a/ipa-core/src/protocol/ipa_prf/shuffle/malicious.rs b/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,

diff --git a/ipa-core/src/protocol/ipa_prf/shuffle/step.rs b/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)]