feat: Added support for generate vk and to use custom CRS files (#13)

Author: mm-zk

Cargo.toml

@@ -32,6 +32,7 @@ circuit_mersenne_field = {path="circuit_mersenne_field"}
 serde_json = { version = "*" }
 serde = { version = "1", default-features = false, features = ["derive", "alloc"]}
 clap = { version = "4.5.21", features = ["derive"] }
+sha3 = "*"
 
 [dev-dependencies]
 rand = "0.8.4"

README.md

@@ -6,5 +6,23 @@ It takes the zkos proof (in json format), and returns the boojum proof.
 You can use it as a library, or as a cli tool:
 
 ```
-cargo run --release -- --input testing_data/risc_proof --output-dir tmp.json
-```
+cargo run --release -- prove --input testing_data/risc_proof --output-dir tmp.json
+```
+
+## CRS file
+
+For production - please make sure to use the real CRS file:
+
+https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2^24.key
+
+The tool also offers to use the 'fake' file - that works for local testing, but must NEVER be used for production cases.
+
+## Generating verification key
+
+
+```shell
+cargo run --release generate-vk --input-binary ../air_compiler/examples/hashed_fibonacci/app.bin --output-dir /tmp --trusted-setup-file crs/setup.key
+```
+
+This will generate the verification key (and verification key hash) - that can be put into the solidity code.
+Or you can use era-boojum-verifier-cli tool to verify the proofs manually.

crs/.gitignore

@@ -0,0 +1 @@
+*.key

crs/README.md

@@ -0,0 +1,9 @@
+# CRS key directory
+
+You can put a real CRS file here, by downloading it from:
+
+```shell
+curl https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2\^24.key --output setup.key
+```
+
+We don't include it in github, as it is around 1GB in size.

src/lib.rs

@@ -74,6 +74,7 @@ use bellman::plonk::better_better_cs::cs::PlonkCsWidth4WithNextStepAndCustomGate
 use bellman::plonk::better_better_cs::cs::{ProvingAssembly, SetupAssembly};
 use bellman::plonk::better_better_cs::gates::selector_optimized_with_d_next::SelectorOptimizedWidth4MainGateWithDNext;
 use bellman::plonk::better_better_cs::verifier::verify as verify_snark;
+use boojum::ethereum_types::H256;
 use snark_wrapper::implementations::poseidon2::CircuitPoseidon2Sponge;
 use snark_wrapper::implementations::poseidon2::transcript::CircuitPoseidon2Transcript;
 
@@ -392,3 +393,69 @@ pub fn prove_snark_wrapper(
 pub fn verify_snark_wrapper_proof(proof: &SnarkWrapperProof, vk: &SnarkWrapperVK) -> bool {
     verify_snark::<Bn256, SnarkWrapperCircuit, SnarkWrapperTranscript>(vk, proof, None).unwrap()
 }
+
+// TODO: this should probably be moved somewhere into crypto.
+pub fn calculate_verification_key_hash(verification_key: SnarkWrapperVK) -> H256 {
+    use bellman::compact_bn256::Fq;
+    use bellman::{CurveAffine, PrimeField, PrimeFieldRepr};
+    use sha3::{Digest, Keccak256};
+
+    let mut res = vec![];
+
+    // gate setup commitments
+    assert_eq!(8, verification_key.gate_setup_commitments.len());
+
+    for gate_setup in verification_key.gate_setup_commitments {
+        let (x, y) = gate_setup.as_xy();
+        x.into_repr().write_be(&mut res).unwrap();
+        y.into_repr().write_be(&mut res).unwrap();
+    }
+
+    // gate selectors commitments
+    assert_eq!(2, verification_key.gate_selectors_commitments.len());
+
+    for gate_selector in verification_key.gate_selectors_commitments {
+        let (x, y) = gate_selector.as_xy();
+        x.into_repr().write_be(&mut res).unwrap();
+        y.into_repr().write_be(&mut res).unwrap();
+    }
+
+    // permutation commitments
+    assert_eq!(4, verification_key.permutation_commitments.len());
+
+    for permutation in verification_key.permutation_commitments {
+        let (x, y) = permutation.as_xy();
+        x.into_repr().write_be(&mut res).unwrap();
+        y.into_repr().write_be(&mut res).unwrap();
+    }
+
+    // lookup selector commitment
+    let lookup_selector = verification_key.lookup_selector_commitment.unwrap();
+    let (x, y) = lookup_selector.as_xy();
+    x.into_repr().write_be(&mut res).unwrap();
+    y.into_repr().write_be(&mut res).unwrap();
+
+    // lookup tables commitments
+    assert_eq!(4, verification_key.lookup_tables_commitments.len());
+
+    for table_commit in verification_key.lookup_tables_commitments {
+        let (x, y) = table_commit.as_xy();
+        x.into_repr().write_be(&mut res).unwrap();
+        y.into_repr().write_be(&mut res).unwrap();
+    }
+
+    // table type commitment
+    let lookup_table = verification_key.lookup_table_type_commitment.unwrap();
+    let (x, y) = lookup_table.as_xy();
+    x.into_repr().write_be(&mut res).unwrap();
+    y.into_repr().write_be(&mut res).unwrap();
+
+    // flag for using recursive part
+    Fq::default().into_repr().write_be(&mut res).unwrap();
+
+    let mut hasher = Keccak256::new();
+    hasher.update(&res);
+    let computed_vk_hash = hasher.finalize();
+
+    H256::from_slice(&computed_vk_hash)
+}

src/main.rs

@@ -8,7 +8,7 @@ use std::alloc::Global;
 // - wrapping the proof into SNARK.
 use std::path::Path;
 
-use clap::Parser;
+use clap::{Parser, Subcommand};
 
 use bellman::kate_commitment::{Crs, CrsForMonomialForm};
 use bellman::worker::Worker as BellmanWorker;
@@ -20,7 +20,7 @@ use execution_utils::{
 };
 use risc_verifier::prover::worker::Worker;
 use zkos_wrapper::circuits::BinaryCommitment;
-use zkos_wrapper::{Bn256, L1_VERIFIER_DOMAIN_SIZE_LOG};
+use zkos_wrapper::{Bn256, L1_VERIFIER_DOMAIN_SIZE_LOG, calculate_verification_key_hash};
 use zkos_wrapper::{
     circuits::RiscWrapperWitness, get_compression_setup, get_risc_wrapper_setup,
     get_snark_wrapper_setup, prove_compression, prove_risc_wrapper, prove_snark_wrapper,
@@ -30,16 +30,50 @@ use zkos_wrapper::{
 #[derive(Parser)]
 #[command(version, about, long_about = None)]
 struct Cli {
-    #[arg(short, long)]
-    input: String,
-
-    // Binary used to generate the proof.
-    // If not specified, take the default binary (fibonacci hasher).
-    #[arg(long)]
-    input_binary: Option<String>,
+    #[command(subcommand)]
+    command: Commands,
+}
 
-    #[arg(short, long)]
-    output_dir: String,
+#[derive(Subcommand)]
+enum Commands {
+    /// Take the riscV final proof, and create a SNARK proof.
+    Prove {
+        #[arg(short, long)]
+        input: String,
+
+        // Binary used to generate the proof.
+        // If not specified, take the default binary (fibonacci hasher).
+        #[arg(long)]
+        input_binary: Option<String>,
+
+        #[arg(short, long)]
+        output_dir: String,
+
+        /// File with the trusted setup.
+        /// If missing - will use the 'fake' trusted setup.
+        #[arg(long)]
+        trusted_setup_file: Option<String>,
+    },
+    /// Generate verification key for the SNARK proof.
+    GenerateVk {
+        // Binary used to generate the proof.
+        // If not specified, take the default binary (fibonacci hasher).
+        #[arg(long)]
+        input_binary: String,
+
+        #[arg(short, long)]
+        output_dir: String,
+
+        /// File with the trusted setup.
+        /// If missing - will use the 'fake' trusted setup.
+        #[arg(long)]
+        trusted_setup_file: String,
+
+        /// If true, then create VK for universal verifier program.
+        /// If false then for the separate verifiers.
+        #[arg(long)]
+        universal_verifier: bool,
+    },
 }
 
 fn serialize_to_file<T: serde::ser::Serialize>(content: &T, filename: &Path) {
@@ -52,6 +86,14 @@ fn deserialize_from_file<T: serde::de::DeserializeOwned>(filename: &str) -> T {
     serde_json::from_reader(src).unwrap()
 }
 
+/// Uploads trusted setup file to the RAM
+pub fn get_trusted_setup(crs_file_str: &String) -> Crs<Bn256, CrsForMonomialForm> {
+    let crs_file_path = std::path::Path::new(crs_file_str);
+    let crs_file = std::fs::File::open(&crs_file_path)
+        .expect(format!("Trying to open CRS FILE: {:?}", crs_file_path).as_str());
+    Crs::read(&crs_file).expect(format!("Trying to read CRS FILE: {:?}", crs_file_path).as_str())
+}
+
 pub fn create_binary_commitment(
     binary_path: String,
     expected_end_params: &[u32; 8],
@@ -103,12 +145,106 @@ pub fn create_binary_commitment(
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let cli = Cli::parse();
 
+    match cli.command {
+        Commands::Prove {
+            input,
+            input_binary,
+            output_dir,
+            trusted_setup_file,
+        } => {
+            println!("=== Phase 0: Proving");
+            prove(input, input_binary, output_dir, trusted_setup_file)?;
+        }
+        Commands::GenerateVk {
+            input_binary,
+            output_dir,
+            trusted_setup_file,
+            universal_verifier,
+        } => {
+            println!("=== Phase 0: Generating the verification key");
+            generate_vk(
+                input_binary,
+                output_dir,
+                trusted_setup_file,
+                universal_verifier,
+            )?;
+        }
+    }
+    Ok(())
+}
+
+fn generate_vk(
+    input_binary: String,
+    output_dir: String,
+    trusted_setup_file: String,
+    universal_verifier: bool,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let worker = BellmanWorker::new_with_cpus(4);
+    let boojum_worker = boojum::worker::Worker::new_with_num_threads(4);
+
+    let trusted_setup_file = Some(trusted_setup_file);
+
+    let crs_mons = match trusted_setup_file {
+        Some(ref crs_file_str) => get_trusted_setup(crs_file_str),
+        None => Crs::<Bn256, CrsForMonomialForm>::crs_42(
+            1 << L1_VERIFIER_DOMAIN_SIZE_LOG,
+            &BellmanWorker::new(),
+        ),
+    };
+    println!("=== Phase 1: Creating the Risc wrapper key");
+
+    let verifier_params = if universal_verifier {
+        universal_circuit_no_delegation_verifier_vk().params
+    } else {
+        final_recursion_layer_verifier_vk().params
+    };
+
+    let binary_commitment = create_binary_commitment(input_binary, &verifier_params);
+
+    let (_, _, _, risc_wrapper_vk, _, _, _) =
+        get_risc_wrapper_setup(&boojum_worker, binary_commitment.clone());
+
+    println!("=== Phase 2: Creating the Compression key");
+    let (_, _, _, compression_vk, _, _, _) =
+        get_compression_setup(risc_wrapper_vk.clone(), &boojum_worker);
+
+    println!("=== Phase 3: Creating the SNARK key");
+
+    let (_, snark_wrapper_vk) = get_snark_wrapper_setup(compression_vk.clone(), &crs_mons, &worker);
+
+    serialize_to_file(
+        &snark_wrapper_vk,
+        &Path::new(&output_dir.clone()).join("snark_vk_expected.json"),
+    );
+
+    println!(
+        "VK key hash: {:?}",
+        calculate_verification_key_hash(snark_wrapper_vk)
+    );
+
+    Ok(())
+}
+
+fn prove(
+    input: String,
+    input_binary: Option<String>,
+    output_dir: String,
+    trusted_setup_file: Option<String>,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let crs_mons = match trusted_setup_file {
+        Some(ref crs_file_str) => get_trusted_setup(crs_file_str),
+        None => Crs::<Bn256, CrsForMonomialForm>::crs_42(
+            1 << L1_VERIFIER_DOMAIN_SIZE_LOG,
+            &BellmanWorker::new(),
+        ),
+    };
+
     println!("=== Phase 1: Creating the Risc wrapper proof");
 
     let worker = boojum::worker::Worker::new_with_num_threads(4);
 
-    let program_proof: zkos_wrapper::ProgramProof = deserialize_from_file(&cli.input);
-    let binary_commitment = match cli.input_binary {
+    let program_proof: zkos_wrapper::ProgramProof = deserialize_from_file(&input);
+    let binary_commitment = match input_binary {
         Some(binary_path) => create_binary_commitment(binary_path, &program_proof.end_params),
         None => BinaryCommitment::from_default_binary(),
     };
@@ -143,7 +279,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     serialize_to_file(
         &risc_wrapper_proof,
-        &Path::new(&cli.output_dir.clone()).join("risc_proof.json"),
+        &Path::new(&output_dir.clone()).join("risc_proof.json"),
     );
 
     println!("=== Phase 2: Creating compression proof");
@@ -175,17 +311,14 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     serialize_to_file(
         &compression_proof,
-        &Path::new(&cli.output_dir.clone()).join("compression_proof.json"),
+        &Path::new(&output_dir.clone()).join("compression_proof.json"),
     );
 
     println!("=== Phase 3: Creating SNARK proof");
 
     {
         let worker = BellmanWorker::new_with_cpus(4);
 
-        let crs_mons =
-            Crs::<Bn256, CrsForMonomialForm>::crs_42(1 << L1_VERIFIER_DOMAIN_SIZE_LOG, &worker);
-
         let (snark_setup, snark_wrapper_vk) =
             get_snark_wrapper_setup(compression_vk.clone(), &crs_mons, &worker);
 
@@ -203,11 +336,11 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
         serialize_to_file(
             &snark_wrapper_proof,
-            &Path::new(&cli.output_dir.clone()).join("snark_proof.json"),
+            &Path::new(&output_dir.clone()).join("snark_proof.json"),
         );
         serialize_to_file(
             &snark_wrapper_vk,
-            &Path::new(&cli.output_dir.clone()).join("snark_vk.json"),
+            &Path::new(&output_dir.clone()).join("snark_vk.json"),
         );
     }
 

src/wrapper_inner_verifier/imports/circuit_layout.rs

@@ -378,7 +378,7 @@ const COMPILED_WITNESS_LAYOUT: CompiledWitnessSubtree<Mersenne31Field> = Compile
 };
 const COMPILED_MEMORY_LAYOUT: CompiledMemorySubtree<'static> = CompiledMemorySubtree {
     shuffle_ram_inits_and_teardowns: Some(ShuffleRamInitAndTeardownLayout {
-        lazy_init_addesses_columns: ColumnSet::<2usize> {
+        lazy_init_addresses_columns: ColumnSet::<2usize> {
             start: 0usize,
             num_elements: 1usize,
         },