circom_scotia/
reader.rs

1// Copyright (c) 2022 Nalin
2// Copyright (c) Lurk Lab
3// SPDX-License-Identifier: MIT
4//
5// Contributors:
6//
7// - Hanting Zhang (winston@lurk-lab.com)
8//   - Adapted the original work here: https://github.com/nalinbhardwaj/Nova-Scotia/blob/main/src/circom/reader.rs
9
10use anyhow::bail;
11use crypto_bigint::U256;
12use ff::PrimeField;
13use serde::{Deserialize, Serialize};
14use std::collections::BTreeMap;
15use std::collections::HashMap;
16use std::fs::File;
17use std::fs::OpenOptions;
18use std::io::{BufReader, Read, Seek, SeekFrom};
19use std::io::{Error, ErrorKind};
20use std::path::Path;
21
22use byteorder::{LittleEndian, ReadBytesExt};
23use itertools::Itertools;
24
25use crate::r1cs::Constraint;
26use crate::r1cs::R1CS;
27
28#[derive(Serialize, Deserialize)]
29pub(crate) struct CircuitJson {
30    constraints: Vec<Vec<BTreeMap<String, String>>>,
31    #[serde(rename = "nPubInputs")]
32    num_inputs: usize,
33    #[serde(rename = "nOutputs")]
34    num_outputs: usize,
35    #[serde(rename = "nVars")]
36    num_variables: usize,
37}
38
39// R1CSFile's header
40#[allow(dead_code)]
41#[derive(Debug, Default)]
42struct Header {
43    field_size: u32,
44    prime_size: Vec<u8>,
45    n_wires: u32,
46    n_pub_out: u32,
47    n_pub_in: u32,
48    n_prv_in: u32,
49    n_labels: u64,
50    n_constraints: u32,
51}
52
53// R1CSFile parse result
54#[allow(dead_code)]
55#[derive(Debug, Default)]
56pub struct R1CSFile<F: PrimeField> {
57    version: u32,
58    header: Header,
59    constraints: Vec<Constraint<F>>,
60    wire_mapping: Vec<u64>,
61}
62
63/// load witness file by filename with autodetect encoding (bin or json).
64pub(crate) fn load_witness_from_file<Fr: PrimeField>(filename: impl AsRef<Path>) -> Vec<Fr> {
65    if filename.as_ref().ends_with("json") {
66        load_witness_from_json_file::<Fr>(filename)
67    } else {
68        load_witness_from_bin_file::<Fr>(filename)
69    }
70}
71
72/// load witness from bin file by filename
73fn load_witness_from_bin_file<Fr: PrimeField>(filename: impl AsRef<Path>) -> Vec<Fr> {
74    let reader = OpenOptions::new()
75        .read(true)
76        .open(filename)
77        .expect("unable to open.");
78    load_witness_from_bin_reader::<Fr, BufReader<File>>(BufReader::new(reader))
79        .expect("read witness failed")
80}
81
82/// load witness from u8 array by a reader
83fn load_witness_from_bin_reader<Fr: PrimeField, R: Read>(
84    mut reader: R,
85) -> Result<Vec<Fr>, anyhow::Error> {
86    let mut wtns_header = [0u8; 4];
87    reader.read_exact(&mut wtns_header)?;
88    if wtns_header != [119, 116, 110, 115] {
89        // ruby -e 'p "wtns".bytes' => [119, 116, 110, 115]
90        bail!("invalid file header");
91    }
92    let version = reader.read_u32::<LittleEndian>()?;
93    // println!("wtns version {}", version);
94    if version > 2 {
95        bail!("unsupported file version");
96    }
97    let num_sections = reader.read_u32::<LittleEndian>()?;
98    if num_sections != 2 {
99        bail!("invalid num sections");
100    }
101    // read the first section
102    let sec_type = reader.read_u32::<LittleEndian>()?;
103    if sec_type != 1 {
104        bail!("invalid section type");
105    }
106    let sec_size = reader.read_u64::<LittleEndian>()?;
107    if sec_size != 4 + 32 + 4 {
108        bail!("invalid section len")
109    }
110    let field_size = reader.read_u32::<LittleEndian>()?;
111    if field_size != 32 {
112        bail!("invalid field byte size");
113    }
114    let mut prime = vec![0u8; field_size as usize];
115    reader.read_exact(&mut prime)?;
116    // if prime != hex!("010000f093f5e1439170b97948e833285d588181b64550b829a031e1724e6430") {
117    //     bail!("invalid curve prime {:?}", prime);
118    // }
119    let witness_len = reader.read_u32::<LittleEndian>()?;
120    //println!("witness len {}", witness_len);
121    let sec_type = reader.read_u32::<LittleEndian>()?;
122    if sec_type != 2 {
123        bail!("invalid section type");
124    }
125    let sec_size = reader.read_u64::<LittleEndian>()?;
126    if sec_size != u64::from(witness_len * field_size) {
127        bail!("invalid witness section size {}", sec_size);
128    }
129    let mut result = Vec::with_capacity(witness_len as usize);
130    for _ in 0..witness_len {
131        result.push(read_field::<&mut R, Fr>(&mut reader)?);
132    }
133    Ok(result)
134}
135
136/// load witness from json file by filename
137fn load_witness_from_json_file<Fr: PrimeField>(filename: impl AsRef<Path>) -> Vec<Fr> {
138    let reader = OpenOptions::new()
139        .read(true)
140        .open(filename)
141        .expect("unable to open.");
142    load_witness_from_json::<Fr, BufReader<File>>(BufReader::new(reader))
143}
144
145/// load witness from json by a reader
146fn load_witness_from_json<Fr: PrimeField, R: Read>(reader: R) -> Vec<Fr> {
147    let witness: Vec<String> = serde_json::from_reader(reader).expect("unable to read.");
148    witness
149        .into_iter()
150        .map(|x| Fr::from_str_vartime(&x).unwrap())
151        .collect::<Vec<Fr>>()
152}
153
154/// load r1cs from bin file by filename
155fn load_r1cs_from_bin_file<F: PrimeField>(filename: impl AsRef<Path>) -> R1CS<F> {
156    let reader = OpenOptions::new()
157        .read(true)
158        .open(filename.as_ref())
159        .unwrap_or_else(|_| panic!("unable to open {:?}", filename.as_ref()));
160    load_r1cs_from_bin(BufReader::new(reader))
161}
162
163fn read_field<R: Read, Fr: PrimeField>(mut reader: R) -> Result<Fr, Error> {
164    let mut repr = Fr::ZERO.to_repr();
165    for digit in repr.as_mut().iter_mut() {
166        // TODO: may need to reverse order?
167        *digit = reader.read_u8()?;
168    }
169    let fr = Fr::from_repr(repr).unwrap();
170    Ok(fr)
171}
172
173fn read_header<R: Read>(mut reader: R, size: u64, expected_prime: &str) -> Result<Header, Error> {
174    let field_size = reader.read_u32::<LittleEndian>()?;
175
176    if size != 32 + u64::from(field_size) {
177        return Err(Error::new(
178            ErrorKind::InvalidData,
179            "Invalid header section size",
180        ));
181    }
182
183    let mut prime_size = vec![0u8; field_size as usize];
184    reader.read_exact(&mut prime_size)?;
185    let prime = U256::from_le_slice(&prime_size);
186    let prime = &prime.to_string().to_ascii_lowercase();
187
188    if prime != &expected_prime[2..] {
189        // get rid of '0x' in the front
190        return Err(Error::new(
191            ErrorKind::InvalidData,
192            format!("Mismatched prime field. Expected {expected_prime}, read {prime} in the header instead."),
193        ));
194    }
195
196    Ok(Header {
197        field_size,
198        prime_size,
199        n_wires: reader.read_u32::<LittleEndian>()?,
200        n_pub_out: reader.read_u32::<LittleEndian>()?,
201        n_pub_in: reader.read_u32::<LittleEndian>()?,
202        n_prv_in: reader.read_u32::<LittleEndian>()?,
203        n_labels: reader.read_u64::<LittleEndian>()?,
204        n_constraints: reader.read_u32::<LittleEndian>()?,
205    })
206}
207
208fn read_constraint_vec<R: Read, Fr: PrimeField>(
209    mut reader: R,
210    _header: &Header,
211) -> Result<Vec<(usize, Fr)>, Error> {
212    let n_vec = reader.read_u32::<LittleEndian>()? as usize;
213    let mut vec = Vec::with_capacity(n_vec);
214    for _ in 0..n_vec {
215        vec.push((
216            reader.read_u32::<LittleEndian>()? as usize,
217            read_field::<&mut R, Fr>(&mut reader)?,
218        ));
219    }
220    Ok(vec)
221}
222
223fn read_constraints<R: Read, Fr: PrimeField>(
224    mut reader: R,
225    _size: u64,
226    header: &Header,
227) -> Result<Vec<Constraint<Fr>>, Error> {
228    // todo check section size
229    let mut vec = Vec::with_capacity(header.n_constraints as usize);
230    for _ in 0..header.n_constraints {
231        vec.push((
232            read_constraint_vec::<&mut R, Fr>(&mut reader, header)?,
233            read_constraint_vec::<&mut R, Fr>(&mut reader, header)?,
234            read_constraint_vec::<&mut R, Fr>(&mut reader, header)?,
235        ));
236    }
237    Ok(vec)
238}
239
240fn read_map<R: Read>(mut reader: R, size: u64, header: &Header) -> Result<Vec<u64>, Error> {
241    if size != u64::from(header.n_wires) * 8 {
242        return Err(Error::new(
243            ErrorKind::InvalidData,
244            "Invalid map section size",
245        ));
246    }
247    let mut vec = Vec::with_capacity(header.n_wires as usize);
248    for _ in 0..header.n_wires {
249        vec.push(reader.read_u64::<LittleEndian>()?);
250    }
251    if vec[0] != 0 {
252        return Err(Error::new(
253            ErrorKind::InvalidData,
254            "Wire 0 should always be mapped to 0",
255        ));
256    }
257    Ok(vec)
258}
259
260fn from_reader<Fr: PrimeField, R: Read + Seek>(mut reader: R) -> Result<R1CSFile<Fr>, Error> {
261    let mut magic = [0u8; 4];
262    reader.read_exact(&mut magic)?;
263    if magic != [0x72, 0x31, 0x63, 0x73] {
264        // magic = "r1cs"
265        return Err(Error::new(ErrorKind::InvalidData, "Invalid magic number"));
266    }
267
268    let version = reader.read_u32::<LittleEndian>()?;
269    if version != 1 {
270        return Err(Error::new(ErrorKind::InvalidData, "Unsupported version"));
271    }
272
273    let num_sections = reader.read_u32::<LittleEndian>()?;
274
275    // section type -> file offset
276    let mut section_offsets = HashMap::<u32, u64>::new();
277    let mut section_sizes = HashMap::<u32, u64>::new();
278
279    // get file offset of each section
280    for _ in 0..num_sections {
281        let section_type = reader.read_u32::<LittleEndian>()?;
282        let section_size = reader.read_u64::<LittleEndian>()?;
283        let offset = reader.stream_position()?;
284        section_offsets.insert(section_type, offset);
285        section_sizes.insert(section_type, section_size);
286        reader.seek(SeekFrom::Current(section_size as i64))?;
287    }
288
289    let header_type = 1;
290    let constraint_type = 2;
291    let wire2label_type = 3;
292
293    reader.seek(SeekFrom::Start(*section_offsets.get(&header_type).unwrap()))?;
294    let header = read_header(
295        &mut reader,
296        *section_sizes.get(&header_type).unwrap(),
297        Fr::MODULUS,
298    )?;
299    if header.field_size != 32 {
300        return Err(Error::new(
301            ErrorKind::InvalidData,
302            "This parser only supports 32-byte fields",
303        ));
304    }
305    // if header.prime_size != hex!("010000f093f5e1439170b97948e833285d588181b64550b829a031e1724e6430") {
306    //     return Err(Error::new(ErrorKind::InvalidData, "This parser only supports bn256"));
307    // }
308
309    reader.seek(SeekFrom::Start(
310        *section_offsets.get(&constraint_type).unwrap(),
311    ))?;
312    let constraints = read_constraints::<&mut R, Fr>(
313        &mut reader,
314        *section_sizes.get(&constraint_type).unwrap(),
315        &header,
316    )?;
317
318    reader.seek(SeekFrom::Start(
319        *section_offsets.get(&wire2label_type).unwrap(),
320    ))?;
321    let wire_mapping = read_map(
322        &mut reader,
323        *section_sizes.get(&wire2label_type).unwrap(),
324        &header,
325    )?;
326
327    Ok(R1CSFile {
328        version,
329        header,
330        constraints,
331        wire_mapping,
332    })
333}
334
335/// load r1cs from bin by a reader
336fn load_r1cs_from_bin<Fr: PrimeField, R: Read + Seek>(reader: R) -> R1CS<Fr> {
337    let file = from_reader(reader).expect("unable to read.");
338    let num_inputs = (1 + file.header.n_pub_in + file.header.n_pub_out) as usize;
339    let num_variables = file.header.n_wires as usize;
340    let num_aux = num_variables - num_inputs;
341    R1CS {
342        num_aux,
343        num_inputs,
344        num_variables,
345        constraints: file.constraints,
346    }
347}
348
349/// load r1cs file by filename with autodetect encoding (bin or json)
350pub fn load_r1cs<Fr: PrimeField>(filename: impl AsRef<Path>) -> R1CS<Fr> {
351    if filename.as_ref().ends_with("json") {
352        load_r1cs_from_json_file(filename)
353    } else {
354        load_r1cs_from_bin_file(filename)
355    }
356}
357
358/// load r1cs from json file by filename
359fn load_r1cs_from_json_file<Fr: PrimeField>(filename: impl AsRef<Path>) -> R1CS<Fr> {
360    let reader = OpenOptions::new()
361        .read(true)
362        .open(filename)
363        .expect("unable to open.");
364    load_r1cs_from_json(BufReader::new(reader))
365}
366
367/// load r1cs from json by a reader
368fn load_r1cs_from_json<Fr: PrimeField, R: Read>(reader: R) -> R1CS<Fr> {
369    let circuit_json: CircuitJson = serde_json::from_reader(reader).expect("unable to read.");
370
371    let num_inputs = circuit_json.num_inputs + circuit_json.num_outputs + 1;
372    let num_aux = circuit_json.num_variables - num_inputs;
373
374    let convert_constraint = |lc: &BTreeMap<String, String>| {
375        lc.iter()
376            .map(|(index, coeff)| (index.parse().unwrap(), Fr::from_str_vartime(coeff).unwrap()))
377            .collect_vec()
378    };
379
380    let constraints = circuit_json
381        .constraints
382        .iter()
383        .map(|c| {
384            (
385                convert_constraint(&c[0]),
386                convert_constraint(&c[1]),
387                convert_constraint(&c[2]),
388            )
389        })
390        .collect_vec();
391
392    R1CS {
393        num_inputs,
394        num_aux,
395        num_variables: circuit_json.num_variables,
396        constraints,
397    }
398}