feat: allow passing custom binary to the wrapper (#12)

Author: mm-zk

Cargo.toml

@@ -14,18 +14,20 @@ name = "cli"
 path = "src/main.rs"
 
 [dependencies]
-boojum = { git = "https://github.com/matter-labs/zksync-crypto.git", package="boojum", branch = "mmzk_0418_without_field" , features = ["log_tracing"]}
+boojum = { git = "https://github.com/matter-labs/zksync-crypto.git", package="boojum", branch = "main" , features = ["log_tracing"]}
 # boojum = { path = "../zksync-crypto/crates/boojum"}
-bellman = { git = "https://github.com/matter-labs/zksync-crypto.git", package="zksync_bellman", branch = "mmzk_0418_without_field" }
+bellman = { git = "https://github.com/matter-labs/zksync-crypto.git", package="zksync_bellman", branch = "main" }
 # bellman = { path = "../zksync-crypto/crates/bellman", package="zksync_bellman"}
-rescue_poseidon = { git = "https://github.com/matter-labs/zksync-crypto.git", package="rescue_poseidon", branch = "mmzk_0418_without_field" }
+rescue_poseidon = { git = "https://github.com/matter-labs/zksync-crypto.git", package="rescue_poseidon", branch = "main" }
 # rescue_poseidon = { path = "../zksync-crypto/crates/rescue-poseidon"}
-snark_wrapper = { git = "https://github.com/matter-labs/zksync-crypto.git", package="snark_wrapper", branch = "mmzk_0418_without_field" }
+snark_wrapper = { git = "https://github.com/matter-labs/zksync-crypto.git", package="snark_wrapper", branch = "main" }
 # snark_wrapper = { path = "../zksync-crypto/crates/snark-wrapper"}
 risc_verifier = {git="https://github.com/matter-labs/air_compiler.git", package="final_reduced_risc_v_machine_verifier", features=["proof_utils"], branch="main"}
 # risc_verifier = {path="../air_compiler/circuit_defs/final_reduced_risc_v_machine/verifier", package="final_reduced_risc_v_machine_verifier", features=["proof_utils"]}
 execution_utils = {git="https://github.com/matter-labs/air_compiler.git", package="execution_utils", branch="main"}
 # execution_utils = {path="../air_compiler/execution_utils"}
+setups = {git="https://github.com/matter-labs/air_compiler.git", package="setups", branch="main"}
+# setups = {path="../../air_compiler/circuit_defs/setups"}
 circuit_mersenne_field = {path="circuit_mersenne_field"}
 serde_json = { version = "*" }
 serde = { version = "1", default-features = false, features = ["derive", "alloc"]}

circuit_mersenne_field/Cargo.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-boojum = { git = "https://github.com/matter-labs/zksync-crypto.git", package="boojum", branch = "mmzk_0418_without_field" }
+boojum = { git = "https://github.com/matter-labs/zksync-crypto.git", package="boojum", branch = "main" }
 # boojum = { path = "../../zksync-crypto/crates/boojum"}
 mersenne_field = {git="https://github.com/matter-labs/air_compiler.git", package="field", branch="main"}
 # mersenne_field = {path="../../air_compiler/field", package="field"}

src/circuits/risc_wrapper.rs

@@ -27,13 +27,11 @@ use circuit_mersenne_field::{
 };
 use std::mem::MaybeUninit;
 
-use crate::wrapper_inner_verifier::imports::{
-    FINAL_RISC_CIRCUIT_AUX_REGISTERS_VALUES, FINAL_RISC_CIRCUIT_END_PARAMS,
-};
-use crate::wrapper_inner_verifier::skeleton::{
-    WrappedProofSkeletonInstance, WrappedQueryValuesInstance,
-};
 use crate::wrapper_inner_verifier::*;
+use crate::wrapper_inner_verifier::{
+    imports::{FINAL_RISC_CIRCUIT_AUX_REGISTERS_VALUES, FINAL_RISC_CIRCUIT_END_PARAMS},
+    skeleton::{WrappedProofSkeletonInstance, WrappedQueryValuesInstance},
+};
 use crate::wrapper_utils::prover_structs::*;
 use crate::wrapper_utils::verifier_traits::{CircuitLeafInclusionVerifier, PlaceholderSource};
 use risc_verifier::blake2s_u32::*;
@@ -62,8 +60,28 @@ pub struct RiscWrapperWitness {
     pub proof: RiscProof,
 }
 
+#[derive(Clone, Debug)]
+pub struct BinaryCommitment {
+    pub end_params: [u32; 8],
+    pub aux_params: [u32; 8],
+}
+
+impl BinaryCommitment {
+    // Uses the binary information that was provided by wrapper generator.
+    // In future, this will be the 'active' boojumos binary, but for now this is an example fibonacci code.
+    pub fn from_default_binary() -> Self {
+        Self {
+            end_params: FINAL_RISC_CIRCUIT_END_PARAMS,
+            aux_params: FINAL_RISC_CIRCUIT_AUX_REGISTERS_VALUES,
+        }
+    }
+}
+
 impl RiscWrapperWitness {
-    pub fn from_full_proof(full_proof: execution_utils::ProgramProof) -> Self {
+    pub fn from_full_proof(
+        full_proof: execution_utils::ProgramProof,
+        binary_commitment: &BinaryCommitment,
+    ) -> Self {
         let execution_utils::ProgramProof {
             base_layer_proofs,
             delegation_proofs,
@@ -76,7 +94,7 @@ impl RiscWrapperWitness {
         assert!(base_layer_proofs.len() == 1);
         assert!(delegation_proofs.is_empty());
         assert!(register_final_values.len() == NUM_REGISTERS);
-        assert_eq!(end_params, FINAL_RISC_CIRCUIT_END_PARAMS);
+        assert_eq!(end_params, binary_commitment.end_params);
 
         assert!(recursion_chain_preimage.is_some());
         let mut result_hasher = Blake2sBufferingTranscript::new();
@@ -87,10 +105,7 @@ impl RiscWrapperWitness {
             recursion_chain_hash.unwrap(),
             result_hasher.finalize_reset().0
         );
-        assert_eq!(
-            recursion_chain_hash.unwrap(),
-            FINAL_RISC_CIRCUIT_AUX_REGISTERS_VALUES
-        );
+        assert_eq!(recursion_chain_hash.unwrap(), binary_commitment.aux_params);
 
         let final_registers_state: Vec<_> = register_final_values
             .into_iter()
@@ -113,6 +128,7 @@ impl RiscWrapperWitness {
 
 pub struct RiscWrapperCircuit<F: SmallField, V: CircuitLeafInclusionVerifier<F>> {
     pub witness: Option<RiscWrapperWitness>,
+    pub binary_commitment: BinaryCommitment,
     _phantom: std::marker::PhantomData<(F, V)>,
 }
 
@@ -199,7 +215,11 @@ impl<F: SmallField, V: CircuitLeafInclusionVerifier<F>> CircuitBuilder<F>
 }
 
 impl<F: SmallField, V: CircuitLeafInclusionVerifier<F>> RiscWrapperCircuit<F, V> {
-    pub fn new(witness: Option<RiscWrapperWitness>, verify_inner_proof: bool) -> Self {
+    pub fn new(
+        witness: Option<RiscWrapperWitness>,
+        verify_inner_proof: bool,
+        binary_commitment: BinaryCommitment,
+    ) -> Self {
         if verify_inner_proof {
             if let Some(witness) = &witness {
                 verify_risc_proof::<V::OutOfCircuitImpl>(&witness.proof);
@@ -210,6 +230,7 @@ impl<F: SmallField, V: CircuitLeafInclusionVerifier<F>> RiscWrapperCircuit<F, V>
 
         Self {
             witness,
+            binary_commitment,
             _phantom: std::marker::PhantomData,
         }
     }
@@ -283,7 +304,13 @@ impl<F: SmallField, V: CircuitLeafInclusionVerifier<F>> RiscWrapperCircuit<F, V>
         let (proof_state, proof_input) =
             crate::wrapper_inner_verifier::verify(cs, skeleton, queries);
 
-        check_proof_state(cs, final_registers_state, &proof_state, &proof_input);
+        check_proof_state(
+            cs,
+            final_registers_state,
+            &proof_state,
+            &proof_input,
+            &self.binary_commitment,
+        );
 
         // we carry registers 10-17 to the next layer - those are the output of the base program
         let output_registers_values: Vec<_> = final_registers_state
@@ -405,6 +432,7 @@ pub(crate) fn check_proof_state<F: SmallField, CS: ConstraintSystem<F>>(
         NUM_AUX_BOUNDARY_VALUES,
     >,
     public_input: &WrappedProofPublicInputs<F, NUM_STATE_ELEMENTS>,
+    binary_commitment: &BinaryCommitment,
 ) {
     let mut transcript = Blake2sWrappedBufferingTranscript::new(cs);
 
@@ -502,7 +530,7 @@ pub(crate) fn check_proof_state<F: SmallField, CS: ConstraintSystem<F>>(
 
     for i in 0..8 {
         let end_params_word = UInt32::from_le_bytes(cs, end_params_output.0[i].inner);
-        let expected_word = UInt32::allocate_constant(cs, FINAL_RISC_CIRCUIT_END_PARAMS[i]);
+        let expected_word = UInt32::allocate_constant(cs, binary_commitment.end_params[i]);
         Num::enforce_equal(cs, &expected_word.into_num(), &end_params_word.into_num());
     }
 
@@ -512,8 +540,7 @@ pub(crate) fn check_proof_state<F: SmallField, CS: ConstraintSystem<F>>(
     for i in 0..8 {
         let aux_register_idx = (i + 18) * 3;
         let aux_register = final_registers_state[aux_register_idx];
-        let expected_word =
-            UInt32::allocate_constant(cs, FINAL_RISC_CIRCUIT_AUX_REGISTERS_VALUES[i]);
+        let expected_word = UInt32::allocate_constant(cs, binary_commitment.aux_params[i]);
         Num::enforce_equal(cs, &expected_word.into_num(), &aux_register.into_num());
     }
 }

src/lib.rs

@@ -31,6 +31,7 @@ use boojum::cs::traits::circuit::CircuitBuilderProxy;
 use boojum::gadgets::recursion::recursive_transcript::CircuitAlgebraicSpongeBasedTranscript;
 use boojum::gadgets::recursion::recursive_tree_hasher::CircuitGoldilocksPoseidon2Sponge;
 use boojum::implementations::poseidon2::Poseidon2Goldilocks;
+pub use boojum::worker::Worker as BoojumWorker;
 use boojum::worker::*;
 use circuits::*;
 use wrapper_utils::verifier_traits::CircuitBlake2sForEverythingVerifier;
@@ -52,8 +53,8 @@ pub type CircuitRiscWrapperTreeHasher = CircuitGoldilocksPoseidon2Sponge;
 pub type RiscWrapperCircuitBuilder = CircuitBuilderProxy<GL, RiscWrapper>;
 
 pub use rescue_poseidon::franklin_crypto::bellman::pairing::bn256::{Bn256, Fr};
-use rescue_poseidon::poseidon2::transcript::Poseidon2Transcript;
 use rescue_poseidon::poseidon2::Poseidon2Sponge;
+use rescue_poseidon::poseidon2::transcript::Poseidon2Transcript;
 
 use snark_wrapper::implementations::poseidon2::tree_hasher::AbsorptionModeReplacement;
 
@@ -73,8 +74,8 @@ 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 snark_wrapper::implementations::poseidon2::transcript::CircuitPoseidon2Transcript;
 use snark_wrapper::implementations::poseidon2::CircuitPoseidon2Sponge;
+use snark_wrapper::implementations::poseidon2::transcript::CircuitPoseidon2Transcript;
 
 use bellman::worker::Worker as BellmanWorker;
 
@@ -90,11 +91,14 @@ pub type SnarkWrapperSetup =
 pub type SnarkWrapperTranscript =
     bellman::plonk::commitments::transcript::keccak_transcript::RollingKeccakTranscript<Fr>;
 
+pub use execution_utils::ProgramProof;
+
 //CircuitAlgebraicSpongeBasedTranscript<GoldilocksField, 8, 12, 4, R>,
 
 // RiscV -> Stark Wrapper
 pub fn get_risc_wrapper_setup(
     worker: &Worker,
+    binary_commitment: BinaryCommitment,
 ) -> (
     FinalizationHintsForProver,
     SetupBaseStorage<GL>,
@@ -105,7 +109,7 @@ pub fn get_risc_wrapper_setup(
     DenseWitnessCopyHint,
 ) {
     let verify_inner_proof: bool = false;
-    let circuit = RiscWrapper::new(None, verify_inner_proof);
+    let circuit = RiscWrapper::new(None, verify_inner_proof, binary_commitment);
 
     let geometry = RiscWrapper::geometry();
     let (max_trace_len, num_vars) = circuit.size_hint();
@@ -153,9 +157,14 @@ pub fn prove_risc_wrapper(
     vars_hint: &DenseVariablesCopyHint,
     witness_hints: &DenseWitnessCopyHint,
     worker: &Worker,
+    binary_commitment: BinaryCommitment,
 ) -> RiscWrapperProof {
     let verify_inner_proof = true;
-    let circuit = RiscWrapper::new(Some(risc_wrapper_witness), verify_inner_proof);
+    let circuit = RiscWrapper::new(
+        Some(risc_wrapper_witness),
+        verify_inner_proof,
+        binary_commitment,
+    );
 
     let geometry = RiscWrapper::geometry();
     let (max_trace_len, num_vars) = circuit.size_hint();

src/main.rs

@@ -1,3 +1,5 @@
+#![feature(allocator_api)]
+use std::alloc::Global;
 /// Tool that takes the riscv proof from boojum 2.0, together with the final value of the
 /// registers - and returns the SNARK proof.
 // Inside, it runs 3 submodules:
@@ -11,6 +13,13 @@ use clap::Parser;
 use bellman::kate_commitment::{Crs, CrsForMonomialForm};
 use bellman::worker::Worker as BellmanWorker;
 
+use execution_utils::{
+    final_recursion_layer_verifier_vk, recursion_layer_no_delegation_verifier_vk,
+    recursion_layer_verifier_vk, universal_circuit_no_delegation_verifier_vk,
+    universal_circuit_verifier_vk,
+};
+use risc_verifier::prover::worker::Worker;
+use zkos_wrapper::circuits::BinaryCommitment;
 use zkos_wrapper::{Bn256, L1_VERIFIER_DOMAIN_SIZE_LOG};
 use zkos_wrapper::{
     circuits::RiscWrapperWitness, get_compression_setup, get_risc_wrapper_setup,
@@ -23,6 +32,12 @@ use zkos_wrapper::{
 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>,
+
     #[arg(short, long)]
     output_dir: String,
 }
@@ -37,15 +52,68 @@ fn deserialize_from_file<T: serde::de::DeserializeOwned>(filename: &str) -> T {
     serde_json::from_reader(src).unwrap()
 }
 
+pub fn create_binary_commitment(
+    binary_path: String,
+    expected_end_params: &[u32; 8],
+) -> BinaryCommitment {
+    let bin = std::fs::read(binary_path).unwrap();
+
+    let worker = Worker::new_with_num_threads(8);
+
+    let expected_final_pc = execution_utils::find_binary_exit_point(&bin);
+    let binary: Vec<u32> = execution_utils::get_padded_binary(&bin);
+
+    let base_params = execution_utils::compute_end_parameters(
+        expected_final_pc,
+        &setups::get_main_riscv_circuit_setup::<Global>(&binary, &worker),
+    );
+
+    // Check which verifier was used.
+    if universal_circuit_no_delegation_verifier_vk().params == *expected_end_params {
+        let layers = vec![
+            [0u32; 8],
+            base_params,
+            universal_circuit_verifier_vk().params,
+            universal_circuit_no_delegation_verifier_vk().params,
+        ];
+        BinaryCommitment {
+            end_params: universal_circuit_no_delegation_verifier_vk().params,
+            aux_params: execution_utils::compute_chain_encoding(layers),
+        }
+    } else if final_recursion_layer_verifier_vk().params == *expected_end_params {
+        let layers = vec![
+            [0u32; 8],
+            base_params,
+            recursion_layer_verifier_vk().params,
+            recursion_layer_no_delegation_verifier_vk().params,
+            final_recursion_layer_verifier_vk().params,
+        ];
+        BinaryCommitment {
+            end_params: final_recursion_layer_verifier_vk().params,
+            aux_params: execution_utils::compute_chain_encoding(layers),
+        }
+    } else {
+        panic!(
+            "Cannot find a verifier for the proof end parameters: {:?}",
+            expected_end_params
+        );
+    }
+}
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let cli = Cli::parse();
 
     println!("=== Phase 1: Creating the Risc wrapper proof");
 
     let worker = boojum::worker::Worker::new_with_num_threads(4);
 
-    let program_proof = deserialize_from_file(&cli.input);
-    let risc_wrapper_witness = RiscWrapperWitness::from_full_proof(program_proof);
+    let program_proof: zkos_wrapper::ProgramProof = deserialize_from_file(&cli.input);
+    let binary_commitment = match cli.input_binary {
+        Some(binary_path) => create_binary_commitment(binary_path, &program_proof.end_params),
+        None => BinaryCommitment::from_default_binary(),
+    };
+    let risc_wrapper_witness =
+        RiscWrapperWitness::from_full_proof(program_proof, &binary_commitment);
 
     let (
         finalization_hint,
@@ -55,7 +123,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         setup_tree,
         vars_hint,
         witness_hints,
-    ) = get_risc_wrapper_setup(&worker);
+    ) = get_risc_wrapper_setup(&worker, binary_commitment.clone());
 
     let risc_wrapper_proof = prove_risc_wrapper(
         risc_wrapper_witness,
@@ -67,6 +135,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         &vars_hint,
         &witness_hints,
         &worker,
+        binary_commitment,
     );
     let is_valid = verify_risc_wrapper_proof(&risc_wrapper_proof, &risc_wrapper_vk);