feat: optimize pipeline for nodes_fold and reuse of vk [skip-line-limit]

crates/aggregator/src/domain/publickey_aggregation.rs

@@ -206,6 +206,12 @@ pub enum PublicKeyAggregatorState {
         dkg_aggregated_proof: Option<Proof>,
         c5_proof_pending: Option<Proof>,
         last_ec: Option<e3_events::EventContext<e3_events::Sequenced>>,
+        /// Accumulated nodes_fold proof after `nodes_fold_completed_slots` streaming steps.
+        nodes_fold_accumulator: Option<Proof>,
+        /// Number of slots folded so far; equals the next slot index to dispatch.
+        nodes_fold_completed_slots: u32,
+        /// Correlation ID of the in-flight [`ZkRequest::NodesFoldStep`], if any.
+        nodes_fold_step_correlation: Option<e3_events::CorrelationId>,
     },
     Complete {
         public_key: ArcBytes,

crates/events/src/enclave_event/compute_request/mod.rs

@@ -92,6 +92,7 @@ impl ToString for ComputeRequest {
                 ZkRequest::ThresholdShareDecryption(_) => "ZkThresholdShareDecryption",
                 ZkRequest::DecryptedSharesAggregation(_) => "ZkDecryptedSharesAggregation",
                 ZkRequest::NodeDkgFold(_) => "ZkNodeDkgFold",
+                ZkRequest::NodesFoldStep(_) => "ZkNodesFoldStep",
                 ZkRequest::DkgAggregation(_) => "ZkDkgAggregation",
                 ZkRequest::DecryptionAggregation(_) => "ZkDecryptionAggregation",
             },

crates/events/src/enclave_event/compute_request/zk.rs

@@ -38,6 +38,8 @@ pub enum ZkRequest {
     DecryptedSharesAggregation(DecryptedSharesAggregationProofRequest),
     /// Per-node DKG recursive fold (C2abFold → … → NodeFold).
     NodeDkgFold(NodeDkgFoldRequest),
+    /// Single step of the streaming cross-node nodes_fold accumulation.
+    NodesFoldStep(NodesFoldStepRequest),
     /// Cross-node DKG aggregator (NodesFold + C5 + DkgAggregator).
     DkgAggregation(DkgAggregationRequest),
     /// Phase-7 decryption aggregator (C6Fold + C7 + DecryptionAggregator).
@@ -71,10 +73,30 @@ pub struct NodeDkgFoldRequest {
     pub committee_size: CiphernodesCommitteeSize,
 }
 
+/// Single step of the streaming cross-node nodes_fold accumulation.
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub struct NodesFoldStepRequest {
+    /// The `node_fold` proof for this slot.
+    pub inner_proof: Proof,
+    /// The prior accumulator proof, or `None` for the first step.
+    pub prior_accumulator: Option<Proof>,
+    /// Slot index for this honest party (position in ascending `party_id` order).
+    pub slot_index: u32,
+    /// Total honest-party count H.
+    pub total_slots: usize,
+    /// E3 identifier used for job namespacing.
+    pub e3_id: String,
+    pub params_preset: BfvPreset,
+    pub committee_size: CiphernodesCommitteeSize,
+}
+
 /// Cross-node DKG aggregation (NodesFold + C5 + DkgAggregator).
 #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct DkgAggregationRequest {
     pub node_fold_proofs: Vec<Proof>,
+    /// Pre-computed nodes_fold accumulator proof. When present the prover skips the
+    /// sequential fold and uses this directly as input to the DkgAggregator circuit.
+    pub nodes_fold_proof: Option<Proof>,
     pub c5_proof: Proof,
     pub party_ids: Vec<u64>,
     /// Ordered committee addresses (`topNodes`) for `committee_hash_*` public inputs.
@@ -293,6 +315,8 @@ pub enum ZkResponse {
     DecryptedSharesAggregation(DecryptedSharesAggregationProofResponse),
     /// Output of [`ZkRequest::NodeDkgFold`].
     NodeDkgFold(NodeDkgFoldResponse),
+    /// Output of [`ZkRequest::NodesFoldStep`].
+    NodesFoldStep(NodesFoldStepResponse),
     /// Output of [`ZkRequest::DkgAggregation`].
     DkgAggregation(DkgAggregationResponse),
     /// Output of [`ZkRequest::DecryptionAggregation`].
@@ -305,6 +329,12 @@ pub struct NodeDkgFoldResponse {
     pub proof: Proof,
 }
 
+/// Response from [`ZkRequest::NodesFoldStep`].
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub struct NodesFoldStepResponse {
+    pub accumulator_proof: Proof,
+}
+
 /// Response from [`ZkRequest::DkgAggregation`].
 #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct DkgAggregationResponse {

crates/multithread/src/multithread.rs

@@ -29,10 +29,10 @@ use e3_events::{
     DecryptionAggregationResponse, DkgAggregationRequest, DkgAggregationResponse,
     DkgShareDecryptionProofRequest, DkgShareDecryptionProofResponse, EnclaveEvent,
     EnclaveEventData, EventPublisher, EventSubscriber, EventType, NodeDkgFoldRequest,
-    NodeDkgFoldResponse, PartyVerificationResult, PkAggregationProofRequest,
-    PkAggregationProofResponse, PkBfvProofRequest, PkBfvProofResponse, PkGenerationProofRequest,
-    PkGenerationProofResponse, Proof, ShareComputationProofRequest, ShareComputationProofResponse,
-    ShareEncryptionProofRequest, ShareEncryptionProofResponse,
+    NodeDkgFoldResponse, NodesFoldStepRequest, NodesFoldStepResponse, PartyVerificationResult,
+    PkAggregationProofRequest, PkAggregationProofResponse, PkBfvProofRequest, PkBfvProofResponse,
+    PkGenerationProofRequest, PkGenerationProofResponse, Proof, ShareComputationProofRequest,
+    ShareComputationProofResponse, ShareEncryptionProofRequest, ShareEncryptionProofResponse,
     ThresholdShareDecryptionProofRequest, ThresholdShareDecryptionProofResponse, TypedEvent,
     VerifyShareDecryptionProofsRequest, VerifyShareDecryptionProofsResponse,
     VerifyShareProofsRequest, VerifyShareProofsResponse, ZkError as ZkEventError, ZkRequest,
@@ -72,9 +72,9 @@ use e3_zk_helpers::threshold::pk_aggregation::PkAggregationCircuitData;
 use e3_zk_helpers::CiphernodesCommittee;
 use e3_zk_helpers::CiphernodesCommitteeSize;
 use e3_zk_prover::{
-    prove_decryption_aggregation_jobs, prove_dkg_aggregation, prove_node_dkg_fold, CircuitVariant,
-    DecryptionAggregationJob, DkgAggregationInput, NodeDkgFoldInput, NodeDkgFoldProveResult,
-    Provable, ZkBackend, ZkError, ZkProver,
+    generate_nodes_fold_step, prove_decryption_aggregation_jobs, prove_dkg_aggregation,
+    prove_node_dkg_fold, CircuitVariant, DecryptionAggregationJob, DkgAggregationInput,
+    NodeDkgFoldInput, NodeDkgFoldProveResult, Provable, ZkBackend, ZkError, ZkProver,
 };
 use fhe::bfv::{Ciphertext, Encoding, Plaintext, PublicKey, SecretKey};
 use fhe::mbfv::PublicKeyShare;
@@ -693,6 +693,9 @@ fn handle_zk_request(
         ZkRequest::NodeDkgFold(req) => timefunc("zk_node_dkg_fold", id, || {
             handle_node_dkg_fold_proof(&prover, req, request.clone(), report.clone())
         }),
+        ZkRequest::NodesFoldStep(req) => timefunc("zk_nodes_fold_step", id, || {
+            handle_nodes_fold_step_proof(&prover, req, request.clone())
+        }),
         ZkRequest::DkgAggregation(req) => timefunc("zk_dkg_aggregation", id, || {
             handle_dkg_aggregation_proof(&prover, req, request.clone())
         }),
@@ -749,6 +752,35 @@ fn handle_node_dkg_fold_proof(
     ))
 }
 
+fn handle_nodes_fold_step_proof(
+    prover: &ZkProver,
+    req: NodesFoldStepRequest,
+    request: ComputeRequest,
+) -> Result<ComputeResponse, ComputeRequestError> {
+    let artifacts_dir =
+        prover.resolve_artifacts_dir(req.params_preset, req.committee_size.as_str());
+    let accumulator_proof = generate_nodes_fold_step(
+        prover,
+        &req.inner_proof,
+        req.prior_accumulator.as_ref(),
+        req.slot_index,
+        req.total_slots,
+        &format!("{}-nodesfold-step-{}", req.e3_id, req.slot_index),
+        artifacts_dir.as_str(),
+    )
+    .map_err(|e| {
+        ComputeRequestError::new(
+            ComputeRequestErrorKind::Zk(ZkEventError::ProofGenerationFailed(e.to_string())),
+            request.clone(),
+        )
+    })?;
+    Ok(ComputeResponse::zk(
+        ZkResponse::NodesFoldStep(NodesFoldStepResponse { accumulator_proof }),
+        request.correlation_id,
+        request.e3_id,
+    ))
+}
+
 fn handle_dkg_aggregation_proof(
     prover: &ZkProver,
     req: DkgAggregationRequest,
@@ -757,6 +789,7 @@ fn handle_dkg_aggregation_proof(
     let job_id = zk_bb_work_id(&request);
     let input = DkgAggregationInput {
         node_fold_proofs: &req.node_fold_proofs,
+        nodes_fold_proof: req.nodes_fold_proof.as_ref(),
         c5_proof: &req.c5_proof,
         party_ids: &req.party_ids,
         committee_addresses: &req.committee_addresses,

crates/zk-prover/src/circuits/aggregation/c3_accumulator.rs

@@ -32,6 +32,31 @@ fn c3_fold_public_input_field_count(total_slots: usize) -> usize {
 const C3_FOLD_PREFIX_LEN: usize = 4;
 const C3_FOLD_SLOT_WIDTH: usize = 3;
 
+struct C3FoldVks {
+    inner_vk: vk::VkArtifacts,
+    fold_vk: vk::VkArtifacts,
+    kernel_vk: vk::VkArtifacts,
+}
+
+impl C3FoldVks {
+    fn load(prover: &ZkProver, artifacts_dir: &str) -> Result<Self, ZkError> {
+        Ok(Self {
+            inner_vk: vk::load_vk_artifacts(
+                &prover.circuits_dir(CircuitVariant::Recursive, artifacts_dir),
+                CircuitName::ShareEncryption,
+            )?,
+            fold_vk: vk::load_vk_artifacts(
+                &prover.circuits_dir(CircuitVariant::Default, artifacts_dir),
+                CircuitName::C3Fold,
+            )?,
+            kernel_vk: vk::load_vk_artifacts(
+                &prover.circuits_dir(CircuitVariant::Default, artifacts_dir),
+                CircuitName::C3FoldKernel,
+            )?,
+        })
+    }
+}
+
 /// Proves [`CircuitName::C3FoldKernel`] for the same `inner` / `total_slots` as the fold step.
 ///
 /// Uses work dir `job_id` (caller should use a suffix of the fold `e3_id` so jobs stay distinct).
@@ -130,38 +155,22 @@ fn parse_c3_fold_public_field_strings(proof: &Proof) -> Result<Vec<String>, ZkEr
     )
 }
 
-/// One sequential `c3_fold` step.
-///
-/// `prior_fold` is `None` on the first step and `Some` on all subsequent steps.
-/// `total_slots` is `N_PARTIES * L_THRESHOLD` — one slot per (party, threshold-modulus) pair.
-/// On the first step this sets the accumulator size; on subsequent steps it is cross-checked
-/// against the slot count already encoded in `prior_fold`.
-///
-/// Used only by [`generate_sequential_c3_fold`]; callers should use that entry point.
-fn generate_c3_fold_step(
+fn generate_c3_fold_step_with_vks(
     prover: &ZkProver,
     inner: &Proof,
     prior_fold: Option<&Proof>,
     slot_index: u32,
     total_slots: usize,
     e3_id: &str,
     artifacts_dir: &str,
+    vks: &C3FoldVks,
 ) -> Result<Proof, ZkError> {
     let is_first_step = prior_fold.is_none();
 
-    let inner_vk = vk::load_vk_artifacts(
-        &prover.circuits_dir(CircuitVariant::Recursive, artifacts_dir),
-        CircuitName::ShareEncryption,
-    )?;
     let c3_public_inputs = share_encryption_inner_public_inputs(inner)?;
-
     let expected_acc_pub = c3_fold_public_input_field_count(total_slots);
 
-    let (acc_vk_art, acc_proof, acc_public_inputs) = if is_first_step {
-        let kernel_vk = vk::load_vk_artifacts(
-            &prover.circuits_dir(CircuitVariant::Default, artifacts_dir),
-            CircuitName::C3FoldKernel,
-        )?;
+    let (acc_vk_fields, acc_vk_hash, acc_proof, acc_public_inputs) = if is_first_step {
         let kernel_job_id = format!("{e3_id}-c3fold-kernel");
         let kernel_proof = generate_c3_fold_kernel_genesis_proof(
             prover,
@@ -180,13 +189,13 @@ fn generate_c3_fold_step(
             )));
         }
         (
-            kernel_vk,
+            vks.kernel_vk.verification_key.clone(),
+            vks.kernel_vk.key_hash.clone(),
             bytes_to_field_strings(&kernel_proof.data)?,
             acc_pi,
         )
     } else {
         let p = prior_fold.expect("prior_fold required when is_first_step is false");
-        // Parse once; derive slot count from field count to avoid a second parse.
         let acc_pi = parse_c3_fold_public_field_strings(p)?;
         let prior_slots = (acc_pi.len() - 4) / 3;
         if prior_slots == 0 {
@@ -209,24 +218,22 @@ fn generate_c3_fold_step(
             )));
         }
         (
-            vk::load_vk_artifacts(
-                &prover.circuits_dir(CircuitVariant::Default, artifacts_dir),
-                CircuitName::C3Fold,
-            )?,
+            vks.fold_vk.verification_key.clone(),
+            vks.fold_vk.key_hash.clone(),
             bytes_to_field_strings(&p.data)?,
             acc_pi,
         )
     };
 
     let full_input = C3FoldStepInput {
-        inner_vk: inner_vk.verification_key,
+        inner_vk: vks.inner_vk.verification_key.clone(),
         inner_proof: bytes_to_field_strings(&inner.data)?,
         c3_public_inputs,
-        acc_vk: acc_vk_art.verification_key,
+        acc_vk: acc_vk_fields,
         acc_proof,
         acc_public_inputs,
-        inner_key_hash: inner_vk.key_hash,
-        acc_key_hash: acc_vk_art.key_hash,
+        inner_key_hash: vks.inner_vk.key_hash.clone(),
+        acc_key_hash: acc_vk_hash,
         is_first_step,
         slot_index,
     };
@@ -291,19 +298,21 @@ pub fn generate_sequential_c3_fold(
         }
         seen[idx] = true;
     }
+    let vks = C3FoldVks::load(prover, artifacts_dir)?;
     sequential_fold(
         "generate_sequential_c3_fold",
         inner_proofs,
         slot_indices,
         |inner, prior, slot| {
-            generate_c3_fold_step(
+            generate_c3_fold_step_with_vks(
                 prover,
                 inner,
                 prior,
                 slot,
                 total_slots,
                 e3_id,
                 artifacts_dir,
+                &vks,
             )
         },
     )