From 75161ad275d02f6add55c016fa81f5cc3b4c943d Mon Sep 17 00:00:00 2001
From: Alex Koshelev <koshelev@meta.com>
Date: Fri, 27 Sep 2024 10:59:49 -0700
Subject: [PATCH] Increase number of steps for DZKP proof to 600 from 192

Running Draft on 50M records revealed that we don't have sufficient number of steps for zero-knowledge proof runs.

I tried to keep the step limit under 20k, so I did some cleanup as well, but I think I failed and I had to bump it.
---
 ipa-core/src/protocol/context/step.rs             |  2 +-
 ipa-core/src/protocol/ipa_prf/aggregation/step.rs |  8 --------
 ipa-core/src/protocol/ipa_prf/mod.rs              |  2 +-
 ipa-core/src/protocol/step.rs                     | 10 ----------
 4 files changed, 2 insertions(+), 20 deletions(-)

diff --git a/ipa-core/src/protocol/context/step.rs b/ipa-core/src/protocol/context/step.rs
index d650340f8..aeb6bd76f 100644
--- a/ipa-core/src/protocol/context/step.rs
+++ b/ipa-core/src/protocol/context/step.rs
@@ -31,7 +31,7 @@ pub(crate) enum ValidateStep {
 // This really is only for DZKPs and not for MACs. The MAC protocol uses record IDs to
 // count batches. DZKP probably should do the same to avoid the fixed upper limit.
 #[derive(CompactStep)]
-#[step(count = 192, child = DzkpValidationProtocolStep)]
+#[step(count = 600, child = DzkpValidationProtocolStep)]
 pub(crate) struct DzkpBatchStep(pub usize);
 
 // This is used when we don't do batched verification, to avoid paying for x256 as many
diff --git a/ipa-core/src/protocol/ipa_prf/aggregation/step.rs b/ipa-core/src/protocol/ipa_prf/aggregation/step.rs
index fcb619a45..3c7b5da95 100644
--- a/ipa-core/src/protocol/ipa_prf/aggregation/step.rs
+++ b/ipa-core/src/protocol/ipa_prf/aggregation/step.rs
@@ -12,16 +12,8 @@ pub(crate) enum AggregationStep {
     RevealStep,
     #[step(child = AggregateChunkStep)]
     SumContributions,
-    #[step(child = crate::protocol::context::step::DzkpBatchStep)]
-    AggregateValidate,
 }
 
-/// the number of steps must be kept in sync with `MAX_BREAKDOWNS` defined
-/// [here](https://tinyurl.com/mwnbbnj6)
-#[derive(CompactStep)]
-#[step(count = 512, child = crate::protocol::boolean::step::EightBitStep, name = "b")]
-pub struct BucketStep(usize);
-
 #[derive(CompactStep)]
 #[step(count = 32, child = AggregateValuesStep, name = "depth")]
 pub(crate) struct AggregateChunkStep(usize);
diff --git a/ipa-core/src/protocol/ipa_prf/mod.rs b/ipa-core/src/protocol/ipa_prf/mod.rs
index 0462e31d6..cc3fa2633 100644
--- a/ipa-core/src/protocol/ipa_prf/mod.rs
+++ b/ipa-core/src/protocol/ipa_prf/mod.rs
@@ -744,7 +744,7 @@ mod compact_gate_tests {
     fn step_count_limit() {
         // This is an arbitrary limit intended to catch changes that unintentionally
         // blow up the step count. It can be increased, within reason.
-        const STEP_COUNT_LIMIT: u32 = 20_000;
+        const STEP_COUNT_LIMIT: u32 = 24_000;
         assert!(
             ProtocolStep::STEP_COUNT < STEP_COUNT_LIMIT,
             "Step count of {actual} exceeds limit of {STEP_COUNT_LIMIT}.",
diff --git a/ipa-core/src/protocol/step.rs b/ipa-core/src/protocol/step.rs
index c19bdb53b..8346557d2 100644
--- a/ipa-core/src/protocol/step.rs
+++ b/ipa-core/src/protocol/step.rs
@@ -31,16 +31,6 @@ impl<'de> serde::Deserialize<'de> for ProtocolGate {
 
 #[derive(CompactStep)]
 pub enum DeadCodeStep {
-    #[step(child = crate::protocol::basics::step::CheckZeroStep)]
-    CheckZero,
-    #[step(child = crate::protocol::basics::mul::step::MaliciousMultiplyStep)]
-    MaliciousMultiply,
-    #[step(child = crate::protocol::context::step::UpgradeStep)]
-    UpgradeShare,
-    #[step(child = crate::protocol::context::step::MaliciousProtocolStep)]
-    MaliciousProtocol,
-    #[step(child = crate::protocol::context::step::ValidateStep)]
-    MaliciousValidation,
     #[step(child = crate::protocol::ipa_prf::boolean_ops::step::SaturatedSubtractionStep)]
     SaturatedSubtraction,
     #[step(child = crate::protocol::ipa_prf::prf_sharding::step::FeatureLabelDotProductStep)]