From 147161500b209e4dc05fca252cbf23f61c47a7a9 Mon Sep 17 00:00:00 2001 From: Marcin M <128217157+mm-zk@users.noreply.github.com> Date: Mon, 19 Feb 2024 11:18:16 +0100 Subject: [PATCH] feat(prover): adding keystore object to handle reading and writing of prover keys (#1132) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What ❔ * Adding keystore object as one interface to handle reading and writing of the verification and setup keys * As a side effect, also removing dependency on the old test_harness version from key generator ## Why ❔ * This will allow us to specify different (non default) directories for verification and setup keys in the future. --- prover/Cargo.lock | 3 +- prover/proof_fri_compressor/Cargo.toml | 3 + prover/proof_fri_compressor/src/compressor.rs | 27 +- .../src/gpu_prover_job_processor.rs | 17 +- prover/prover_fri/src/main.rs | 2 - prover/prover_fri/src/prover_job_processor.rs | 15 +- prover/prover_fri/src/socket_listener.rs | 21 +- prover/prover_fri/tests/basic_test.rs | 7 +- prover/prover_fri_types/src/lib.rs | 32 +- .../Cargo.toml | 1 - .../src/commitment_generator.rs | 7 +- .../src/commitment_utils.rs | 30 +- .../src/keystore.rs | 419 ++++++++++++++++++ .../src/lib.rs | 337 +------------- .../src/main.rs | 165 ++----- .../src/setup_data_generator.rs | 179 +++++--- .../src/tests.rs | 45 +- .../src/utils.rs | 45 +- .../witness_generator/src/leaf_aggregation.rs | 11 +- .../witness_generator/src/node_aggregation.rs | 19 +- prover/witness_generator/src/scheduler.rs | 16 +- .../witness_vector_generator/src/generator.rs | 6 +- 22 files changed, 756 insertions(+), 651 deletions(-) create mode 100644 prover/vk_setup_data_generator_server_fri/src/keystore.rs diff --git a/prover/Cargo.lock b/prover/Cargo.lock index c9f1185ca4a6..46981afb799c 100644 --- a/prover/Cargo.lock +++ b/prover/Cargo.lock @@ -6131,7 +6131,6 @@ dependencies = [ "tracing", "tracing-subscriber", "vlog", - "zkevm_test_harness 1.3.3", "zkevm_test_harness 1.4.1", "zksync_config", "zksync_env_config", @@ -7092,6 +7091,8 @@ dependencies = [ "futures 0.3.30", "prometheus_exporter", "reqwest", + "serde", + "serde_json", "structopt", "tokio", "tracing", diff --git a/prover/proof_fri_compressor/Cargo.toml b/prover/proof_fri_compressor/Cargo.toml index 465579d059c0..b1ec5965e0b8 100644 --- a/prover/proof_fri_compressor/Cargo.toml +++ b/prover/proof_fri_compressor/Cargo.toml @@ -33,3 +33,6 @@ ctrlc = { version = "3.1", features = ["termination"] } async-trait = "0.1" bincode = "1.0" reqwest = { version = "0.11", features = ["blocking"] } +serde_json = "1.0" +serde = { version = "1.0", features = ["derive"] } + diff --git a/prover/proof_fri_compressor/src/compressor.rs b/prover/proof_fri_compressor/src/compressor.rs index ecf26cacd4d1..521b184c2d6a 100644 --- a/prover/proof_fri_compressor/src/compressor.rs +++ b/prover/proof_fri_compressor/src/compressor.rs @@ -8,7 +8,10 @@ use zkevm_test_harness_1_3_3::{ abstract_zksync_circuit::concrete_circuits::{ ZkSyncCircuit, ZkSyncProof, ZkSyncVerificationKey, }, - bellman::{bn256::Bn256, plonk::better_better_cs::proof::Proof}, + bellman::{ + bn256::Bn256, + plonk::better_better_cs::{proof::Proof, setup::VerificationKey as SnarkVerificationKey}, + }, witness::oracle::VmWitnessOracle, }; use zksync_dal::ConnectionPool; @@ -26,7 +29,7 @@ use zksync_prover_fri_types::{ use zksync_prover_interface::outputs::L1BatchProofForL1; use zksync_queued_job_processor::JobProcessor; use zksync_types::L1BatchNumber; -use zksync_vk_setup_data_server_fri::{get_recursive_layer_vk_for_circuit_type, get_snark_vk}; +use zksync_vk_setup_data_server_fri::keystore::Keystore; use crate::metrics::METRICS; @@ -60,10 +63,12 @@ impl ProofCompressor { compression_mode: u8, verify_wrapper_proof: bool, ) -> anyhow::Result>>> { - let scheduler_vk = get_recursive_layer_vk_for_circuit_type( - ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8, - ) - .context("get_recursiver_layer_vk_for_circuit_type()")?; + let keystore = Keystore::default(); + let scheduler_vk = keystore + .load_recursive_layer_verification_key( + ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8, + ) + .context("get_recursiver_layer_vk_for_circuit_type()")?; let config = WrapperConfig::new(compression_mode); let (wrapper_proof, _) = wrap_proof(proof, scheduler_vk, config); @@ -76,7 +81,15 @@ impl ProofCompressor { bincode::deserialize(&serialized) .expect("Failed to deserialize proof with ZkSyncCircuit"); if verify_wrapper_proof { - let existing_vk = get_snark_vk().context("get_snark_vk()")?; + // We're fetching the key as String and deserializing it here + // as we don't want to include the old version of prover in the main libraries. + let existing_vk_serialized = keystore + .load_snark_verification_key() + .context("get_snark_vk()")?; + let existing_vk = serde_json::from_str::< + SnarkVerificationKey>>, + >(&existing_vk_serialized)?; + let vk = ZkSyncVerificationKey::from_verification_key_and_numeric_type(0, existing_vk); let scheduler_proof = ZkSyncProof::from_proof_and_numeric_type(0, proof.clone()); match vk.verify_proof(&scheduler_proof) { diff --git a/prover/prover_fri/src/gpu_prover_job_processor.rs b/prover/prover_fri/src/gpu_prover_job_processor.rs index 82b78024a98d..fc0030bfb6fc 100644 --- a/prover/prover_fri/src/gpu_prover_job_processor.rs +++ b/prover/prover_fri/src/gpu_prover_job_processor.rs @@ -27,9 +27,7 @@ pub mod gpu_prover { }; use zksync_queued_job_processor::{async_trait, JobProcessor}; use zksync_types::basic_fri_types::CircuitIdRoundTuple; - use zksync_vk_setup_data_server_fri::{ - get_setup_data_for_circuit_type, GoldilocksGpuProverSetupData, - }; + use zksync_vk_setup_data_server_fri::{keystore::Keystore, GoldilocksGpuProverSetupData}; use crate::{ metrics::METRICS, @@ -103,9 +101,10 @@ pub mod gpu_prover { .clone(), SetupLoadMode::FromDisk => { let started_at = Instant::now(); - let artifact: GoldilocksGpuProverSetupData = - get_setup_data_for_circuit_type(key.clone()) - .context("get_setup_data_for_circuit_type()")?; + let keystore = Keystore::default(); + let artifact: GoldilocksGpuProverSetupData = keystore + .load_gpu_setup_data_for_circuit_type(key.clone()) + .context("load_gpu_setup_data_for_circuit_type()")?; METRICS.gpu_setup_data_load_time[&key.circuit_id.to_string()] .observe(started_at.elapsed()); @@ -326,10 +325,12 @@ pub mod gpu_prover { &config.specialized_group_id, prover_setup_metadata_list ); + let keystore = Keystore::default(); for prover_setup_metadata in prover_setup_metadata_list { let key = setup_metadata_to_setup_data_key(&prover_setup_metadata); - let setup_data = get_setup_data_for_circuit_type(key.clone()) - .context("get_setup_data_for_circuit_type()")?; + let setup_data = keystore + .load_gpu_setup_data_for_circuit_type(key.clone()) + .context("load_gpu_setup_data_for_circuit_type()")?; cache.insert(key, Arc::new(setup_data)); } SetupLoadMode::FromMemory(cache) diff --git a/prover/prover_fri/src/main.rs b/prover/prover_fri/src/main.rs index d867fd5e93c9..e8c298c0663d 100644 --- a/prover/prover_fri/src/main.rs +++ b/prover/prover_fri/src/main.rs @@ -208,8 +208,6 @@ async fn get_prover_tasks( pool: ConnectionPool, circuit_ids_for_round_to_be_proven: Vec, ) -> anyhow::Result>>> { - use std::sync::Arc; - use gpu_prover_job_processor::gpu_prover; use socket_listener::gpu_socket_listener; use tokio::sync::Mutex; diff --git a/prover/prover_fri/src/prover_job_processor.rs b/prover/prover_fri/src/prover_job_processor.rs index 40275b681b70..77c32ce73f34 100644 --- a/prover/prover_fri/src/prover_job_processor.rs +++ b/prover/prover_fri/src/prover_job_processor.rs @@ -25,9 +25,7 @@ use zksync_prover_fri_types::{ use zksync_prover_fri_utils::fetch_next_circuit; use zksync_queued_job_processor::{async_trait, JobProcessor}; use zksync_types::{basic_fri_types::CircuitIdRoundTuple, protocol_version::L1VerifierConfig}; -use zksync_vk_setup_data_server_fri::{ - get_cpu_setup_data_for_circuit_type, GoldilocksProverSetupData, -}; +use zksync_vk_setup_data_server_fri::{keystore::Keystore, GoldilocksProverSetupData}; use crate::{ metrics::{CircuitLabels, Layer, METRICS}, @@ -88,9 +86,10 @@ impl Prover { .clone(), SetupLoadMode::FromDisk => { let started_at = Instant::now(); - let artifact: GoldilocksProverSetupData = - get_cpu_setup_data_for_circuit_type(key.clone()) - .context("get_cpu_setup_data_for_circuit_type()")?; + let keystore = Keystore::default(); + let artifact: GoldilocksProverSetupData = keystore + .load_cpu_setup_data_for_circuit_type(key.clone()) + .context("get_cpu_setup_data_for_circuit_type()")?; METRICS.gpu_setup_data_load_time[&key.circuit_id.to_string()] .observe(started_at.elapsed()); @@ -299,9 +298,11 @@ pub fn load_setup_data_cache(config: &FriProverConfig) -> anyhow::Result anyhow::Result { let started_at = Instant::now(); + let keystore = Keystore::default(); let cs = match circuit_wrapper { CircuitWrapper::Base(base_circuit) => { let key = ProverServiceDataKey::new( base_circuit.numeric_circuit_type(), AggregationRound::BasicCircuits, ); - let finalization_hint = - get_finalization_hints(key).context("get_finalization_hints()")?; + let finalization_hint = keystore + .load_finalization_hints(key) + .context("get_finalization_hints()")?; init_base_layer_cs_for_repeated_proving(base_circuit, &finalization_hint) } CircuitWrapper::Recursive(recursive_circuit) => { - let key = ProverServiceDataKey::new( - recursive_circuit.numeric_circuit_type(), - get_round_for_recursive_circuit_type(recursive_circuit.numeric_circuit_type()), - ); - let finalization_hint = - get_finalization_hints(key).context("get_finalization_hints()")?; + let key = + ProverServiceDataKey::new_recursive(recursive_circuit.numeric_circuit_type()); + let finalization_hint = keystore + .load_finalization_hints(key) + .context("get_finalization_hints()")?; init_recursive_layer_cs_for_repeated_proving(recursive_circuit, &finalization_hint) } }; diff --git a/prover/prover_fri/tests/basic_test.rs b/prover/prover_fri/tests/basic_test.rs index d8ce27b58507..d9c407b32efb 100644 --- a/prover/prover_fri/tests/basic_test.rs +++ b/prover/prover_fri/tests/basic_test.rs @@ -10,7 +10,9 @@ use zksync_prover_fri_types::{ keys::FriCircuitKey, CircuitWrapper, ProverJob, ProverServiceDataKey, }; use zksync_types::{basic_fri_types::AggregationRound, L1BatchNumber}; -use zksync_vk_setup_data_server_fri::generate_cpu_base_layer_setup_data; +use zksync_vk_setup_data_server_fri::{ + keystore::Keystore, setup_data_generator::generate_cpu_base_layer_setup_data, +}; fn compare_serialized(expected: &T, actual: &T) { let serialized_expected = bincode::serialize(expected).unwrap(); @@ -54,8 +56,9 @@ async fn prover_and_assert_base_layer( CircuitWrapper::Base(base) => base.clone(), CircuitWrapper::Recursive(_) => anyhow::bail!("Expected base layer circuit"), }; + let keystore = Keystore::default(); let setup_data = Arc::new( - generate_cpu_base_layer_setup_data(circuit) + generate_cpu_base_layer_setup_data(&keystore, circuit) .context("generate_cpu_base_layers_setup_data()")?, ); let setup_key = ProverServiceDataKey::new(circuit_id, aggregation_round); diff --git a/prover/prover_fri_types/src/lib.rs b/prover/prover_fri_types/src/lib.rs index a1572ee2a2c2..0fc979849f2f 100644 --- a/prover/prover_fri_types/src/lib.rs +++ b/prover/prover_fri_types/src/lib.rs @@ -6,7 +6,9 @@ use circuit_definitions::{ boojum::{cs::implementations::witness::WitnessVec, field::goldilocks::GoldilocksField}, circuit_definitions::{ base_layer::{ZkSyncBaseLayerCircuit, ZkSyncBaseLayerProof}, - recursion_layer::{ZkSyncRecursionLayerProof, ZkSyncRecursiveLayerCircuit}, + recursion_layer::{ + ZkSyncRecursionLayerProof, ZkSyncRecursionLayerStorageType, ZkSyncRecursiveLayerCircuit, + }, }, zkevm_circuits::scheduler::block_header::BlockAuxilaryOutputWitness, ZkSyncDefaultRoundFunction, @@ -112,10 +114,38 @@ pub struct ProverServiceDataKey { pub round: AggregationRound, } +fn get_round_for_recursive_circuit_type(circuit_type: u8) -> AggregationRound { + match circuit_type { + circuit_type if circuit_type == ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8 => { + AggregationRound::Scheduler + } + circuit_type if circuit_type == ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8 => { + AggregationRound::NodeAggregation + } + _ => AggregationRound::LeafAggregation, + } +} + impl ProverServiceDataKey { pub fn new(circuit_id: u8, round: AggregationRound) -> Self { Self { circuit_id, round } } + + /// Creates a new data key for recursive type - with auto selection of the aggregation round. + pub fn new_recursive(circuit_id: u8) -> Self { + Self { + circuit_id, + round: get_round_for_recursive_circuit_type(circuit_id), + } + } + + /// Data key for snark wrapper. + pub fn snark() -> Self { + Self { + circuit_id: 1, + round: AggregationRound::Scheduler, + } + } } #[derive(serde::Serialize, serde::Deserialize)] diff --git a/prover/vk_setup_data_generator_server_fri/Cargo.toml b/prover/vk_setup_data_generator_server_fri/Cargo.toml index 0e57dc8b7ad5..1f532cb28f3f 100644 --- a/prover/vk_setup_data_generator_server_fri/Cargo.toml +++ b/prover/vk_setup_data_generator_server_fri/Cargo.toml @@ -18,7 +18,6 @@ vlog = { path = "../../core/lib/vlog" } zksync_types = { path = "../../core/lib/types" } zksync_prover_fri_types = { path = "../prover_fri_types" } -zkevm_test_harness_1_3_3 = { git = "https://github.com/matter-labs/era-zkevm_test_harness.git", branch = "v1.3.3", package = "zkevm_test_harness" } zkevm_test_harness = { git = "https://github.com/matter-labs/era-zkevm_test_harness.git", branch = "v1.4.1" } circuit_definitions = { git = "https://github.com/matter-labs/era-zkevm_test_harness.git", branch = "v1.4.1", features = [ "log_tracing", diff --git a/prover/vk_setup_data_generator_server_fri/src/commitment_generator.rs b/prover/vk_setup_data_generator_server_fri/src/commitment_generator.rs index 7557a476ab70..73705e5f96c1 100644 --- a/prover/vk_setup_data_generator_server_fri/src/commitment_generator.rs +++ b/prover/vk_setup_data_generator_server_fri/src/commitment_generator.rs @@ -1,12 +1,13 @@ use anyhow::Context; use zksync_vk_setup_data_server_fri::{ commitment_utils::generate_commitments, + keystore::Keystore, vk_commitment_helper::{get_toml_formatted_value, read_contract_toml, write_contract_toml}, }; -pub fn read_and_update_contract_toml(dryrun: bool) -> anyhow::Result<()> { +pub fn read_and_update_contract_toml(keystore: &Keystore, dryrun: bool) -> anyhow::Result<()> { let mut contract_doc = read_contract_toml().context("read_contract_toml()")?; - let vk_commitments = generate_commitments().context("generate_commitments()")?; + let vk_commitments = generate_commitments(keystore).context("generate_commitments()")?; contract_doc["contracts"]["FRI_RECURSION_LEAF_LEVEL_VK_HASH"] = get_toml_formatted_value(vk_commitments.leaf); @@ -29,6 +30,6 @@ mod test { #[test] fn test_read_and_update_contract_toml() { - read_and_update_contract_toml(true).unwrap(); + read_and_update_contract_toml(&Keystore::default(), true).unwrap(); } } diff --git a/prover/vk_setup_data_generator_server_fri/src/commitment_utils.rs b/prover/vk_setup_data_generator_server_fri/src/commitment_utils.rs index 9a3bf5be3399..7900296b8f80 100644 --- a/prover/vk_setup_data_generator_server_fri/src/commitment_utils.rs +++ b/prover/vk_setup_data_generator_server_fri/src/commitment_utils.rs @@ -15,11 +15,11 @@ use zksync_types::{ H256, }; -use crate::{get_recursive_layer_vk_for_circuit_type, utils::get_leaf_vk_params}; +use crate::{keystore::Keystore, utils::get_leaf_vk_params}; lazy_static! { // TODO: do not initialize a static const with data read in runtime. - static ref COMMITMENTS: Lazy = Lazy::new(|| { circuit_commitments().unwrap() }); + static ref COMMITMENTS: Lazy = Lazy::new(|| { circuit_commitments(&Keystore::default()).unwrap() }); } #[derive(Debug)] @@ -29,8 +29,8 @@ pub struct VkCommitments { pub scheduler: String, } -fn circuit_commitments() -> anyhow::Result { - let commitments = generate_commitments().context("generate_commitments()")?; +fn circuit_commitments(keystore: &Keystore) -> anyhow::Result { + let commitments = generate_commitments(keystore).context("generate_commitments()")?; let snark_wrapper_vk = std::env::var("CONTRACTS_SNARK_WRAPPER_VK_HASH") .context("SNARK wrapper VK not found in the config")?; Ok(L1VerifierConfig { @@ -55,8 +55,8 @@ fn circuit_commitments() -> anyhow::Result { }) } -pub fn generate_commitments() -> anyhow::Result { - let leaf_vk_params = get_leaf_vk_params().context("get_leaf_vk_params()")?; +pub fn generate_commitments(keystore: &Keystore) -> anyhow::Result { + let leaf_vk_params = get_leaf_vk_params(keystore).context("get_leaf_vk_params()")?; let leaf_layer_params = leaf_vk_params .iter() .map(|el| el.1.clone()) @@ -65,16 +65,18 @@ pub fn generate_commitments() -> anyhow::Result { .unwrap(); let leaf_vk_commitment = compute_leaf_vks_and_params_commitment(leaf_layer_params); - let node_vk = get_recursive_layer_vk_for_circuit_type( - ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, - ) - .context("get_recursive_layer_vk_for_circuit_type(NodeLayerCircuit)")?; + let node_vk = keystore + .load_recursive_layer_verification_key( + ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, + ) + .context("get_recursive_layer_vk_for_circuit_type(NodeLayerCircuit)")?; let node_vk_commitment = compute_node_vk_commitment(node_vk.clone()); - let scheduler_vk = get_recursive_layer_vk_for_circuit_type( - ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8, - ) - .context("get_recursive_layer_vk_for_circuit_type(SchedulerCircuit)")?; + let scheduler_vk = keystore + .load_recursive_layer_verification_key( + ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8, + ) + .context("get_recursive_layer_vk_for_circuit_type(SchedulerCircuit)")?; let scheduler_vk_commitment = compute_node_vk_commitment(scheduler_vk.clone()); let hex_concatenator = |hex_array: [GoldilocksField; 4]| { diff --git a/prover/vk_setup_data_generator_server_fri/src/keystore.rs b/prover/vk_setup_data_generator_server_fri/src/keystore.rs new file mode 100644 index 000000000000..a684d2c121e0 --- /dev/null +++ b/prover/vk_setup_data_generator_server_fri/src/keystore.rs @@ -0,0 +1,419 @@ +use std::{fs, fs::File, io::Read}; + +use anyhow::Context as _; +use circuit_definitions::{ + boojum::cs::implementations::setup::FinalizationHintsForProver, + circuit_definitions::{ + aux_layer::ZkSyncSnarkWrapperVK, + base_layer::ZkSyncBaseLayerVerificationKey, + recursion_layer::{ZkSyncRecursionLayerStorageType, ZkSyncRecursionLayerVerificationKey}, + }, + zkevm_circuits::scheduler::aux::BaseLayerCircuitType, +}; +use serde::{Deserialize, Serialize}; +use zkevm_test_harness::data_source::{in_memory_data_source::InMemoryDataSource, SetupDataSource}; +use zksync_config::configs::FriProverConfig; +use zksync_env_config::FromEnv; +use zksync_prover_fri_types::ProverServiceDataKey; +use zksync_types::basic_fri_types::AggregationRound; + +#[cfg(feature = "gpu")] +use crate::GoldilocksGpuProverSetupData; +use crate::GoldilocksProverSetupData; + +pub enum ProverServiceDataType { + VerificationKey, + SetupData, + FinalizationHints, + SnarkVerificationKey, +} + +/// Key store manages all the prover keys. +/// There are 2 types: +/// - small verification, finalization keys (used only during verification) +/// - large setup keys, used during proving. +pub struct Keystore { + /// Directory to store all the small keys. + basedir: String, + /// Directory to store large setup keys. + setup_data_path: Option, +} + +fn get_base_path_from_env() -> String { + let zksync_home = std::env::var("ZKSYNC_HOME").unwrap_or_else(|_| "/".into()); + format!( + "{}/prover/vk_setup_data_generator_server_fri/data", + zksync_home + ) +} + +impl Default for Keystore { + fn default() -> Self { + Self { + basedir: get_base_path_from_env(), + setup_data_path: Some( + FriProverConfig::from_env() + .expect("FriProverConfig::from_env()") + .setup_data_path, + ), + } + } +} + +impl Keystore { + /// Base-dir is the location of smaller keys (like verification keys and finalization hints). + /// Setup data path is used for the large setup keys. + pub fn new(basedir: String, setup_data_path: String) -> Self { + Keystore { + basedir, + setup_data_path: Some(setup_data_path), + } + } + pub fn new_with_optional_setup_path(basedir: String, setup_data_path: Option) -> Self { + Keystore { + basedir, + setup_data_path, + } + } + + pub fn get_base_path(&self) -> &str { + &self.basedir + } + + fn get_file_path( + &self, + key: ProverServiceDataKey, + service_data_type: ProverServiceDataType, + ) -> String { + let name = match key.round { + AggregationRound::BasicCircuits => { + format!("basic_{}", key.circuit_id) + } + AggregationRound::LeafAggregation => { + format!("leaf_{}", key.circuit_id) + } + AggregationRound::NodeAggregation => "node".to_string(), + AggregationRound::Scheduler => "scheduler".to_string(), + }; + match service_data_type { + ProverServiceDataType::VerificationKey => { + format!("{}/verification_{}_key.json", self.basedir, name) + } + ProverServiceDataType::SetupData => { + format!( + "{}/setup_{}_data.bin", + self.setup_data_path + .as_ref() + .expect("Setup data path not set"), + name + ) + } + ProverServiceDataType::FinalizationHints => { + format!("{}/finalization_hints_{}.bin", self.basedir, name) + } + ProverServiceDataType::SnarkVerificationKey => { + format!("{}/snark_verification_{}_key.json", self.basedir, name) + } + } + } + + fn load_json_from_file Deserialize<'a>>(filepath: String) -> anyhow::Result { + let text = std::fs::read_to_string(&filepath) + .with_context(|| format!("Failed reading verification key from path: {filepath}"))?; + serde_json::from_str::(&text) + .with_context(|| format!("Failed deserializing verification key from path: {filepath}")) + } + fn save_json_pretty(filepath: String, data: &T) -> anyhow::Result<()> { + std::fs::write(&filepath, serde_json::to_string_pretty(data).unwrap()) + .with_context(|| format!("writing to '{filepath}' failed")) + } + + fn load_bincode_from_file Deserialize<'a>>(filepath: String) -> anyhow::Result { + let mut file = File::open(filepath.clone()) + .with_context(|| format!("Failed reading setup-data from path: {filepath:?}"))?; + let mut buffer = Vec::new(); + file.read_to_end(&mut buffer).with_context(|| { + format!("Failed reading setup-data to buffer from path: {filepath:?}") + })?; + bincode::deserialize::(&buffer) + .with_context(|| format!("Failed deserializing setup-data at path: {filepath:?}")) + } + + /// + /// Verification keys + /// + + pub fn load_base_layer_verification_key( + &self, + circuit_type: u8, + ) -> anyhow::Result { + Self::load_json_from_file(self.get_file_path( + ProverServiceDataKey::new(circuit_type, AggregationRound::BasicCircuits), + ProverServiceDataType::VerificationKey, + )) + } + + pub fn load_recursive_layer_verification_key( + &self, + circuit_type: u8, + ) -> anyhow::Result { + Self::load_json_from_file(self.get_file_path( + ProverServiceDataKey::new_recursive(circuit_type), + ProverServiceDataType::VerificationKey, + )) + } + + pub fn save_base_layer_verification_key( + &self, + vk: ZkSyncBaseLayerVerificationKey, + ) -> anyhow::Result<()> { + let filepath = self.get_file_path( + ProverServiceDataKey::new(vk.numeric_circuit_type(), AggregationRound::BasicCircuits), + ProverServiceDataType::VerificationKey, + ); + tracing::info!("saving basic verification key to: {}", filepath); + Self::save_json_pretty(filepath, &vk) + } + + pub fn save_recursive_layer_verification_key( + &self, + vk: ZkSyncRecursionLayerVerificationKey, + ) -> anyhow::Result<()> { + let filepath = self.get_file_path( + ProverServiceDataKey::new_recursive(vk.numeric_circuit_type()), + ProverServiceDataType::VerificationKey, + ); + tracing::info!("saving recursive layer verification key to: {}", filepath); + Self::save_json_pretty(filepath, &vk) + } + + /// + /// Finalization hints + /// + + pub fn save_finalization_hints( + &self, + key: ProverServiceDataKey, + hint: &FinalizationHintsForProver, + ) -> anyhow::Result<()> { + let filepath = self.get_file_path(key.clone(), ProverServiceDataType::FinalizationHints); + + tracing::info!("saving finalization hints for {:?} to: {}", key, filepath); + let serialized = + bincode::serialize(&hint).context("Failed to serialize finalization hints")?; + fs::write(filepath, serialized).context("Failed to write finalization hints to file") + } + + pub fn load_finalization_hints( + &self, + key: ProverServiceDataKey, + ) -> anyhow::Result { + let mut key = key; + // For `NodeAggregation` round we have only 1 finalization hints for all circuit type. + // TODO: is this needed?? + if key.round == AggregationRound::NodeAggregation { + key.circuit_id = ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8; + } + Self::load_bincode_from_file( + self.get_file_path(key, ProverServiceDataType::FinalizationHints), + ) + } + + /// + /// Snark wrapper + /// + + /// Loads snark verification key + // For snark wrapper - we're actually returning a raw serialized string, and the parsing happens + // on the reader's side (in proof compressor). + // This way, we avoid including the old 1.3.3 test harness to our main library. + pub fn load_snark_verification_key(&self) -> anyhow::Result { + let filepath = self.get_file_path( + ProverServiceDataKey::snark(), + ProverServiceDataType::SnarkVerificationKey, + ); + std::fs::read_to_string(&filepath) + .with_context(|| format!("Failed reading Snark verification key from path: {filepath}")) + } + + pub fn save_snark_verification_key(&self, vk: ZkSyncSnarkWrapperVK) -> anyhow::Result<()> { + let filepath = self.get_file_path( + ProverServiceDataKey::snark(), + ProverServiceDataType::SnarkVerificationKey, + ); + tracing::info!("saving snark verification key to: {}", filepath); + Self::save_json_pretty(filepath, &vk.into_inner()) + } + + /// + /// Setup keys + /// + + pub fn load_cpu_setup_data_for_circuit_type( + &self, + key: ProverServiceDataKey, + ) -> anyhow::Result { + let filepath = self.get_file_path(key.clone(), ProverServiceDataType::SetupData); + + let mut file = File::open(filepath.clone()) + .with_context(|| format!("Failed reading setup-data from path: {filepath:?}"))?; + let mut buffer = Vec::new(); + file.read_to_end(&mut buffer).with_context(|| { + format!("Failed reading setup-data to buffer from path: {filepath:?}") + })?; + tracing::info!("loading {:?} setup data from path: {}", key, filepath); + bincode::deserialize::(&buffer).with_context(|| { + format!("Failed deserializing setup-data at path: {filepath:?} for circuit: {key:?}") + }) + } + + #[cfg(feature = "gpu")] + pub fn load_gpu_setup_data_for_circuit_type( + &self, + key: ProverServiceDataKey, + ) -> anyhow::Result { + let filepath = self.get_file_path(key.clone(), ProverServiceDataType::SetupData); + + let mut file = File::open(filepath.clone()) + .with_context(|| format!("Failed reading setup-data from path: {filepath:?}"))?; + let mut buffer = Vec::new(); + file.read_to_end(&mut buffer).with_context(|| { + format!("Failed reading setup-data to buffer from path: {filepath:?}") + })?; + tracing::info!("loading {:?} setup data from path: {}", key, filepath); + bincode::deserialize::(&buffer).with_context(|| { + format!("Failed deserializing setup-data at path: {filepath:?} for circuit: {key:?}") + }) + } + + pub fn save_setup_data_for_circuit_type( + &self, + key: ProverServiceDataKey, + serialized_setup_data: &Vec, + ) -> anyhow::Result<()> { + let filepath = self.get_file_path(key.clone(), ProverServiceDataType::SetupData); + tracing::info!("saving {:?} setup data to: {}", key, filepath); + std::fs::write(filepath.clone(), serialized_setup_data) + .with_context(|| format!("Failed saving setup-data at path: {filepath:?}")) + } + + /// Loads all the verification keys into the Data Source. + /// Keys are loaded from the default 'base path' files. + pub fn load_keys_to_data_source(&self) -> anyhow::Result { + let mut data_source = InMemoryDataSource::new(); + for base_circuit_type in + (BaseLayerCircuitType::VM as u8)..=(BaseLayerCircuitType::L1MessagesHasher as u8) + { + data_source + .set_base_layer_vk(self.load_base_layer_verification_key(base_circuit_type)?) + .unwrap(); + } + + for circuit_type in ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8 + ..=ZkSyncRecursionLayerStorageType::LeafLayerCircuitForL1MessagesHasher as u8 + { + data_source + .set_recursion_layer_vk(self.load_recursive_layer_verification_key(circuit_type)?) + .unwrap(); + } + data_source + .set_recursion_layer_node_vk(self.load_recursive_layer_verification_key( + ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, + )?) + .unwrap(); + Ok(data_source) + } + + pub fn save_keys_from_data_source(&self, source: &dyn SetupDataSource) -> anyhow::Result<()> { + for base_circuit_type in + (BaseLayerCircuitType::VM as u8)..=(BaseLayerCircuitType::L1MessagesHasher as u8) + { + let vk = source.get_base_layer_vk(base_circuit_type).map_err(|err| { + anyhow::anyhow!("No vk exist for circuit type: {base_circuit_type}: {err}") + })?; + self.save_base_layer_verification_key(vk) + .context("save_base_layer_vk()")?; + + let hint = source + .get_base_layer_finalization_hint(base_circuit_type) + .map_err(|err| { + anyhow::anyhow!( + "No finalization_hint exist for circuit type: {base_circuit_type}: {err}" + ) + })? + .into_inner(); + let key = ProverServiceDataKey::new(base_circuit_type, AggregationRound::BasicCircuits); + self.save_finalization_hints(key, &hint) + .context("save_finalization_hints()")?; + } + for leaf_circuit_type in (ZkSyncRecursionLayerStorageType::LeafLayerCircuitForMainVM as u8) + ..=(ZkSyncRecursionLayerStorageType::LeafLayerCircuitForL1MessagesHasher as u8) + { + let vk = source + .get_recursion_layer_vk(leaf_circuit_type) + .map_err(|err| { + anyhow::anyhow!("No vk exist for circuit type: {leaf_circuit_type}: {err}") + })?; + self.save_recursive_layer_verification_key(vk) + .context("save_recursive_layer_vk()")?; + + let hint = source + .get_recursion_layer_finalization_hint(leaf_circuit_type) + .map_err(|err| { + anyhow::anyhow!( + "No finalization hint exist for circuit type: {leaf_circuit_type}: {err}" + ) + })? + .into_inner(); + let key = ProverServiceDataKey::new_recursive(leaf_circuit_type); + self.save_finalization_hints(key, &hint) + .context("save_finalization_hints()")?; + } + self.save_recursive_layer_verification_key( + source + .get_recursion_layer_node_vk() + .map_err(|err| anyhow::anyhow!("No vk exist for node layer circuit: {err}"))?, + ) + .context("save_recursive_layer_vk")?; + + let node_hint = source + .get_recursion_layer_node_finalization_hint() + .map_err(|err| { + anyhow::anyhow!("No finalization hint exist for node layer circuit: {err}") + })? + .into_inner(); + self.save_finalization_hints( + ProverServiceDataKey::new_recursive( + ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, + ), + &node_hint, + ) + .context("save_finalization_hints()")?; + + self.save_recursive_layer_verification_key( + source + .get_recursion_layer_vk(ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8) + .map_err(|err| anyhow::anyhow!("No vk exist for scheduler circuit: {err}"))?, + ) + .context("save_recursive_layer_vk")?; + + let scheduler_hint = source + .get_recursion_layer_finalization_hint( + ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8, + ) + .map_err(|err| { + anyhow::anyhow!("No finalization hint exist for scheduler layer circuit: {err}") + })? + .into_inner(); + + self.save_finalization_hints( + ProverServiceDataKey::new_recursive( + ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8, + ), + &scheduler_hint, + ) + .context("save_finalization_hints()")?; + + Ok(()) + } +} diff --git a/prover/vk_setup_data_generator_server_fri/src/lib.rs b/prover/vk_setup_data_generator_server_fri/src/lib.rs index b3d286314582..dd88d1145617 100644 --- a/prover/vk_setup_data_generator_server_fri/src/lib.rs +++ b/prover/vk_setup_data_generator_server_fri/src/lib.rs @@ -1,61 +1,30 @@ #![feature(generic_const_exprs)] #![feature(allocator_api)] -use std::{fs, fs::File, io::Read}; - -use anyhow::Context as _; -use circuit_definitions::circuit_definitions::aux_layer::{ - ZkSyncCompressionLayerStorageType, ZkSyncSnarkWrapperVK, -}; -use serde::{de::DeserializeOwned, Deserialize, Serialize}; -use zkevm_test_harness::prover_utils::create_base_layer_setup_data; -use zkevm_test_harness_1_3_3::{ - abstract_zksync_circuit::concrete_circuits::ZkSyncCircuit, - bellman::{ - bn256::Bn256, plonk::better_better_cs::setup::VerificationKey as SnarkVerificationKey, - }, - witness::oracle::VmWitnessOracle as SnarkWitnessOracle, -}; -use zksync_config::configs::FriProverConfig; -use zksync_env_config::FromEnv; -use zksync_prover_fri_types::{ - circuit_definitions::{ - aux_definitions::witness_oracle::VmWitnessOracle, - boojum::{ - algebraic_props::{ - round_function::AbsorptionModeOverwrite, sponge::GenericAlgebraicSponge, - }, - cs::{ - implementations::{ - hints::{DenseVariablesCopyHint, DenseWitnessCopyHint}, - polynomial_storage::{SetupBaseStorage, SetupStorage}, - setup::FinalizationHintsForProver, - verifier::VerificationKey, - }, - oracle::{merkle_tree::MerkleTreeWithCap, TreeHasher}, - }, - field::{ - goldilocks::GoldilocksField, traits::field_like::PrimeFieldLikeVectorized, - PrimeField, SmallField, - }, - implementations::poseidon2::Poseidon2Goldilocks, - worker::Worker, - }, - circuit_definitions::{ - base_layer::{ZkSyncBaseLayerCircuit, ZkSyncBaseLayerVerificationKey}, - recursion_layer::{ - ZkSyncRecursionLayerStorageType, ZkSyncRecursionLayerVerificationKey, - }, +use serde::{Deserialize, Serialize}; +use zksync_prover_fri_types::circuit_definitions::boojum::{ + algebraic_props::{round_function::AbsorptionModeOverwrite, sponge::GenericAlgebraicSponge}, + cs::{ + implementations::{ + hints::{DenseVariablesCopyHint, DenseWitnessCopyHint}, + polynomial_storage::{SetupBaseStorage, SetupStorage}, + setup::FinalizationHintsForProver, + verifier::VerificationKey, }, - ZkSyncDefaultRoundFunction, BASE_LAYER_CAP_SIZE, BASE_LAYER_FRI_LDE_FACTOR, + oracle::{merkle_tree::MerkleTreeWithCap, TreeHasher}, + }, + field::{ + goldilocks::GoldilocksField, traits::field_like::PrimeFieldLikeVectorized, PrimeField, + SmallField, }, - ProverServiceDataKey, + implementations::poseidon2::Poseidon2Goldilocks, }; -use zksync_types::basic_fri_types::AggregationRound; #[cfg(feature = "gpu")] use {shivini::cs::GpuSetup, std::alloc::Global}; pub mod commitment_utils; +pub mod keystore; +pub mod setup_data_generator; pub mod utils; pub mod vk_commitment_helper; @@ -125,275 +94,3 @@ pub type GoldilocksGpuProverSetupData = GpuProverSetupData< AbsorptionModeOverwrite, >, >; - -pub enum ProverServiceDataType { - VerificationKey, - SetupData, - FinalizationHints, - SnarkVerificationKey, -} - -pub fn get_base_path() -> String { - let zksync_home = std::env::var("ZKSYNC_HOME").unwrap_or_else(|_| "/".into()); - format!( - "{}/prover/vk_setup_data_generator_server_fri/data", - zksync_home - ) -} - -pub fn get_file_path( - key: ProverServiceDataKey, - service_data_type: ProverServiceDataType, -) -> anyhow::Result { - let name = match key.round { - AggregationRound::BasicCircuits => { - format!("basic_{}", key.circuit_id) - } - AggregationRound::LeafAggregation => { - format!("leaf_{}", key.circuit_id) - } - AggregationRound::NodeAggregation => "node".to_string(), - AggregationRound::Scheduler => "scheduler".to_string(), - }; - Ok(match service_data_type { - ProverServiceDataType::VerificationKey => { - format!("{}/verification_{}_key.json", get_base_path(), name) - } - ProverServiceDataType::SetupData => { - format!( - "{}/setup_{}_data.bin", - FriProverConfig::from_env() - .context("FriProverConfig::from_env()")? - .setup_data_path, - name - ) - } - ProverServiceDataType::FinalizationHints => { - format!("{}/finalization_hints_{}.bin", get_base_path(), name) - } - ProverServiceDataType::SnarkVerificationKey => { - format!("{}/snark_verification_{}_key.json", get_base_path(), name) - } - }) -} - -pub fn get_base_layer_vk_for_circuit_type( - circuit_type: u8, -) -> anyhow::Result { - let filepath = get_file_path( - ProverServiceDataKey::new(circuit_type, AggregationRound::BasicCircuits), - ProverServiceDataType::VerificationKey, - ) - .context("get_file_path()")?; - tracing::info!("Fetching verification key from path: {}", filepath); - let text = std::fs::read_to_string(&filepath) - .with_context(|| format!("Failed reading verification key from path: {filepath}"))?; - serde_json::from_str::(&text) - .with_context(|| format!("Failed deserializing verification key from path: {filepath}")) -} - -pub fn get_recursive_layer_vk_for_circuit_type( - circuit_type: u8, -) -> anyhow::Result { - let round = get_round_for_recursive_circuit_type(circuit_type); - let filepath = get_file_path( - ProverServiceDataKey::new(circuit_type, round), - ProverServiceDataType::VerificationKey, - ) - .context("get_file_path()")?; - tracing::info!("Fetching verification key from path: {}", filepath); - let text = std::fs::read_to_string(&filepath) - .with_context(|| format!("Failed reading verification key from path: {filepath}"))?; - serde_json::from_str::(&text) - .with_context(|| format!("Failed deserializing verification key from path: {filepath}")) -} - -pub fn get_round_for_recursive_circuit_type(circuit_type: u8) -> AggregationRound { - match circuit_type { - circuit_type if circuit_type == ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8 => { - AggregationRound::Scheduler - } - circuit_type if circuit_type == ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8 => { - AggregationRound::NodeAggregation - } - _ => AggregationRound::LeafAggregation, - } -} - -pub fn save_base_layer_vk(vk: ZkSyncBaseLayerVerificationKey) -> anyhow::Result<()> { - let circuit_type = vk.numeric_circuit_type(); - let filepath = get_file_path( - ProverServiceDataKey::new(circuit_type, AggregationRound::BasicCircuits), - ProverServiceDataType::VerificationKey, - ) - .context("get_file_path()")?; - tracing::info!("saving basic verification key to: {}", filepath); - std::fs::write(&filepath, serde_json::to_string_pretty(&vk).unwrap()) - .with_context(|| format!("writing to '{filepath}' failed")) -} - -pub fn save_recursive_layer_vk(vk: ZkSyncRecursionLayerVerificationKey) -> anyhow::Result<()> { - let circuit_type = vk.numeric_circuit_type(); - let round = get_round_for_recursive_circuit_type(circuit_type); - let filepath = get_file_path( - ProverServiceDataKey::new(circuit_type, round), - ProverServiceDataType::VerificationKey, - ) - .context("get_file_path()")?; - tracing::info!("saving recursive layer verification key to: {}", filepath); - std::fs::write(&filepath, serde_json::to_string_pretty(&vk).unwrap()) - .with_context(|| format!("writing to '{filepath}' failed")) -} - -pub fn save_snark_vk(vk: ZkSyncSnarkWrapperVK) -> anyhow::Result<()> { - let filepath = get_file_path( - ProverServiceDataKey::new(vk.numeric_circuit_type(), AggregationRound::Scheduler), - ProverServiceDataType::SnarkVerificationKey, - ) - .context("get_file_path()")?; - tracing::info!("saving snark verification key to: {}", filepath); - fs::write( - &filepath, - serde_json::to_string_pretty(&vk.into_inner()).unwrap(), - ) - .with_context(|| format!("writing to '{filepath}' failed")) -} - -pub fn get_cpu_setup_data_for_circuit_type( - key: ProverServiceDataKey, -) -> anyhow::Result> -where - F: PrimeField + SmallField + Serialize + DeserializeOwned, - P: PrimeFieldLikeVectorized + Serialize + DeserializeOwned, - H: TreeHasher, - >::Output: Serialize + DeserializeOwned, -{ - let filepath = - get_file_path(key.clone(), ProverServiceDataType::SetupData).context("get_file_path()")?; - let mut file = File::open(filepath.clone()) - .with_context(|| format!("Failed reading setup-data from path: {filepath:?}"))?; - let mut buffer = Vec::new(); - file.read_to_end(&mut buffer) - .with_context(|| format!("Failed reading setup-data to buffer from path: {filepath:?}"))?; - tracing::info!("loading {:?} setup data from path: {}", key, filepath); - bincode::deserialize::>(&buffer).with_context(|| { - format!("Failed deserializing setup-data at path: {filepath:?} for circuit: {key:?}") - }) -} - -#[cfg(feature = "gpu")] -pub fn get_setup_data_for_circuit_type( - key: ProverServiceDataKey, -) -> anyhow::Result> -where - F: PrimeField + SmallField + Serialize + DeserializeOwned, - H: TreeHasher, - >::Output: Serialize + DeserializeOwned, -{ - let filepath = - get_file_path(key.clone(), ProverServiceDataType::SetupData).context("get_file_path()")?; - let mut file = File::open(filepath.clone()) - .with_context(|| format!("Failed reading setup-data from path: {filepath:?}"))?; - let mut buffer = Vec::new(); - file.read_to_end(&mut buffer) - .with_context(|| format!("Failed reading setup-data to buffer from path: {filepath:?}"))?; - tracing::info!("loading {:?} setup data from path: {}", key, filepath); - bincode::deserialize::>(&buffer).with_context(|| { - format!("Failed deserializing setup-data at path: {filepath:?} for circuit: {key:?}") - }) -} - -pub fn save_setup_data( - key: ProverServiceDataKey, - serialized_setup_data: &Vec, -) -> anyhow::Result<()> { - let filepath = - get_file_path(key.clone(), ProverServiceDataType::SetupData).context("get_file_path()")?; - tracing::info!("saving {:?} setup data to: {}", key, filepath); - std::fs::write(filepath.clone(), serialized_setup_data) - .with_context(|| format!("Failed saving setup-data at path: {filepath:?}")) -} - -pub fn generate_cpu_base_layer_setup_data( - circuit: ZkSyncBaseLayerCircuit< - GoldilocksField, - VmWitnessOracle, - ZkSyncDefaultRoundFunction, - >, -) -> anyhow::Result { - let circuit_type = circuit.numeric_circuit_type(); - tracing::info!( - "starting setup data generator for base layer circuit: {}.", - circuit_type - ); - let worker = Worker::new(); - let (setup_base, setup, vk, setup_tree, vars_hint, wits_hint, finalization_hint) = - create_base_layer_setup_data( - circuit.clone(), - &worker, - BASE_LAYER_FRI_LDE_FACTOR, - BASE_LAYER_CAP_SIZE, - ); - let key = ProverServiceDataKey::new(circuit_type, AggregationRound::BasicCircuits); - let existing_finalization_hint = - get_finalization_hints(key).context("get_finalization_hints()")?; - if existing_finalization_hint != finalization_hint { - anyhow::bail!("finalization hint mismatch for circuit: {circuit_type}"); - } - let existing_vk = get_base_layer_vk_for_circuit_type(circuit_type) - .with_context(|| format!("get_base_layer_vk_for_circuit_type({circuit_type})"))?; - if existing_vk.into_inner() != vk { - anyhow::bail!("vk mismatch for circuit: {circuit_type}"); - } - Ok(ProverSetupData { - setup_base, - setup, - vk, - setup_tree, - vars_hint, - wits_hint, - finalization_hint, - }) -} - -pub fn save_finalization_hints( - key: ProverServiceDataKey, - hint: &FinalizationHintsForProver, -) -> anyhow::Result<()> { - let filepath = get_file_path(key.clone(), ProverServiceDataType::FinalizationHints) - .context("get_file_path()")?; - tracing::info!("saving finalization hints for {:?} to: {}", key, filepath); - let serialized = bincode::serialize(&hint).context("Failed to serialize finalization hints")?; - fs::write(filepath, serialized).context("Failed to write finalization hints to file") -} -pub fn get_finalization_hints( - key: ProverServiceDataKey, -) -> anyhow::Result { - let mut key = key; - // For `NodeAggregation` round we have only 1 finalization hints for all circuit type. - if key.round == AggregationRound::NodeAggregation { - key.circuit_id = ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8; - } - let filepath = - get_file_path(key, ProverServiceDataType::FinalizationHints).context("get_file_path")?; - let file = fs::read(filepath).context("Failed to read finalization hints from file")?; - bincode::deserialize::(&file) - .context("Finalization hint deserialization failed") -} - -pub fn get_snark_vk( -) -> anyhow::Result>>> { - let circuit_id = ZkSyncCompressionLayerStorageType::CompressionMode1Circuit as u8; - let filepath = get_file_path( - ProverServiceDataKey::new(circuit_id, AggregationRound::Scheduler), - ProverServiceDataType::SnarkVerificationKey, - ) - .context("get_file_path()")?; - tracing::info!("Fetching verification key from path: {}", filepath); - let text = fs::read_to_string(&filepath) - .with_context(|| format!("Failed reading verification key from path: {filepath}"))?; - serde_json::from_str::< - SnarkVerificationKey>>, - >(&text) - .with_context(|| format!("Failed deserializing verification key from path: {filepath}")) -} diff --git a/prover/vk_setup_data_generator_server_fri/src/main.rs b/prover/vk_setup_data_generator_server_fri/src/main.rs index 3abe32ee1935..a14807a7347c 100644 --- a/prover/vk_setup_data_generator_server_fri/src/main.rs +++ b/prover/vk_setup_data_generator_server_fri/src/main.rs @@ -1,13 +1,8 @@ use std::collections::HashMap; use anyhow::Context as _; -use circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerVerificationKey; use clap::{Parser, Subcommand}; use commitment_generator::read_and_update_contract_toml; -use setup_data_generator::{ - generate_all_cpu_setup_data, generate_all_gpu_setup_data, generate_cpu_setup_data, - generate_gpu_setup_data, -}; use tracing::level_filters::LevelFilter; use zkevm_test_harness::{ compute_setups::{generate_base_layer_vks_and_proofs, generate_recursive_layer_vks_and_proofs}, @@ -17,134 +12,21 @@ use zkevm_test_harness::{ WrapperConfig, }, }; -use zksync_prover_fri_types::{ - circuit_definitions::{ - circuit_definitions::recursion_layer::ZkSyncRecursionLayerStorageType, - zkevm_circuits::scheduler::aux::BaseLayerCircuitType, - }, - ProverServiceDataKey, -}; -use zksync_types::basic_fri_types::AggregationRound; +use zksync_prover_fri_types::circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerStorageType; use zksync_vk_setup_data_server_fri::{ - get_base_path, get_round_for_recursive_circuit_type, save_base_layer_vk, - save_finalization_hints, save_recursive_layer_vk, save_snark_vk, + keystore::Keystore, + setup_data_generator::{ + generate_all_cpu_setup_data, generate_all_gpu_setup_data, generate_cpu_setup_data, + generate_gpu_setup_data, + }, }; mod commitment_generator; -mod setup_data_generator; #[cfg(test)] mod tests; -fn save_vks(source: &dyn SetupDataSource) -> anyhow::Result<()> { - for base_circuit_type in - (BaseLayerCircuitType::VM as u8)..=(BaseLayerCircuitType::L1MessagesHasher as u8) - { - let vk = source.get_base_layer_vk(base_circuit_type).map_err(|err| { - anyhow::anyhow!("No vk exist for circuit type: {base_circuit_type}: {err}") - })?; - save_base_layer_vk(vk).context("save_base_layer_vk()")?; - } - for leaf_circuit_type in (ZkSyncRecursionLayerStorageType::LeafLayerCircuitForMainVM as u8) - ..=(ZkSyncRecursionLayerStorageType::LeafLayerCircuitForL1MessagesHasher as u8) - { - let vk = source - .get_recursion_layer_vk(leaf_circuit_type) - .map_err(|err| { - anyhow::anyhow!("No vk exist for circuit type: {leaf_circuit_type}: {err}") - })?; - save_recursive_layer_vk(vk).context("save_recursive_layer_vk()")?; - } - save_recursive_layer_vk( - source - .get_recursion_layer_node_vk() - .map_err(|err| anyhow::anyhow!("No vk exist for node layer circuit: {err}"))?, - ) - .context("save_recursive_layer_vk")?; - save_recursive_layer_vk( - source - .get_recursion_layer_vk(ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8) - .map_err(|err| anyhow::anyhow!("No vk exist for scheduler circuit: {err}"))?, - ) - .context("save_recursive_layer_vk")?; - Ok(()) -} - -fn save_finalization_hints_using_source(source: &dyn SetupDataSource) -> anyhow::Result<()> { - for base_circuit_type in - (BaseLayerCircuitType::VM as u8)..=(BaseLayerCircuitType::L1MessagesHasher as u8) - { - let hint = source - .get_base_layer_finalization_hint(base_circuit_type) - .map_err(|err| { - anyhow::anyhow!( - "No finalization_hint exist for circuit type: {base_circuit_type}: {err}" - ) - })? - .into_inner(); - let key = ProverServiceDataKey::new(base_circuit_type, AggregationRound::BasicCircuits); - save_finalization_hints(key, &hint).context("save_finalization_hints()")?; - } - for leaf_circuit_type in (ZkSyncRecursionLayerStorageType::LeafLayerCircuitForMainVM as u8) - ..=(ZkSyncRecursionLayerStorageType::LeafLayerCircuitForL1MessagesHasher as u8) - { - let hint = source - .get_recursion_layer_finalization_hint(leaf_circuit_type) - .map_err(|err| { - anyhow::anyhow!( - "No finalization hint exist for circuit type: {leaf_circuit_type}: {err}" - ) - })? - .into_inner(); - let key = ProverServiceDataKey::new( - leaf_circuit_type, - get_round_for_recursive_circuit_type(leaf_circuit_type), - ); - save_finalization_hints(key, &hint).context("save_finalization_hints()")?; - } - - let node_hint = source - .get_recursion_layer_node_finalization_hint() - .map_err(|err| anyhow::anyhow!("No finalization hint exist for node layer circuit: {err}"))? - .into_inner(); - save_finalization_hints( - ProverServiceDataKey::new( - ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, - AggregationRound::NodeAggregation, - ), - &node_hint, - ) - .context("save_finalization_hints()")?; - - let scheduler_hint = source - .get_recursion_layer_finalization_hint( - ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8, - ) - .map_err(|err| { - anyhow::anyhow!("No finalization hint exist for scheduler layer circuit: {err}") - })? - .into_inner(); - save_finalization_hints( - ProverServiceDataKey::new( - ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8, - AggregationRound::Scheduler, - ), - &scheduler_hint, - ) - .context("save_finalization_hints()") -} - -fn generate_snark_vk( - scheduler_vk: ZkSyncRecursionLayerVerificationKey, - compression_mode: u8, -) -> anyhow::Result<()> { - let config = WrapperConfig::new(compression_mode); - - let (_, vk) = get_wrapper_setup_and_vk_from_scheduler_vk(scheduler_vk, config); - save_snark_vk(vk).context("save_snark_vk") -} - -fn generate_vks() -> anyhow::Result<()> { +fn generate_vks(keystore: &Keystore) -> anyhow::Result<()> { // Start by checking the trusted setup existence. // This is used at the last step, but we want to fail early if user didn't configure everything // correctly. @@ -157,18 +39,23 @@ fn generate_vks() -> anyhow::Result<()> { tracing::info!("Generating verification keys for Recursive layer."); generate_recursive_layer_vks_and_proofs(&mut in_memory_source) .map_err(|err| anyhow::anyhow!("Failed generating recursive vk's: {err}"))?; - tracing::info!("Saving finalization hints."); - save_finalization_hints_using_source(&in_memory_source) - .context("save_finalization_hints_using_source()")?; - tracing::info!("Saving circuit verification keys."); - save_vks(&in_memory_source).context("save_vks()")?; + tracing::info!("Saving keys & hints"); + + keystore.save_keys_from_data_source(&in_memory_source)?; // Generate snark VK let scheduler_vk = in_memory_source .get_recursion_layer_vk(ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8) .map_err(|err| anyhow::anyhow!("Failed to get scheduler vk: {err}"))?; tracing::info!("Generating verification keys for snark wrapper."); - generate_snark_vk(scheduler_vk, 1).context("generate_snark_vk") + + // Compression mode is 1 + let config = WrapperConfig::new(1); + + let (_, vk) = get_wrapper_setup_and_vk_from_scheduler_vk(scheduler_vk, config); + keystore + .save_snark_verification_key(vk) + .context("save_snark_vk") } #[derive(Debug, Parser)] @@ -255,16 +142,18 @@ fn main() -> anyhow::Result<()> { .init(); let opt = Cli::parse(); + // Setting key store from argument flags will come in next PR. + let keystore = Keystore::default(); match opt.command { Command::GenerateVerificationKeys {} => { tracing::info!( "Generating verification keys and storing them inside {:?}", - get_base_path() + keystore.get_base_path() ); - generate_vks().context("generate_vks()") + generate_vks(&keystore).context("generate_vks()") } - Command::UpdateCommitments { dryrun } => read_and_update_contract_toml(dryrun), + Command::UpdateCommitments { dryrun } => read_and_update_contract_toml(&keystore, dryrun), Command::GenerateSetupKeys { circuits_type, @@ -272,12 +161,13 @@ fn main() -> anyhow::Result<()> { dry_run, } => match circuits_type { CircuitSelector::All => { - let digests = generate_all_cpu_setup_data(dry_run)?; + let digests = generate_all_cpu_setup_data(&keystore, dry_run)?; tracing::info!("CPU Setup keys md5(s):"); print_stats(digests) } CircuitSelector::Recursive => { let digest = generate_cpu_setup_data( + &keystore, false, numeric_circuit.expect("--numeric-circuit must be provided"), dry_run, @@ -288,6 +178,7 @@ fn main() -> anyhow::Result<()> { } CircuitSelector::Basic => { let digest = generate_cpu_setup_data( + &keystore, true, numeric_circuit.expect("--numeric-circuit must be provided"), dry_run, @@ -303,12 +194,13 @@ fn main() -> anyhow::Result<()> { dry_run, } => match circuits_type { CircuitSelector::All => { - let digests = generate_all_gpu_setup_data(dry_run)?; + let digests = generate_all_gpu_setup_data(&keystore, dry_run)?; tracing::info!("GPU Setup keys md5(s):"); print_stats(digests) } CircuitSelector::Recursive => { let digest = generate_gpu_setup_data( + &keystore, false, numeric_circuit.expect("--numeric-circuit must be provided"), dry_run, @@ -319,6 +211,7 @@ fn main() -> anyhow::Result<()> { } CircuitSelector::Basic => { let digest = generate_gpu_setup_data( + &keystore, true, numeric_circuit.expect("--numeric-circuit must be provided"), dry_run, diff --git a/prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs b/prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs index 56e10ca67de0..830ec5bb5805 100644 --- a/prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs +++ b/prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs @@ -6,7 +6,8 @@ use circuit_definitions::{ zkevm_circuits::scheduler::aux::BaseLayerCircuitType, }; use zkevm_test_harness::{ - geometry_config::get_geometry_config, prover_utils::create_recursive_layer_setup_data, + geometry_config::get_geometry_config, + prover_utils::{create_base_layer_setup_data, create_recursive_layer_setup_data}, }; use zksync_prover_fri_types::{ circuit_definitions::{ @@ -20,30 +21,28 @@ use zksync_prover_fri_types::{ ProverServiceDataKey, }; use zksync_types::basic_fri_types::AggregationRound; -use zksync_vk_setup_data_server_fri::{ - generate_cpu_base_layer_setup_data, get_finalization_hints, - get_recursive_layer_vk_for_circuit_type, get_round_for_recursive_circuit_type, save_setup_data, +#[cfg(feature = "gpu")] +use {crate::GpuProverSetupData, shivini::cs::setup::GpuSetup, shivini::ProverContext}; + +use crate::{ + keystore::Keystore, utils::{ get_basic_circuits, get_leaf_circuits, get_node_circuit, get_scheduler_circuit, CYCLE_LIMIT, }, GoldilocksProverSetupData, ProverSetupData, }; -#[cfg(feature = "gpu")] -use { - shivini::cs::setup::GpuSetup, shivini::ProverContext, - zksync_vk_setup_data_server_fri::GpuProverSetupData, -}; fn generate_all( + keystore: &Keystore, dry_run: bool, - f: fn(bool, u8, bool) -> anyhow::Result, + f: fn(&Keystore, bool, u8, bool) -> anyhow::Result, ) -> anyhow::Result> { let mut result = HashMap::new(); for numeric_circuit in BaseLayerCircuitType::VM as u8..=BaseLayerCircuitType::L1MessagesHasher as u8 { - let digest = f(true, numeric_circuit, dry_run) + let digest = f(keystore, true, numeric_circuit, dry_run) .context(format!("base layer, circuit {:?}", numeric_circuit))?; result.insert(format!("base_{}", numeric_circuit), digest); } @@ -52,31 +51,38 @@ fn generate_all( for numeric_circuit in ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8 ..=ZkSyncRecursionLayerStorageType::LeafLayerCircuitForL1MessagesHasher as u8 { - let digest = f(false, numeric_circuit, dry_run) + let digest = f(keystore, false, numeric_circuit, dry_run) .context(format!("recursive layer, circuit {:?}", numeric_circuit))?; result.insert(format!("recursive_{}", numeric_circuit), digest); } Ok(result) } -pub fn generate_all_cpu_setup_data(dry_run: bool) -> anyhow::Result> { - generate_all(dry_run, generate_cpu_setup_data) +pub fn generate_all_cpu_setup_data( + keystore: &Keystore, + dry_run: bool, +) -> anyhow::Result> { + generate_all(keystore, dry_run, generate_cpu_setup_data) } -pub fn generate_all_gpu_setup_data(dry_run: bool) -> anyhow::Result> { - generate_all(dry_run, generate_gpu_setup_data) +pub fn generate_all_gpu_setup_data( + keystore: &Keystore, + dry_run: bool, +) -> anyhow::Result> { + generate_all(keystore, dry_run, generate_gpu_setup_data) } pub fn generate_cpu_setup_data( + keystore: &Keystore, is_base_layer: bool, numeric_circuit: u8, dry_run: bool, ) -> anyhow::Result { match is_base_layer { true => { - let circuit = - get_base_layer_circuit(numeric_circuit).context("get_base_layer_circuit")?; - let prover_setup_data = generate_cpu_base_layer_setup_data(circuit) + let circuit = get_base_layer_circuit(keystore, numeric_circuit) + .context("get_base_layer_circuit")?; + let prover_setup_data = generate_cpu_base_layer_setup_data(keystore, circuit) .context("generate_cpu_base_layer_setup_data()")?; // Serialization should always succeed. let serialized = @@ -84,32 +90,34 @@ pub fn generate_cpu_setup_data( let digest = md5::compute(&serialized); if !dry_run { - save_setup_data( - ProverServiceDataKey::new(numeric_circuit, AggregationRound::BasicCircuits), - &serialized, - ) - .context("save_setup_data()")?; + keystore + .save_setup_data_for_circuit_type( + ProverServiceDataKey::new(numeric_circuit, AggregationRound::BasicCircuits), + &serialized, + ) + .context("save_setup_data_for_circuit_type()")?; } else { tracing::warn!("Dry run - not writing the key"); } Ok(format!("{:?}", digest)) } false => { - let circuit = - get_recursive_circuit(numeric_circuit).context("get_recursive_circuit()")?; - let prover_setup_data = generate_cpu_recursive_layer_setup_data(circuit) + let circuit = get_recursive_circuit(keystore, numeric_circuit) + .context("get_recursive_circuit()")?; + let prover_setup_data = generate_cpu_recursive_layer_setup_data(keystore, circuit) .context("generate_cpu_recursive_layer_setup_data()")?; // Serialization should always succeed. let serialized = bincode::serialize(&prover_setup_data).expect("Failed serializing setup data"); - let round = get_round_for_recursive_circuit_type(numeric_circuit); + let digest = md5::compute(&serialized); if !dry_run { - save_setup_data( - ProverServiceDataKey::new(numeric_circuit, round), - &serialized, - ) - .context("save_setup_data()")?; + keystore + .save_setup_data_for_circuit_type( + ProverServiceDataKey::new_recursive(numeric_circuit), + &serialized, + ) + .context("save_setup_data_for_circuit_type()")?; } else { tracing::warn!("Dry run - not writing the key"); } @@ -119,6 +127,7 @@ pub fn generate_cpu_setup_data( } fn get_base_layer_circuit( + keystore: &Keystore, id: u8, ) -> anyhow::Result< ZkSyncBaseLayerCircuit< @@ -127,24 +136,73 @@ fn get_base_layer_circuit( ZkSyncDefaultRoundFunction, >, > { - get_basic_circuits(CYCLE_LIMIT, get_geometry_config()) + get_basic_circuits(keystore, CYCLE_LIMIT, get_geometry_config()) .context("get_basic_circuits()")? .into_iter() .find(|circuit| id == circuit.numeric_circuit_type()) .with_context(|| format!("No basic circuit found for id: {id}")) } -fn get_recursive_circuit(id: u8) -> anyhow::Result { - let mut recursive_circuits = get_leaf_circuits().context("get_leaf_circuits()")?; - recursive_circuits.push(get_node_circuit().context("get_node_circuit()")?); - recursive_circuits.push(get_scheduler_circuit().context("get_scheduler_circuit()")?); +fn get_recursive_circuit( + keystore: &Keystore, + id: u8, +) -> anyhow::Result { + let mut recursive_circuits = get_leaf_circuits(keystore).context("get_leaf_circuits()")?; + recursive_circuits.push(get_node_circuit(keystore).context("get_node_circuit()")?); + recursive_circuits.push(get_scheduler_circuit(keystore).context("get_scheduler_circuit()")?); recursive_circuits .into_iter() .find(|circuit| id == circuit.numeric_circuit_type()) .with_context(|| format!("No recursive circuit found for id: {id}")) } +pub fn generate_cpu_base_layer_setup_data( + keystore: &Keystore, + circuit: ZkSyncBaseLayerCircuit< + GoldilocksField, + VmWitnessOracle, + ZkSyncDefaultRoundFunction, + >, +) -> anyhow::Result { + let circuit_type = circuit.numeric_circuit_type(); + tracing::info!( + "starting setup data generator for base layer circuit: {}.", + circuit_type + ); + let worker = Worker::new(); + let (setup_base, setup, vk, setup_tree, vars_hint, wits_hint, finalization_hint) = + create_base_layer_setup_data( + circuit.clone(), + &worker, + BASE_LAYER_FRI_LDE_FACTOR, + BASE_LAYER_CAP_SIZE, + ); + let key = ProverServiceDataKey::new(circuit_type, AggregationRound::BasicCircuits); + let existing_finalization_hint = keystore + .load_finalization_hints(key) + .context("load_finalization_hints()")?; + if existing_finalization_hint != finalization_hint { + anyhow::bail!("finalization hint mismatch for circuit: {circuit_type}"); + } + let existing_vk = keystore + .load_base_layer_verification_key(circuit_type) + .with_context(|| format!("load_base_layer_verification_key({circuit_type})"))?; + if existing_vk.into_inner() != vk { + anyhow::bail!("vk mismatch for circuit: {circuit_type}"); + } + Ok(ProverSetupData { + setup_base, + setup, + vk, + setup_tree, + vars_hint, + wits_hint, + finalization_hint, + }) +} + fn generate_cpu_recursive_layer_setup_data( + keystore: &Keystore, circuit: ZkSyncRecursiveLayerCircuit, ) -> anyhow::Result { let circuit_type = circuit.numeric_circuit_type(); @@ -160,17 +218,16 @@ fn generate_cpu_recursive_layer_setup_data( BASE_LAYER_FRI_LDE_FACTOR, BASE_LAYER_CAP_SIZE, ); - let key = ProverServiceDataKey::new( - circuit_type, - get_round_for_recursive_circuit_type(circuit_type), - ); - let existing_finalization_hint = - get_finalization_hints(key).context("get_finalization_hints()")?; + let key = ProverServiceDataKey::new_recursive(circuit_type); + let existing_finalization_hint = keystore + .load_finalization_hints(key) + .context("load_finalization_hints()")?; if existing_finalization_hint != finalization_hint { anyhow::bail!("finalization hint mismatch for circuit: {circuit_type}"); } - let existing_vk = get_recursive_layer_vk_for_circuit_type(circuit_type) - .context("get_recursive_layer_vk_for_circuit_type()")?; + let existing_vk = keystore + .load_recursive_layer_verification_key(circuit_type) + .context("load_recursive_layer_verification_key()")?; if existing_vk.into_inner() != vk { anyhow::bail!("vk mismatch for circuit: {circuit_type}"); } @@ -187,6 +244,7 @@ fn generate_cpu_recursive_layer_setup_data( #[cfg(not(feature = "gpu"))] pub fn generate_gpu_setup_data( + _keystore: &Keystore, _is_base_layer: bool, _numeric_circuit: u8, _dry_run: bool, @@ -196,28 +254,29 @@ pub fn generate_gpu_setup_data( #[cfg(feature = "gpu")] pub fn generate_gpu_setup_data( + keystore: &Keystore, is_base_layer: bool, numeric_circuit: u8, dry_run: bool, -) -> anyhow::Result<(String)> { +) -> anyhow::Result { let _context = ProverContext::create().context("failed initializing gpu prover context")?; - let (cpu_setup_data, round) = match is_base_layer { + let (cpu_setup_data, key) = match is_base_layer { true => { - let circuit = - get_base_layer_circuit(numeric_circuit).context("get_base_layer_circuit()")?; + let circuit = get_base_layer_circuit(keystore, numeric_circuit) + .context("get_base_layer_circuit()")?; ( - generate_cpu_base_layer_setup_data(circuit.clone()) + generate_cpu_base_layer_setup_data(keystore, circuit.clone()) .context("generate_cpu_base_layer_setup_data()")?, - AggregationRound::BasicCircuits, + ProverServiceDataKey::new(numeric_circuit, AggregationRound::BasicCircuits), ) } false => { - let circuit = - get_recursive_circuit(numeric_circuit).context("get_recursive_circuit()")?; + let circuit = get_recursive_circuit(keystore, numeric_circuit) + .context("get_recursive_circuit()")?; ( - generate_cpu_recursive_layer_setup_data(circuit.clone()) + generate_cpu_recursive_layer_setup_data(keystore, circuit.clone()) .context("generate_cpu_recursive_layer_setup_data()")?, - get_round_for_recursive_circuit_type(numeric_circuit), + ProverServiceDataKey::new_recursive(numeric_circuit), ) } }; @@ -240,11 +299,9 @@ pub fn generate_gpu_setup_data( bincode::serialize(&gpu_prover_setup_data).expect("Failed serializing setup data"); let digest = md5::compute(&serialized); if !dry_run { - save_setup_data( - ProverServiceDataKey::new(numeric_circuit, round), - &serialized, - ) - .context("save_setup_data")?; + keystore + .save_setup_data_for_circuit_type(key, &serialized) + .context("save_setup_data_for_circuit_type")?; } else { tracing::warn!("Dry run - not writing the key"); } diff --git a/prover/vk_setup_data_generator_server_fri/src/tests.rs b/prover/vk_setup_data_generator_server_fri/src/tests.rs index 0059a646fd87..41aba88f784c 100644 --- a/prover/vk_setup_data_generator_server_fri/src/tests.rs +++ b/prover/vk_setup_data_generator_server_fri/src/tests.rs @@ -1,19 +1,13 @@ use proptest::prelude::*; use zksync_prover_fri_types::{ circuit_definitions::{ - circuit_definitions::recursion_layer::{ - base_circuit_type_into_recursive_leaf_circuit_type, ZkSyncRecursionLayerStorageType, - }, + circuit_definitions::recursion_layer::base_circuit_type_into_recursive_leaf_circuit_type, zkevm_circuits::scheduler::aux::BaseLayerCircuitType, }, ProverServiceDataKey, }; use zksync_types::basic_fri_types::AggregationRound; -use zksync_vk_setup_data_server_fri::{ - get_base_layer_vk_for_circuit_type, get_base_path, get_file_path, get_finalization_hints, - get_recursive_layer_vk_for_circuit_type, get_round_for_recursive_circuit_type, - ProverServiceDataType, -}; +use zksync_vk_setup_data_server_fri::keystore::Keystore; fn all_possible_prover_service_data_key() -> impl Strategy { let mut keys = Vec::with_capacity(30); @@ -42,19 +36,23 @@ fn all_possible_prover_service_data_key() -> impl Strategy anyhow::Result { } pub fn get_basic_circuits( + keystore: &Keystore, cycle_limit: usize, geometry: GeometryConfig, ) -> anyhow::Result< @@ -85,7 +84,7 @@ pub fn get_basic_circuits( >, >, > { - let path = format!("{}/witness_artifacts.json", get_base_path()); + let path = format!("{}/witness_artifacts.json", keystore.get_base_path()); let mut test_artifact = read_witness_artifact(&path).context("read_withess_artifact()")?; let mut storage_impl = InMemoryStorage::new(); @@ -177,14 +176,15 @@ pub fn get_basic_circuits( .collect()) } -pub fn get_leaf_circuits() -> anyhow::Result> { +pub fn get_leaf_circuits(keystore: &Keystore) -> anyhow::Result> { let mut circuits = vec![]; for base_circuit_type in (BaseLayerCircuitType::VM as u8)..=(BaseLayerCircuitType::L1MessagesHasher as u8) { let input = RecursionLeafInput::placeholder_witness(); - let vk = get_base_layer_vk_for_circuit_type(base_circuit_type) - .with_context(|| format!("get_base_layer_vk_for_circuit_type({base_circuit_type})"))?; + let vk = keystore + .load_base_layer_verification_key(base_circuit_type) + .with_context(|| format!("load_base_layer_verification_key({base_circuit_type})"))?; let witness = RecursionLeafInstanceWitness { input, @@ -219,13 +219,14 @@ pub fn get_leaf_circuits() -> anyhow::Result> { Ok(circuits) } -pub fn get_node_circuit() -> anyhow::Result { +pub fn get_node_circuit(keystore: &Keystore) -> anyhow::Result { let input = RecursionNodeInput::placeholder_witness(); - let input_vk = get_recursive_layer_vk_for_circuit_type( - ZkSyncRecursionLayerStorageType::LeafLayerCircuitForMainVM as u8, - ) - .context("get_recursive_layer_vk_for_circuit_type(LeafLAyerCircyutFromMainVM")?; + let input_vk = keystore + .load_recursive_layer_verification_key( + ZkSyncRecursionLayerStorageType::LeafLayerCircuitForMainVM as u8, + ) + .context("load_recursive_layer_verification_key(LeafLAyerCircyutFromMainVM")?; let witness = RecursionNodeInstanceWitness { input, vk_witness: input_vk.clone().into_inner(), @@ -249,15 +250,16 @@ pub fn get_node_circuit() -> anyhow::Result { Ok(ZkSyncRecursiveLayerCircuit::NodeLayerCircuit(circuit)) } -pub fn get_scheduler_circuit() -> anyhow::Result { +pub fn get_scheduler_circuit(keystore: &Keystore) -> anyhow::Result { let mut scheduler_witness = SchedulerCircuitInstanceWitness::placeholder(); // node VK - let node_vk = get_recursive_layer_vk_for_circuit_type( - ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, - ) - .context("get_recursive_layer_vk_for_circuit_type(NodeLayerCircuit)")? - .into_inner(); + let node_vk = keystore + .load_recursive_layer_verification_key( + ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, + ) + .context("load_recursive_layer_verification_key(NodeLayerCircuit)")? + .into_inner(); scheduler_witness.node_layer_vk_witness = node_vk.clone(); let config = SchedulerConfig { @@ -299,6 +301,7 @@ fn get_recursive_layer_proofs() -> Vec { } pub fn get_leaf_vk_params( + keystore: &Keystore, ) -> anyhow::Result)>> { let mut leaf_vk_commits = vec![]; @@ -308,9 +311,11 @@ pub fn get_leaf_vk_params( let recursive_circuit_type = base_circuit_type_into_recursive_leaf_circuit_type( BaseLayerCircuitType::from_numeric_value(circuit_type), ); - let base_vk = get_base_layer_vk_for_circuit_type(circuit_type) + let base_vk = keystore + .load_base_layer_verification_key(circuit_type) .with_context(|| format!("get_base_layer_vk_for_circuit_type({circuit_type})"))?; - let leaf_vk = get_recursive_layer_vk_for_circuit_type(recursive_circuit_type as u8) + let leaf_vk = keystore + .load_recursive_layer_verification_key(recursive_circuit_type as u8) .with_context(|| { format!("get_recursive_layer_vk_for_circuit_type({recursive_circuit_type:?})") })?; diff --git a/prover/witness_generator/src/leaf_aggregation.rs b/prover/witness_generator/src/leaf_aggregation.rs index dd2b5805e42d..da2711e7a268 100644 --- a/prover/witness_generator/src/leaf_aggregation.rs +++ b/prover/witness_generator/src/leaf_aggregation.rs @@ -30,9 +30,7 @@ use zksync_queued_job_processor::JobProcessor; use zksync_types::{ basic_fri_types::AggregationRound, protocol_version::FriProtocolVersionId, L1BatchNumber, }; -use zksync_vk_setup_data_server_fri::{ - get_base_layer_vk_for_circuit_type, get_recursive_layer_vk_for_circuit_type, -}; +use zksync_vk_setup_data_server_fri::keystore::Keystore; use crate::{ metrics::WITNESS_GENERATOR_METRICS, @@ -215,11 +213,14 @@ pub async fn prepare_leaf_aggregation_job( .observe(started_at.elapsed()); let started_at = Instant::now(); - let base_vk = get_base_layer_vk_for_circuit_type(metadata.circuit_id) + let keystore = Keystore::default(); + let base_vk = keystore + .load_base_layer_verification_key(metadata.circuit_id) .context("get_base_layer_vk_for_circuit_type()")?; // this is a temp solution to unblock shadow proving. // we should have a method that converts basic circuit id to leaf circuit id as they are different. - let leaf_vk = get_recursive_layer_vk_for_circuit_type(metadata.circuit_id + 2) + let leaf_vk = keystore + .load_recursive_layer_verification_key(metadata.circuit_id + 2) .context("get_recursive_layer_vk_for_circuit_type()")?; let mut base_proofs = vec![]; for wrapper in proofs { diff --git a/prover/witness_generator/src/node_aggregation.rs b/prover/witness_generator/src/node_aggregation.rs index 3c46ed98d50e..324b05ace8aa 100644 --- a/prover/witness_generator/src/node_aggregation.rs +++ b/prover/witness_generator/src/node_aggregation.rs @@ -26,9 +26,7 @@ use zksync_queued_job_processor::JobProcessor; use zksync_types::{ basic_fri_types::AggregationRound, protocol_version::FriProtocolVersionId, L1BatchNumber, }; -use zksync_vk_setup_data_server_fri::{ - get_recursive_layer_vk_for_circuit_type, utils::get_leaf_vk_params, -}; +use zksync_vk_setup_data_server_fri::{keystore::Keystore, utils::get_leaf_vk_params}; use crate::{ metrics::WITNESS_GENERATOR_METRICS, @@ -241,12 +239,15 @@ pub async fn prepare_job( .observe(started_at.elapsed()); let started_at = Instant::now(); - let leaf_vk = get_recursive_layer_vk_for_circuit_type(metadata.circuit_id) + let keystore = Keystore::default(); + let leaf_vk = keystore + .load_recursive_layer_verification_key(metadata.circuit_id) .context("get_recursive_layer_vk_for_circuit_type")?; - let node_vk = get_recursive_layer_vk_for_circuit_type( - ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, - ) - .context("get_recursive_layer_vk_for_circuit_type()")?; + let node_vk = keystore + .load_recursive_layer_verification_key( + ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, + ) + .context("get_recursive_layer_vk_for_circuit_type()")?; let mut recursive_proofs = vec![]; for wrapper in proofs { @@ -272,7 +273,7 @@ pub async fn prepare_job( proofs: recursive_proofs, leaf_vk, node_vk, - all_leafs_layer_params: get_leaf_vk_params().context("get_leaf_vk_params()")?, + all_leafs_layer_params: get_leaf_vk_params(&keystore).context("get_leaf_vk_params()")?, }) } diff --git a/prover/witness_generator/src/scheduler.rs b/prover/witness_generator/src/scheduler.rs index 68e48f832890..98fe8338f85f 100644 --- a/prover/witness_generator/src/scheduler.rs +++ b/prover/witness_generator/src/scheduler.rs @@ -26,9 +26,7 @@ use zksync_queued_job_processor::JobProcessor; use zksync_types::{ basic_fri_types::AggregationRound, protocol_version::FriProtocolVersionId, L1BatchNumber, }; -use zksync_vk_setup_data_server_fri::{ - get_recursive_layer_vk_for_circuit_type, utils::get_leaf_vk_params, -}; +use zksync_vk_setup_data_server_fri::{keystore::Keystore, utils::get_leaf_vk_params}; use crate::{ metrics::WITNESS_GENERATOR_METRICS, @@ -256,17 +254,19 @@ pub async fn prepare_job( } let started_at = Instant::now(); - let node_vk = get_recursive_layer_vk_for_circuit_type( - ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, - ) - .context("get_recursive_layer_vk_for_circuit_type()")?; + let keystore = Keystore::default(); + let node_vk = keystore + .load_recursive_layer_verification_key( + ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8, + ) + .context("get_recursive_layer_vk_for_circuit_type()")?; let SchedulerPartialInputWrapper(mut scheduler_witness) = object_store.get(l1_batch_number).await.unwrap(); scheduler_witness.node_layer_vk_witness = node_vk.clone().into_inner(); scheduler_witness.proof_witnesses = recursive_proofs.into(); - let leaf_vk_commits = get_leaf_vk_params().context("get_leaf_vk_params()")?; + let leaf_vk_commits = get_leaf_vk_params(&keystore).context("get_leaf_vk_params()")?; let leaf_layer_params = leaf_vk_commits .iter() .map(|el| el.1.clone()) diff --git a/prover/witness_vector_generator/src/generator.rs b/prover/witness_vector_generator/src/generator.rs index 45a24518ab8c..af400042a9d3 100644 --- a/prover/witness_vector_generator/src/generator.rs +++ b/prover/witness_vector_generator/src/generator.rs @@ -19,7 +19,7 @@ use zksync_prover_fri_utils::{ }; use zksync_queued_job_processor::JobProcessor; use zksync_types::{basic_fri_types::CircuitIdRoundTuple, protocol_version::L1VerifierConfig}; -use zksync_vk_setup_data_server_fri::get_finalization_hints; +use zksync_vk_setup_data_server_fri::keystore::Keystore; use crate::metrics::METRICS; @@ -55,7 +55,9 @@ impl WitnessVectorGenerator { } pub fn generate_witness_vector(job: ProverJob) -> anyhow::Result { - let finalization_hints = get_finalization_hints(job.setup_data_key.clone()) + let keystore = Keystore::default(); + let finalization_hints = keystore + .load_finalization_hints(job.setup_data_key.clone()) .context("get_finalization_hints()")?; let cs = match job.circuit_wrapper.clone() { CircuitWrapper::Base(base_circuit) => {