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            public_inputs_indexes: vec![],
 62            allocate_inputs_as_witnesses: false,
 63        };
 64
 65        // Disable the wire mapping
 66        circom.r1cs.wire_mapping = None;
 67
 68        circom
 69    }
 70
 71    /// Creates the circuit populated with the witness corresponding to the previously
 72    /// provided inputs
 73    pub fn build(mut self) -> Result<CircomCircuit<F>> {
 74        let mut circom = self.setup();
 75
 76        // calculate the witness
 77        let witness = self
 78            .cfg
 79            .wtns
 80            .calculate_witness_element::<F, _>(self.inputs, self.cfg.sanity_check)?;
 81        circom.witness = Some(witness);
 82
 83        // sanity check
 84        debug_assert!({
 85            use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystem};
 86            let cs = ConstraintSystem::<F>::new_ref();
 87            circom.clone().generate_constraints(cs.clone()).unwrap();
 88            let is_satisfied = cs.is_satisfied().unwrap();
 89            if !is_satisfied {
 90                println!(
 91                    "Unsatisfied constraint: {:?}",
 92                    cs.which_is_unsatisfied().unwrap()
 93                );
 94            }
 95
 96            is_satisfied
 97        });
 98
 99        Ok(circom)
100    }
101}