ark_zkey/
lib.rs

1use ark_bn254::{Bn254, Fr};
2use ark_circom::read_zkey;
3//use ark_ec::pairing::Pairing;
4use ark_ff::Field;
5use ark_groth16::ProvingKey;
6//use ark_groth16::VerifyingKey;
7use ark_relations::r1cs::ConstraintMatrices;
8use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
9use color_eyre::eyre::{Result, WrapErr};
10use memmap2::Mmap;
11use std::fs::File;
12//use std::io::Cursor;
13//use std::io::{Read,self};
14use std::io::BufReader;
15use std::path::PathBuf;
16use std::time::Instant;
17
18#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)]
19pub struct SerializableProvingKey(pub ProvingKey<Bn254>);
20
21#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)]
22pub struct SerializableMatrix<F: Field> {
23    pub data: Vec<Vec<(F, usize)>>,
24}
25
26#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)]
27pub struct SerializableConstraintMatrices<F: Field> {
28    pub num_instance_variables: usize,
29    pub num_witness_variables: usize,
30    pub num_constraints: usize,
31    pub a_num_non_zero: usize,
32    pub b_num_non_zero: usize,
33    pub c_num_non_zero: usize,
34    pub a: SerializableMatrix<F>,
35    pub b: SerializableMatrix<F>,
36    pub c: SerializableMatrix<F>,
37}
38
39impl<F: Field> From<Vec<Vec<(F, usize)>>> for SerializableMatrix<F> {
40    fn from(matrix: Vec<Vec<(F, usize)>>) -> Self {
41        SerializableMatrix { data: matrix }
42    }
43}
44
45impl<F: Field> From<SerializableMatrix<F>> for Vec<Vec<(F, usize)>> {
46    fn from(serializable_matrix: SerializableMatrix<F>) -> Self {
47        serializable_matrix.data
48    }
49}
50
51pub fn serialize_proving_key(pk: &SerializableProvingKey) -> Vec<u8> {
52    let mut serialized_data = Vec::new();
53    pk.serialize_compressed(&mut serialized_data)
54        .expect("Serialization failed");
55    serialized_data
56}
57
58pub fn deserialize_proving_key(data: Vec<u8>) -> SerializableProvingKey {
59    SerializableProvingKey::deserialize_compressed_unchecked(&mut &data[..])
60        .expect("Deserialization failed")
61}
62
63pub fn read_arkzkey(arkzkey_path: &str) -> Result<(ProvingKey<Bn254>, ConstraintMatrices<Fr>)> {
64    let now = std::time::Instant::now();
65    let arkzkey_file_path = PathBuf::from(arkzkey_path);
66    let arkzkey_file = File::open(arkzkey_file_path).wrap_err("Failed to open arkzkey file")?;
67    println!("Time to open arkzkey file: {:?}", now.elapsed());
68
69    //let mut buf_reader = BufReader::new(arkzkey_file);
70
71    // Using mmap
72    let now = std::time::Instant::now();
73    let mmap = unsafe { Mmap::map(&arkzkey_file)? };
74    let mut cursor = std::io::Cursor::new(mmap);
75    println!("Time to mmap arkzkey: {:?}", now.elapsed());
76
77    // Was &mut buf_reader
78    let now = std::time::Instant::now();
79    let serialized_proving_key =
80        SerializableProvingKey::deserialize_compressed_unchecked(&mut cursor)
81            .wrap_err("Failed to deserialize proving key")?;
82    println!("Time to deserialize proving key: {:?}", now.elapsed());
83
84    let now = std::time::Instant::now();
85    let serialized_constraint_matrices =
86        SerializableConstraintMatrices::deserialize_compressed_unchecked(&mut cursor)
87            .wrap_err("Failed to deserialize constraint matrices")?;
88    println!("Time to deserialize matrices: {:?}", now.elapsed());
89
90    let proving_key: ProvingKey<Bn254> = serialized_proving_key.0;
91    let constraint_matrices: ConstraintMatrices<Fr> = ConstraintMatrices {
92        num_instance_variables: serialized_constraint_matrices.num_instance_variables,
93        num_witness_variables: serialized_constraint_matrices.num_witness_variables,
94        num_constraints: serialized_constraint_matrices.num_constraints,
95        a_num_non_zero: serialized_constraint_matrices.a_num_non_zero,
96        b_num_non_zero: serialized_constraint_matrices.b_num_non_zero,
97        c_num_non_zero: serialized_constraint_matrices.c_num_non_zero,
98        a: serialized_constraint_matrices.a.data,
99        b: serialized_constraint_matrices.b.data,
100        c: serialized_constraint_matrices.c.data,
101    };
102
103    Ok((proving_key, constraint_matrices))
104}
105
106pub fn read_arkzkey_from_bytes(
107    arkzkey_bytes: &[u8],
108) -> Result<(ProvingKey<Bn254>, ConstraintMatrices<Fr>)> {
109    let mut cursor = std::io::Cursor::new(arkzkey_bytes);
110
111    let now = std::time::Instant::now();
112    let serialized_proving_key =
113        SerializableProvingKey::deserialize_compressed_unchecked(&mut cursor)
114            .wrap_err("Failed to deserialize proving key")?;
115    println!("Time to deserialize proving key: {:?}", now.elapsed());
116
117    let now = std::time::Instant::now();
118    let serialized_constraint_matrices =
119        SerializableConstraintMatrices::deserialize_compressed_unchecked(&mut cursor)
120            .wrap_err("Failed to deserialize constraint matrices")?;
121    println!("Time to deserialize matrices: {:?}", now.elapsed());
122
123    // Get on right form for API
124    let proving_key: ProvingKey<Bn254> = serialized_proving_key.0;
125    let constraint_matrices: ConstraintMatrices<Fr> = ConstraintMatrices {
126        num_instance_variables: serialized_constraint_matrices.num_instance_variables,
127        num_witness_variables: serialized_constraint_matrices.num_witness_variables,
128        num_constraints: serialized_constraint_matrices.num_constraints,
129        a_num_non_zero: serialized_constraint_matrices.a_num_non_zero,
130        b_num_non_zero: serialized_constraint_matrices.b_num_non_zero,
131        c_num_non_zero: serialized_constraint_matrices.c_num_non_zero,
132        a: serialized_constraint_matrices.a.data,
133        b: serialized_constraint_matrices.b.data,
134        c: serialized_constraint_matrices.c.data,
135    };
136
137    Ok((proving_key, constraint_matrices))
138}
139
140pub fn read_proving_key_and_matrices_from_zkey(
141    zkey_path: &str,
142) -> Result<(SerializableProvingKey, SerializableConstraintMatrices<Fr>)> {
143    println!("Reading zkey from: {}", zkey_path);
144    let now = Instant::now();
145    let zkey_file_path = PathBuf::from(zkey_path);
146    let zkey_file = File::open(zkey_file_path).wrap_err("Failed to open zkey file")?;
147
148    let mut buf_reader = BufReader::new(zkey_file);
149
150    let (proving_key, matrices) =
151        read_zkey(&mut buf_reader).wrap_err("Failed to read zkey file")?;
152    println!("Time to read zkey: {:?}", now.elapsed());
153
154    println!("Serializing proving key and constraint matrices");
155    let now = Instant::now();
156    let serializable_proving_key = SerializableProvingKey(proving_key);
157    let serializable_constrain_matrices = SerializableConstraintMatrices {
158        num_instance_variables: matrices.num_instance_variables,
159        num_witness_variables: matrices.num_witness_variables,
160        num_constraints: matrices.num_constraints,
161        a_num_non_zero: matrices.a_num_non_zero,
162        b_num_non_zero: matrices.b_num_non_zero,
163        c_num_non_zero: matrices.c_num_non_zero,
164        a: SerializableMatrix { data: matrices.a },
165        b: SerializableMatrix { data: matrices.b },
166        c: SerializableMatrix { data: matrices.c },
167    };
168    println!(
169        "Time to serialize proving key and constraint matrices: {:?}",
170        now.elapsed()
171    );
172
173    Ok((serializable_proving_key, serializable_constrain_matrices))
174}
175
176pub fn convert_zkey(
177    proving_key: SerializableProvingKey,
178    constraint_matrices: SerializableConstraintMatrices<Fr>,
179    arkzkey_path: &str,
180) -> Result<()> {
181    let arkzkey_file_path = PathBuf::from(arkzkey_path);
182
183    let serialized_path = PathBuf::from(arkzkey_file_path);
184
185    let mut file =
186        File::create(&serialized_path).wrap_err("Failed to create serialized proving key file")?;
187
188    proving_key
189        .serialize_compressed(&mut file)
190        .wrap_err("Failed to serialize proving key")?;
191
192    constraint_matrices
193        .serialize_compressed(&mut file)
194        .wrap_err("Failed to serialize constraint matrices")?;
195
196    Ok(())
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use std::time::Instant;
203
204    fn test_circuit_serialization_deserialization(zkey_path: &str) -> Result<()> {
205        let arkzkey_path = zkey_path.replace(".zkey", ".arkzkey");
206
207        let (original_proving_key, original_constraint_matrices) =
208            read_proving_key_and_matrices_from_zkey(zkey_path)?;
209
210        println!("[build] Writing arkzkey to: {}", arkzkey_path);
211        let now = Instant::now();
212        convert_zkey(
213            original_proving_key.clone(),
214            original_constraint_matrices.clone(),
215            &arkzkey_path,
216        )?;
217        println!("[build] Time to write arkzkey: {:?}", now.elapsed());
218
219        println!("Reading arkzkey from: {}", arkzkey_path);
220        let now = Instant::now();
221        let (deserialized_proving_key, deserialized_constraint_matrices) =
222            read_arkzkey(&arkzkey_path)?;
223        println!("Time to read arkzkey: {:?}", now.elapsed());
224
225        assert_eq!(
226            original_proving_key.0, deserialized_proving_key,
227            "Original and deserialized proving keys do not match"
228        );
229
230        let original_deserialized_constraint_matrices: ConstraintMatrices<Fr> =
231            ConstraintMatrices {
232                num_instance_variables: original_constraint_matrices.num_instance_variables,
233                num_witness_variables: original_constraint_matrices.num_witness_variables,
234                num_constraints: original_constraint_matrices.num_constraints,
235                a_num_non_zero: original_constraint_matrices.a_num_non_zero,
236                b_num_non_zero: original_constraint_matrices.b_num_non_zero,
237                c_num_non_zero: original_constraint_matrices.c_num_non_zero,
238                a: original_constraint_matrices.a.data,
239                b: original_constraint_matrices.b.data,
240                c: original_constraint_matrices.c.data,
241            };
242        assert_eq!(
243            original_deserialized_constraint_matrices, deserialized_constraint_matrices,
244            "Original and deserialized constraint matrices do not match"
245        );
246
247        Ok(())
248    }
249
250    #[test]
251    fn test_multiplier2_serialization_deserialization() -> Result<()> {
252        test_circuit_serialization_deserialization("./test-vectors/multiplier2_final.zkey")
253    }
254
255    #[test]
256    fn test_keccak256_serialization_deserialization() -> Result<()> {
257        test_circuit_serialization_deserialization("./test-vectors/keccak256_256_test_final.zkey")
258    }
259
260    #[test]
261    #[ignore = "rsa_final.zkey is too large to run on CI"]
262    fn test_rsa_serialization_deserialization() -> Result<()> {
263        test_circuit_serialization_deserialization("./test-vectors/rsa_final.zkey")
264    }
265}
ark_zkey/lib.rs

ark_zkey/
lib.rs
