1use ark_ec::pairing::Pairing;
  2use ark_relations::r1cs::{
  3    ConstraintSynthesizer, ConstraintSystemRef, LinearCombination, SynthesisError, Variable,
  4};
  5
  6use super::R1CS;
  7
  8use color_eyre::Result;
  9
 10#[derive(Clone, Debug)]
 11pub struct CircomCircuit<E: Pairing> {
 12    pub r1cs: R1CS<E>,
 13    pub witness: Option<Vec<E::ScalarField>>,
 14}
 15
 16impl<E: Pairing> CircomCircuit<E> {
 17    pub fn get_public_inputs(&self) -> Option<Vec<E::ScalarField>> {
 18        match &self.witness {
 19            None => None,
 20            Some(w) => match &self.r1cs.wire_mapping {
 21                None => Some(w[1..self.r1cs.num_inputs].to_vec()),
 22                Some(m) => Some(m[1..self.r1cs.num_inputs].iter().map(|i| w[*i]).collect()),
 23            },
 24        }
 25    }
 26}
 27
 28impl<E: Pairing> ConstraintSynthesizer<E::ScalarField> for CircomCircuit<E> {
 29    fn generate_constraints(
 30        self,
 31        cs: ConstraintSystemRef<E::ScalarField>,
 32    ) -> Result<(), SynthesisError> {
 33        let witness = &self.witness;
 34        let wire_mapping = &self.r1cs.wire_mapping;
 35
 36        // Start from 1 because Arkworks implicitly allocates One for the first input
 37        for i in 1..self.r1cs.num_inputs {
 38            cs.new_input_variable(|| {
 39                Ok(match witness {
 40                    None => E::ScalarField::from(1u32),
 41                    Some(w) => match wire_mapping {
 42                        Some(m) => w[m[i]],
 43                        None => w[i],
 44                    },
 45                })
 46            })?;
 47        }
 48
 49        for i in 0..self.r1cs.num_aux {
 50            cs.new_witness_variable(|| {
 51                Ok(match witness {
 52                    None => E::ScalarField::from(1u32),
 53                    Some(w) => match wire_mapping {
 54                        Some(m) => w[m[i + self.r1cs.num_inputs]],
 55                        None => w[i + self.r1cs.num_inputs],
 56                    },
 57                })
 58            })?;
 59        }
 60
 61        let make_index = |index| {
 62            if index < self.r1cs.num_inputs {
 63                Variable::Instance(index)
 64            } else {
 65                Variable::Witness(index - self.r1cs.num_inputs)
 66            }
 67        };
 68        let make_lc = |lc_data: &[(usize, E::ScalarField)]| {
 69            lc_data.iter().fold(
 70                LinearCombination::<E::ScalarField>::zero(),
 71                |lc: LinearCombination<E::ScalarField>, (index, coeff)| {
 72                    lc + (*coeff, make_index(*index))
 73                },
 74            )
 75        };
 76
 77        for constraint in &self.r1cs.constraints {
 78            cs.enforce_constraint(
 79                make_lc(&constraint.0),
 80                make_lc(&constraint.1),
 81                make_lc(&constraint.2),
 82            )?;
 83        }
 84
 85        Ok(())
 86    }
 87}
 88
 89#[cfg(test)]
 90mod tests {
 91    use super::*;
 92    use crate::{CircomBuilder, CircomConfig};
 93    use ark_bn254::{Bn254, Fr};
 94    use ark_relations::r1cs::ConstraintSystem;
 95
 96    #[test]
 97    fn satisfied() {
 98        let cfg = CircomConfig::<Bn254>::new(
 99            "./test-vectors/mycircuit.wasm",
100            "./test-vectors/mycircuit.r1cs",
101        )
102        .unwrap();
103        let mut builder = CircomBuilder::new(cfg);
104        builder.push_input("a", 3);
105        builder.push_input("b", 11);
106
107        let circom = builder.build().unwrap();
108        let cs = ConstraintSystem::<Fr>::new_ref();
109        circom.generate_constraints(cs.clone()).unwrap();
110        assert!(cs.is_satisfied().unwrap());
111    }
112}