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