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