fix verifier circuits keys loading

ChainSafe · Jan 5, 2024 · d089449 · d089449
1 parent 9bdd7b4
commit d089449
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 23 deletions.
diff --git a/prover/src/prover.rs b/prover/src/prover.rs
@@ -13,6 +13,7 @@ use lightclient_circuits::halo2_proofs::{
 };
 use lightclient_circuits::sync_step_circuit::StepCircuit;
 use lightclient_circuits::util::AppCircuit;
+use snark_verifier_sdk::halo2::aggregation::AggregationCircuit;
 use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
 
@@ -44,14 +45,12 @@ impl ProverState {
             config_path: PathBuf,
             pk_path: PathBuf,
             params_map: &mut BTreeMap<u32, ParamsKZG<Bn256>>,
-        ) -> CircuitContext
-        where
-            Circuit::Witness: Default,
-        {
+            default_witness: Circuit::Witness,
+        ) -> CircuitContext {
             let degree = Circuit::get_degree(&config_path);
             let params = gen_srs(degree);
 
-            let pk = Circuit::read_pk(&params, pk_path, &config_path, &Circuit::Witness::default());
+            let pk = Circuit::read_pk(&params, pk_path, &config_path, &default_witness);
 
             params_map.insert(degree, params);
 
@@ -62,26 +61,52 @@ impl ProverState {
             }
         }
 
+        let step = load_ctx::<StepCircuit<S, Fr>>(
+            config_dir.join(format!("sync_step_{}.json", S::NAME)),
+            build_dir.join(format!("sync_step_{}.pkey", S::NAME)),
+            &mut params_map,
+            Default::default(),
+        );
+
+        let step_snark = StepCircuit::<S, Fr>::gen_snark_shplonk(
+            params_map.get(step.degree()).unwrap(),
+            step.pk(),
+            step.config_path(),
+            None::<String>,
+            &Default::default(),
+        )
+        .unwrap();
+
+        let committee_update = load_ctx::<CommitteeUpdateCircuit<S, Fr>>(
+            config_dir.join(format!("committee_update_{}.json", S::NAME)),
+            build_dir.join(format!("committee_update_{}.pkey", S::NAME)),
+            &mut params_map,
+            Default::default(),
+        );
+
+        let committee_update_snark = CommitteeUpdateCircuit::<S, Fr>::gen_snark_shplonk(
+            params_map.get(committee_update.degree()).unwrap(),
+            committee_update.pk(),
+            committee_update.config_path(),
+            None::<String>,
+            &Default::default(),
+        )
+        .unwrap();
+
         Self {
-            step: load_ctx::<StepCircuit<S, Fr>>(
-                config_dir.join(format!("sync_step_{}.json", S::NAME)),
-                build_dir.join(format!("sync_step_{}.pkey", S::NAME)),
-                &mut params_map,
-            ),
-            step_verifier: load_ctx::<StepCircuit<S, Fr>>(
+            step,
+            step_verifier: load_ctx::<AggregationCircuit>(
                 config_dir.join(format!("sync_step_verifier_{}.json", S::NAME)),
                 build_dir.join(format!("sync_step_verifier_{}.pkey", S::NAME)),
                 &mut params_map,
+                vec![step_snark],
             ),
-            committee_update: load_ctx::<CommitteeUpdateCircuit<S, Fr>>(
-                config_dir.join(format!("committee_update_{}.json", S::NAME)),
-                build_dir.join(format!("committee_update_{}.pkey", S::NAME)),
-                &mut params_map,
-            ),
-            committee_update_verifier: load_ctx::<CommitteeUpdateCircuit<S, Fr>>(
+            committee_update,
+            committee_update_verifier: load_ctx::<AggregationCircuit>(
                 config_dir.join(format!("committee_update_verifier_{}.json", S::NAME)),
                 build_dir.join(format!("committee_update_verifier_{}.pkey", S::NAME)),
                 &mut params_map,
+                vec![committee_update_snark],
             ),
             params: params_map,
         }

diff --git a/prover/src/rpc.rs b/prover/src/rpc.rs
@@ -2,6 +2,7 @@
 // Code: https://github.com/ChainSafe/Spectre
 // SPDX-License-Identifier: LGPL-3.0-only
 
+use ark_std::{end_timer, start_timer};
 use axum::{http::StatusCode, response::IntoResponse, routing::post, Router};
 use ethers::prelude::*;
 use itertools::Itertools;
@@ -71,7 +72,7 @@ where
     let mut update = ssz_rs::deserialize(&light_client_update)?;
     let witness = rotation_args_from_update(&mut update).await?;
     let params = state.params.get(state.committee_update.degree()).unwrap();
-    
+
     let snark = gen_uncompressed_snark::<CommitteeUpdateCircuit<S, Fr>>(
         state.committee_update.config_path(),
         params,
@@ -131,7 +132,7 @@ where
     let pubkeys = ssz_rs::deserialize(&pubkeys)?;
     let witness = step_args_from_finality_update(update, pubkeys, domain).await?;
     let params = state.params.get(state.step.degree()).unwrap();
-    
+
     let snark = gen_uncompressed_snark::<StepCircuit<S, Fr>>(
         state.step.config_path(),
         params,
@@ -196,10 +197,10 @@ where
     [(); S::FINALIZED_HEADER_INDEX]:,
 {
     let tcp_listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", port)).await?;
-    let rpc_server = Arc::new(jsonrpc_server::<S>(ProverState::new::<S>(
-        config_dir.as_ref(),
-        build_dir.as_ref(),
-    )));
+    let timer = start_timer!(|| "Load proving keys");
+    let state = ProverState::new::<S>(config_dir.as_ref(), build_dir.as_ref());
+    end_timer!(timer);
+    let rpc_server = Arc::new(jsonrpc_server::<S>(state));
 
     let router = Router::new()
         .route("/rpc", post(handler))