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