Feat/ark zkey

ark-zkey/.gitignore

@@ -0,0 +1,14 @@
+# Generated by Cargo
+# will have compiled files and executables
+debug/
+target/
+
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+# MSVC Windows builds of rustc generate these, which store debugging information
+*.pdb

ark-zkey/Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+name = "ark-zkey"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+# XXX: Shouldn't be necessary, but this way we stay consistent with wasmer version and fix
+# error[E0432]: unresolved import `wasmer` error
+# (likely due to other packages)
+[patch.crates-io]
+# NOTE: Forked wasmer to work around memory limits
+# See https://github.com/wasmerio/wasmer/commit/09c7070
+wasmer = { git = "https://github.com/oskarth/wasmer.git", rev = "09c7070" }
+
+[dependencies]
+color-eyre = "0.6"
+memmap2 = "0.9"
+flame = "0.2"
+flamer = "0.5"
+
+ark-serialize = { version = "=0.4.1", features = ["derive"] }
+ark-bn254 = { version = "=0.4.0" }
+ark-groth16 = { version = "=0.4.0" }
+ark-circom = { git = "https://github.com/arkworks-rs/circom-compat.git" }
+ark-relations = { version = "=0.4.0" }
+ark-ff = { version = "=0.4.1" }

ark-zkey/src/lib.rs

