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
 12pub struct CircomBuilder<E: PairingEngine> {
 13    pub cfg: CircomConfig<E>,
 14    pub inputs: HashMap<String, Vec<BigInt>>,
 15}
 16
 17// Add utils for creating this from files / directly from bytes
 18pub struct CircomConfig<E: PairingEngine> {
 19    pub r1cs: R1CS<E>,
 20    pub wtns: WitnessCalculator,
 21    pub sanity_check: bool,
 22}
 23
 24impl<E: PairingEngine> CircomConfig<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: CircomConfig<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        use ark_ff::{PrimeField, FpParameters};
 79        let modulus = <<E::Fr as PrimeField>::Params as FpParameters>::MODULUS;
 80
 81        // convert it to field elements
 82        use num_traits::Signed;
 83        let witness = witness
 84            .into_iter()
 85            .map(|w| {
 86                let w = if w.sign() == num_bigint::Sign::Minus {
 87                    // Need to negate the witness element if negative
 88                    modulus.into() - w.abs().to_biguint().unwrap()
 89                } else {
 90                    w.to_biguint().unwrap()
 91                };
 92                E::Fr::from(w)
 93            }).collect::<Vec<_>>();
 94        circom.witness = Some(witness);
 95
 96        // sanity check
 97        debug_assert!({
 98            use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystem};
 99            let cs = ConstraintSystem::<E::Fr>::new_ref();
100            circom.clone().generate_constraints(cs.clone()).unwrap();
101            cs.is_satisfied().unwrap()
102        });
103
104        Ok(circom)
105    }
106}