fix(prover): fix degrees logic in prover new (#1512)

scroll-tech · Aug 29, 2024 · ae8c858 · ae8c858
1 parent 2ee1c89
commit ae8c858
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 51 deletions.
diff --git a/prover/src/zk_circuits_handler/common.rs b/prover/src/zk_circuits_handler/common.rs
@@ -1,12 +1,14 @@
 use std::{collections::BTreeMap, rc::Rc};
 
+use crate::types::ProverType;
+
 use once_cell::sync::OnceCell;
 
 use halo2_proofs::{halo2curves::bn256::Bn256, poly::kzg::commitment::ParamsKZG};
 
 static mut PARAMS_MAP: OnceCell<Rc<BTreeMap<u32, ParamsKZG<Bn256>>>> = OnceCell::new();
 
-pub fn get_params_map<'a, F>(load_params_func: F) -> &'a BTreeMap<u32, ParamsKZG<Bn256>>
+pub fn get_params_map_instance<'a, F>(load_params_func: F) -> &'a BTreeMap<u32, ParamsKZG<Bn256>>
 where
     F: FnOnce() -> BTreeMap<u32, ParamsKZG<Bn256>>,
 {
@@ -17,3 +19,15 @@ where
         })
     }
 }
+
+pub fn get_degrees<F>(prover_types: &std::collections::HashSet<ProverType>, f: F) -> Vec<u32>
+where
+    F: FnMut(&ProverType) -> Vec<u32>,
+{
+    prover_types
+        .iter()
+        .flat_map(f)
+        .collect::<std::collections::HashSet<u32>>()
+        .into_iter()
+        .collect()
+}
diff --git a/prover/src/zk_circuits_handler/darwin.rs b/prover/src/zk_circuits_handler/darwin.rs
@@ -1,4 +1,4 @@
-use super::CircuitsHandler;
+use super::{common::*, CircuitsHandler};
 use crate::{
     geth_client::GethClient,
     types::{ProverType, TaskType},
@@ -11,9 +11,13 @@ use crate::types::{CommonHash, Task};
 use std::{cell::RefCell, cmp::Ordering, env, rc::Rc};
 
 use prover_darwin::{
-    aggregator::Prover as BatchProver, check_chunk_hashes, common::Prover as CommonProver,
-    zkevm::Prover as ChunkProver, BatchProof, BatchProvingTask, BlockTrace, BundleProof,
-    BundleProvingTask, ChunkInfo, ChunkProof, ChunkProvingTask,
+    aggregator::Prover as BatchProver,
+    check_chunk_hashes,
+    common::Prover as CommonProver,
+    config::{AGG_DEGREES, ZKEVM_DEGREES},
+    zkevm::Prover as ChunkProver,
+    BatchProof, BatchProvingTask, BlockTrace, BundleProof, BundleProvingTask, ChunkInfo,
+    ChunkProof, ChunkProvingTask,
 };
 
 // Only used for debugging.
@@ -61,35 +65,27 @@ impl DarwinHandler {
             chunk_prover: None,
             geth_client,
         };
+        let degrees: Vec<u32> = get_degrees(&prover_types_set, |prover_type| match prover_type {
+            ProverType::Chunk => ZKEVM_DEGREES.clone(),
+            ProverType::Batch => AGG_DEGREES.clone(),
+        });
+        let params_map = get_params_map_instance(|| {
+            log::info!(
+                "calling get_params_map from {}, prover_types: {:?}, degrees: {:?}",
+                class_name,
+                prover_types_set,
+                degrees
+            );
+            CommonProver::load_params_map(params_dir, &degrees)
+        });
         for prover_type in prover_types_set {
             match prover_type {
                 ProverType::Chunk => {
-                    let degrees = prover_darwin::config::ZKEVM_DEGREES.clone();
-                    let params_map = super::common::get_params_map(|| {
-                        log::info!(
-                            "calling get_params_map from {}, prover_type: {:?}, degrees: {:?}",
-                            class_name,
-                            prover_type,
-                            degrees
-                        );
-                        CommonProver::load_params_map(params_dir, &degrees)
-                    });
                     handler.chunk_prover = Some(RefCell::new(ChunkProver::from_params_and_assets(
                         params_map, assets_dir,
                     )));
                 }
-
                 ProverType::Batch => {
-                    let degrees = prover_darwin::config::AGG_DEGREES.clone();
-                    let params_map = super::common::get_params_map(|| {
-                        log::info!(
-                            "calling get_params_map from {}, prover_type: {:?}, degrees: {:?}",
-                            class_name,
-                            prover_type,
-                            degrees
-                        );
-                        CommonProver::load_params_map(params_dir, &degrees)
-                    });
                     handler.batch_prover = Some(RefCell::new(BatchProver::from_params_and_assets(
                         params_map, assets_dir,
                     )))

diff --git a/prover/src/zk_circuits_handler/darwin_v2.rs b/prover/src/zk_circuits_handler/darwin_v2.rs
@@ -1,4 +1,4 @@
-use super::CircuitsHandler;
+use super::{common::*, CircuitsHandler};
 use crate::{
     geth_client::GethClient,
     types::{ProverType, TaskType},
@@ -11,9 +11,13 @@ use crate::types::{CommonHash, Task};
 use std::{cell::RefCell, cmp::Ordering, env, rc::Rc};
 
 use prover_darwin_v2::{
-    aggregator::Prover as BatchProver, check_chunk_hashes, common::Prover as CommonProver,
-    zkevm::Prover as ChunkProver, BatchProof, BatchProvingTask, BlockTrace, BundleProof,
-    BundleProvingTask, ChunkInfo, ChunkProof, ChunkProvingTask,
+    aggregator::Prover as BatchProver,
+    check_chunk_hashes,
+    common::Prover as CommonProver,
+    config::{AGG_DEGREES, ZKEVM_DEGREES},
+    zkevm::Prover as ChunkProver,
+    BatchProof, BatchProvingTask, BlockTrace, BundleProof, BundleProvingTask, ChunkInfo,
+    ChunkProof, ChunkProvingTask,
 };
 
 // Only used for debugging.
@@ -61,35 +65,27 @@ impl DarwinV2Handler {
             chunk_prover: None,
             geth_client,
         };
+        let degrees: Vec<u32> = get_degrees(&prover_types_set, |prover_type| match prover_type {
+            ProverType::Chunk => ZKEVM_DEGREES.clone(),
+            ProverType::Batch => AGG_DEGREES.clone(),
+        });
+        let params_map = get_params_map_instance(|| {
+            log::info!(
+                "calling get_params_map from {}, prover_types: {:?}, degrees: {:?}",
+                class_name,
+                prover_types_set,
+                degrees
+            );
+            CommonProver::load_params_map(params_dir, &degrees)
+        });
         for prover_type in prover_types_set {
             match prover_type {
                 ProverType::Chunk => {
-                    let degrees = prover_darwin::config::ZKEVM_DEGREES.clone();
-                    let params_map = super::common::get_params_map(|| {
-                        log::info!(
-                            "calling get_params_map from {}, prover_type: {:?}, degrees: {:?}",
-                            class_name,
-                            prover_type,
-                            degrees
-                        );
-                        CommonProver::load_params_map(params_dir, &degrees)
-                    });
                     handler.chunk_prover = Some(RefCell::new(ChunkProver::from_params_and_assets(
                         params_map, assets_dir,
                     )));
                 }
-
                 ProverType::Batch => {
-                    let degrees = prover_darwin::config::AGG_DEGREES.clone();
-                    let params_map = super::common::get_params_map(|| {
-                        log::info!(
-                            "calling get_params_map from {}, prover_type: {:?}, degrees: {:?}",
-                            class_name,
-                            prover_type,
-                            degrees
-                        );
-                        CommonProver::load_params_map(params_dir, &degrees)
-                    });
                     handler.batch_prover = Some(RefCell::new(BatchProver::from_params_and_assets(
                         params_map, assets_dir,
                     )))