@@ -0,0 +1,208 @@
+use ark_bn254::{Bn254, Fr};
+use ark_circom::read_zkey;
+use ark_ff::Field;
+use ark_groth16::ProvingKey;
+use ark_relations::r1cs::ConstraintMatrices;
+use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
+use color_eyre::eyre::{Result, WrapErr};
+use memmap2::Mmap;
+use std::fs::File;
+use std::io::Cursor;
+use std::io::Read;
+use std::io::{self, BufReader};
+use std::path::PathBuf;
+
+#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)]
+pub struct SerializableProvingKey(pub ProvingKey<Bn254>);
+
+#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)]
+pub struct SerializableMatrix<F: Field> {
+    pub data: Vec<Vec<(F, usize)>>,
+}
+
+#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)]
+pub struct SerializableConstraintMatrices<F: Field> {
+    pub num_instance_variables: usize,
+    pub num_witness_variables: usize,
+    pub num_constraints: usize,
+    pub a_num_non_zero: usize,
+    pub b_num_non_zero: usize,
+    pub c_num_non_zero: usize,
+    pub a: SerializableMatrix<F>,
+    pub b: SerializableMatrix<F>,
+    pub c: SerializableMatrix<F>,
+}
+
+impl<F: Field> From<Vec<Vec<(F, usize)>>> for SerializableMatrix<F> {
+    fn from(matrix: Vec<Vec<(F, usize)>>) -> Self {
+        SerializableMatrix { data: matrix }
+    }
+}
+
+impl<F: Field> From<SerializableMatrix<F>> for Vec<Vec<(F, usize)>> {
+    fn from(serializable_matrix: SerializableMatrix<F>) -> Self {
+        serializable_matrix.data
+    }
+}
+
+pub fn serialize_proving_key(pk: &SerializableProvingKey) -> Vec<u8> {
+    let mut serialized_data = Vec::new();
+    pk.serialize_compressed(&mut serialized_data)
+        .expect("Serialization failed");
+    serialized_data
+}
+
+pub fn deserialize_proving_key(data: Vec<u8>) -> SerializableProvingKey {
+    SerializableProvingKey::deserialize_compressed(&mut &data[..]).expect("Deserialization failed")
+}
+
+pub fn read_arkzkey(
+    arkzkey_path: &str,
+) -> Result<(SerializableProvingKey, SerializableConstraintMatrices<Fr>)> {
+    let now = std::time::Instant::now();
+    let arkzkey_file_path = PathBuf::from(arkzkey_path);
+    let arkzkey_file = File::open(arkzkey_file_path).wrap_err("Failed to open arkzkey file")?;
+    println!("Time to open arkzkey file: {:?}", now.elapsed());
+
+    //let mut buf_reader = BufReader::new(arkzkey_file);
+
+    // Using mmap
+    let now = std::time::Instant::now();
+    let mmap = unsafe { Mmap::map(&arkzkey_file)? };
+    let mut cursor = std::io::Cursor::new(mmap);
+    println!("Time to mmap: {:?}", now.elapsed());
+
+    // Was &mut buf_reader
+    let now = std::time::Instant::now();
+    let proving_key = SerializableProvingKey::deserialize_compressed(&mut cursor)
+        .wrap_err("Failed to deserialize proving key")?;
+    println!("Time to deserialize proving key: {:?}", now.elapsed());
+
+    let now = std::time::Instant::now();
+    let constraint_matrices = SerializableConstraintMatrices::deserialize_compressed(&mut cursor)
+        .wrap_err("Failed to deserialize constraint matrices")?;
+    println!("Time to deserialize matrices: {:?}", now.elapsed());
+
+    Ok((proving_key, constraint_matrices))
+}
+
+pub fn read_proving_key_and_matrices_from_zkey(
+    zkey_path: &str,
+) -> Result<(SerializableProvingKey, SerializableConstraintMatrices<Fr>)> {
+    let zkey_file_path = PathBuf::from(zkey_path);
+    let mut zkey_file = File::open(zkey_file_path).wrap_err("Failed to open zkey file")?;
+
+    let mut buf_reader = BufReader::new(zkey_file);
+
+    let (proving_key, matrices) =
+        read_zkey(&mut buf_reader).wrap_err("Failed to read zkey file")?;
+
+    let serializable_proving_key = SerializableProvingKey(proving_key);
+    let serializable_constrain_matrices = SerializableConstraintMatrices {
+        num_instance_variables: matrices.num_instance_variables,
+        num_witness_variables: matrices.num_witness_variables,
+        num_constraints: matrices.num_constraints,
+        a_num_non_zero: matrices.a_num_non_zero,
+        b_num_non_zero: matrices.b_num_non_zero,
+        c_num_non_zero: matrices.c_num_non_zero,
+        a: SerializableMatrix { data: matrices.a },
+        b: SerializableMatrix { data: matrices.b },
+        c: SerializableMatrix { data: matrices.c },
+    };
+
+    Ok((serializable_proving_key, serializable_constrain_matrices))
+}
+
+pub fn convert_zkey(
+    proving_key: SerializableProvingKey,
+    constraint_matrices: SerializableConstraintMatrices<Fr>,
+    arkzkey_path: &str,
+) -> Result<()> {
+    let arkzkey_file_path = PathBuf::from(arkzkey_path);
+
+    let serialized_path = PathBuf::from(arkzkey_file_path);
+
+    let mut file =
+        File::create(&serialized_path).wrap_err("Failed to create serialized proving key file")?;
+
+    proving_key
+        .serialize_compressed(&mut file)
+        .wrap_err("Failed to serialize proving key")?;
+
+    constraint_matrices
+        .serialize_compressed(&mut file)
+        .wrap_err("Failed to serialize constraint matrices")?;
+
+    Ok(())
+}
+
+fn read_zkey_with_mmap(zkey_path: &str) -> Result<(ProvingKey<Bn254>, ConstraintMatrices<Fr>)> {
+    let file = File::open(zkey_path)?;
+
+    let mmap = unsafe { Mmap::map(&file)? };
+
+    let cursor = Cursor::new(&mmap);
+    let (proving_key, matrices) = read_zkey(&mut cursor.clone())?;
+
+    Ok((proving_key, matrices))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::time::Instant;
+
+    #[test]
+    fn test_serialization_deserialization() -> Result<()> {
+        // multiplier
+        // let dir = "../mopro-core/examples/circom/multiplier2";
+        // let circuit = "multiplier2";
+
+        //keccak256
+        let dir = "../mopro-core/examples/circom/keccak256";
+        let circuit = "keccak256_256_test";
+
+        let zkey_path = format!("{}/target/{}_final.zkey", dir, circuit);
+        let arkzkey_path = format!("{}/target/{}_final.arkzkey", dir, circuit);
+
+        // println!("Reading mmaped zkey from: {}", zkey_path);
+        // let now = Instant::now();
+        // let (original_proving_key, original_constraint_matrices) = read_zkey_with_mmap(&zkey_path)?;
+        // println!("Time to read mmaped zkey: {:?}", now.elapsed());
+
+        println!("Reading zkey from: {}", zkey_path);
+        let now = Instant::now();
+        let (original_proving_key, original_constraint_matrices) =
+            read_proving_key_and_matrices_from_zkey(&zkey_path)?;
+        println!("Time to read zkey: {:?}", now.elapsed());
+
+        println!("Writing arkzkey to: {}", arkzkey_path);
+        let now = Instant::now();
+        convert_zkey(
+            original_proving_key.clone(),
+            original_constraint_matrices.clone(),
+            &arkzkey_path,
+        )?;
+        println!("Time to write zkey: {:?}", now.elapsed());
+
+        println!("Reading arkzkey from: {}", arkzkey_path);
+        let now = Instant::now();
+        let (deserialized_proving_key, deserialized_constraint_matrices) =
+            read_arkzkey(&arkzkey_path)?;
+        println!("Time to read arkzkey: {:?}", now.elapsed());
+
+        assert_eq!(
+            original_proving_key, deserialized_proving_key,
+            "Original and deserialized proving keys do not match"
+        );
+
+        assert_eq!(
+            original_constraint_matrices, deserialized_constraint_matrices,
+            "Original and deserialized constraint matrices do not match"
+        );
+
+        flame::dump_html(&mut std::fs::File::create("flame-graph.html").unwrap()).unwrap();
+
+        Ok(())
+    }
+}

