1use ark_ec::PairingEngine;
 2use std::{fs::File, path::Path};
 3
 4use super::{CircomCircuit, R1CS};
 5
 6use num_bigint::BigInt;
 7use std::collections::HashMap;
 8
 9use crate::{circuit::R1CSFile, WitnessCalculator};
10use color_eyre::Result;
11
12pub struct CircomBuilder<E: PairingEngine> {
13    pub cfg: CircuitConfig<E>,
14    pub inputs: HashMap<String, Vec<BigInt>>,
15}
16
17// Add utils for creating this from files / directly from bytes
18pub struct CircuitConfig<E: PairingEngine> {
19    pub r1cs: R1CS<E>,
20    pub wtns: WitnessCalculator,
21    pub sanity_check: bool,
22}
23
24impl<E: PairingEngine> CircuitConfig<E> {
25    pub fn new(wtns: impl AsRef<Path>, r1cs: impl AsRef<Path>) -> Result<Self> {
26        let wtns = WitnessCalculator::new(wtns).unwrap();
27        let reader = File::open(r1cs)?;
28        let r1cs = R1CSFile::new(reader)?.into();
29        Ok(Self {
30            wtns,
31            r1cs,
32            sanity_check: false,
33        })
34    }
35}
36
37impl<E: PairingEngine> CircomBuilder<E> {
38    /// Instantiates a new builder using the provided WitnessGenerator and R1CS files
39    /// for your circuit
40    pub fn new(cfg: CircuitConfig<E>) -> Self {
41        Self {
42            cfg,
43            inputs: HashMap::new(),
44        }
45    }
46
47    /// Pushes a Circom input at the specified name.
48    pub fn push_input<T: Into<BigInt>>(&mut self, name: impl ToString, val: T) {
49        let values = self.inputs.entry(name.to_string()).or_insert_with(Vec::new);
50        values.push(val.into());
51    }
52
53    /// Generates an empty circom circuit with no witness set, to be used for
54    /// generation of the trusted setup parameters
55    pub fn setup(&self) -> CircomCircuit<E> {
56        let mut circom = CircomCircuit {
57            r1cs: self.cfg.r1cs.clone(),
58            witness: None,
59        };
60
61        // Disable the wire mapping
62        circom.r1cs.wire_mapping = None;
63
64        circom
65    }
66
67    /// Creates the circuit populated with the witness corresponding to the previously
68    /// provided inputs
69    pub fn build(mut self) -> Result<CircomCircuit<E>> {
70        let mut circom = self.setup();
71
72        // calculate the witness
73        let witness = self
74            .cfg
75            .wtns
76            .calculate_witness(self.inputs, self.cfg.sanity_check)?;
77
78        // convert it to field elements
79        let witness = witness
80            .into_iter()
81            .map(|w| E::Fr::from(w.to_biguint().unwrap()))
82            .collect::<Vec<_>>();
83        circom.witness = Some(witness);
84
85        // sanity check
86        debug_assert!({
87            use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystem};
88            let cs = ConstraintSystem::<E::Fr>::new_ref();
89            circom.clone().generate_constraints(cs.clone()).unwrap();
90            cs.is_satisfied().unwrap()
91        });
92
93        Ok(circom)
94    }
95}