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::{circom::R1CSFile, witness::WitnessCalculator};
 10use color_eyre::Result;
 11
 12#[derive(Clone, Debug)]
 13pub struct CircomBuilder<E: PairingEngine> {
 14    pub cfg: CircomConfig<E>,
 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<E: PairingEngine> {
 21    pub r1cs: R1CS<E>,
 22    pub wtns: WitnessCalculator,
 23    pub sanity_check: bool,
 24}
 25
 26impl<E: PairingEngine> CircomConfig<E> {
 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<E: PairingEngine> CircomBuilder<E> {
 40    /// Instantiates a new builder using the provided WitnessGenerator and R1CS files
 41    /// for your circuit
 42    pub fn new(cfg: CircomConfig<E>) -> 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<E> {
 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<E>> {
 72        let mut circom = self.setup();
 73
 74        // calculate the witness
 75        let witness = self
 76            .cfg
 77            .wtns
 78            .calculate_witness(self.inputs, self.cfg.sanity_check)?;
 79
 80        use ark_ff::{FpParameters, PrimeField};
 81        let modulus = <<E::Fr as PrimeField>::Params as FpParameters>::MODULUS;
 82
 83        // convert it to field elements
 84        use num_traits::Signed;
 85        let witness = witness
 86            .into_iter()
 87            .map(|w| {
 88                let w = if w.sign() == num_bigint::Sign::Minus {
 89                    // Need to negate the witness element if negative
 90                    modulus.into() - w.abs().to_biguint().unwrap()
 91                } else {
 92                    w.to_biguint().unwrap()
 93                };
 94                E::Fr::from(w)
 95            })
 96            .collect::<Vec<_>>();
 97        circom.witness = Some(witness);
 98
 99        // sanity check
100        debug_assert!({
101            use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystem};
102            let cs = ConstraintSystem::<E::Fr>::new_ref();
103            circom.clone().generate_constraints(cs.clone()).unwrap();
104            cs.is_satisfied().unwrap()
105        });
106
107        Ok(circom)
108    }
109}