diff --git a/Cargo.lock b/Cargo.lock index 7e2f46d6..9f3dbf0d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -326,6 +326,7 @@ dependencies = [ "flame", "flamer", "memmap2 0.9.4", + "rkyv", ] [[package]] diff --git a/ark-zkey/Cargo.toml b/ark-zkey/Cargo.toml index 6678c4c9..27db30f8 100644 --- a/ark-zkey/Cargo.toml +++ b/ark-zkey/Cargo.toml @@ -14,6 +14,7 @@ color-eyre = "0.6" memmap2 = "0.9" flame = "0.2" flamer = "0.5" +rkyv = "0.7.43" ark-serialize = { version = "=0.4.1", features = ["derive"] } ark-bn254 = { version = "=0.4.0" } diff --git a/ark-zkey/src/lib.rs b/ark-zkey/src/lib.rs index b849f7aa..7048e1f2 100644 --- a/ark-zkey/src/lib.rs +++ b/ark-zkey/src/lib.rs @@ -11,19 +11,52 @@ use memmap2::Mmap; use std::fs::File; //use std::io::Cursor; //use std::io::{Read,self}; +use rkyv::{archived_root, Infallible}; +use rkyv::{Archive, Deserialize, Serialize}; use std::io::BufReader; use std::path::PathBuf; use std::time::Instant; -#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)] +// NOTE: Archive, Serialize, Deserialize traits are for rkyv zero-copy deserialization experiment +// See https://github.com/oskarth/mopro/issues/25 + +// XXX: the trait `Archive` is not implemented for `ProvingKey>` +#[derive( + Archive, + Serialize, + Deserialize, + CanonicalSerialize, + CanonicalDeserialize, + Clone, + Debug, + PartialEq, +)] pub struct SerializableProvingKey(pub ProvingKey); -#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)] +#[derive( + Archive, + Serialize, + Deserialize, + CanonicalSerialize, + CanonicalDeserialize, + Clone, + Debug, + PartialEq, +)] pub struct SerializableMatrix { pub data: Vec>, } -#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)] +#[derive( + Archive, + Serialize, + Deserialize, + CanonicalSerialize, + CanonicalDeserialize, + Clone, + Debug, + PartialEq, +)] pub struct SerializableConstraintMatrices { pub num_instance_variables: usize, pub num_witness_variables: usize, @@ -137,6 +170,23 @@ pub fn read_arkzkey_from_bytes( Ok((proving_key, constraint_matrices)) } +// Function to read the arkzkey file using rkyv for zero-copy deserialization +pub fn read_arkzkey_with_rkyv( + arkzkey_path: &str, +) -> Result<(ProvingKey, ConstraintMatrices), Box> { + let arkzkey_file_path = PathBuf::from(arkzkey_path); + let arkzkey_file = File::open(arkzkey_file_path)?; + + // Using mmap for zero-copy deserialization + let mmap = unsafe { Mmap::map(&arkzkey_file)? }; + let archived_data = unsafe { archived_root::>(&mmap[..]) }; + + // Optionally deserialize if you need to manipulate the data + let constraint_matrices: ConstraintMatrices = archived_data.deserialize(&mut Infallible)?; + + Ok((archived_data.0, constraint_matrices)) +} + pub fn read_proving_key_and_matrices_from_zkey( zkey_path: &str, ) -> Result<(SerializableProvingKey, SerializableConstraintMatrices)> { @@ -247,6 +297,48 @@ mod tests { Ok(()) } + #[test] + fn test_rkyv_serialization_deserialization() -> Result<(), Box> { + let zkey_path = + "../mopro-core/examples/circom/keccak256/target/keccak256_256_test_final.zkey"; + let arkzkey_path = zkey_path.replace(".zkey", ".arkzkey"); + + // Assuming `read_proving_key_and_matrices_from_zkey` reads the original data + let (original_proving_key, original_constraint_matrices) = + read_proving_key_and_matrices_from_zkey(zkey_path)?; + + // Convert to arkzkey if necessary + // Assuming `convert_zkey` serializes the data into the arkzkey format + println!("[build] Writing arkzkey to: {}", arkzkey_path); + let now = Instant::now(); + convert_zkey( + original_proving_key.clone(), + original_constraint_matrices.clone(), + &arkzkey_path, + )?; + println!("[build] Time to write arkzkey: {:?}", now.elapsed()); + + // Read the arkzkey using the rkyv-based function + println!("Reading arkzkey with rkyv from: {}", arkzkey_path); + let now = Instant::now(); + let (deserialized_proving_key, deserialized_constraint_matrices) = + read_arkzkey_with_rkyv(&arkzkey_path)?; + println!("Time to read arkzkey with rkyv: {:?}", now.elapsed()); + + // Compare the original and deserialized data + 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" + ); + + Ok(()) + } + #[test] fn test_multiplier2_serialization_deserialization() -> Result<()> { test_circuit_serialization_deserialization(