mopro-core/Cargo.toml

@@ -5,6 +5,10 @@ edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[[bin]]
+name = "zarkkey-util"
+path = "zarkkey-util.rs"
+
 [patch.crates-io]
 # NOTE: Forked wasmer to work around memory limits
 # See https://github.com/wasmerio/wasmer/commit/09c7070

mopro-core/src/middleware/circom/mod.rs

@@ -34,6 +34,7 @@ use {
 
 pub mod serialization;
 pub mod utils;
+pub mod zarkkey;
 
 type GrothBn = Groth16<Bn254>;
 
@@ -59,6 +60,7 @@ impl Default for CircomState {
 
 const ZKEY_BYTES: &[u8] = include_bytes!(env!("BUILD_RS_ZKEY_FILE"));
 
+// TODO: This can probably be pre-processed
 static ZKEY: Lazy<(ProvingKey<Bn254>, ConstraintMatrices<Fr>)> = Lazy::new(|| {
     let mut reader = Cursor::new(ZKEY_BYTES);
     read_zkey(&mut reader).expect("Failed to read zkey")
@@ -88,6 +90,14 @@ pub fn initialize(dylib_path: &Path) {
     println!("Initializing zkey took: {:.2?}", now.elapsed());
 }
 
+#[cfg(not(feature = "dylib"))]
+pub fn initialize_zkey() {
+    println!("Initializing zkey");
+    let now = std::time::Instant::now();
+    Lazy::force(&ZKEY);
+    println!("Initializing zkey took: {:.2?}", now.elapsed());
+}
+
 /// Creates a `WitnessCalculator` instance from a dylib file.
 #[cfg(feature = "dylib")]
 fn from_dylib(path: &Path) -> Mutex<WitnessCalculator> {
@@ -512,6 +522,12 @@ mod tests {
             initialize(Path::new(&dylib_path));
         }
 
+        #[cfg(not(feature = "dylib"))]
+        {
+            // Do this separately to remove zkey loading from proof generation time
+            initialize_zkey();
+        }
+
         let input_vec = vec![
             116, 101, 115, 116, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0,

mopro-core/src/middleware/circom/zarkkey.rs

@@ -0,0 +1,5 @@
+// Util to convert zkey to zarkkey
+
+pub fn hello() {
+    println!("Hello, world!");
+}

mopro-core/zarkkey-util.rs

@@ -0,0 +1,7 @@
+// Util to convert zkey to zarkkey
+
+use mopro_core::middleware;
+
+fn main() {
+    middleware::circom::zarkkey::hello();
+}