1use ark_circom::{
  2    circom::{r1cs_reader, R1CS},
  3    WitnessCalculator,
  4};
  5use ark_ff::{BigInteger, PrimeField};
  6use ark_serialize::Read;
  7use color_eyre::Result;
  8use num_bigint::{BigInt, Sign};
  9use std::{fs::File, io::Cursor, marker::PhantomData, path::PathBuf};
 10
 11use crate::{utils::PathOrBin, Error};
 12
 13// A struct that wraps Circom functionalities, allowing for extraction of R1CS and witnesses
 14// based on file paths to Circom's .r1cs and .wasm.
 15#[derive(Clone, Debug)]
 16pub struct CircomWrapper<F: PrimeField> {
 17    r1csfile_bytes: Vec<u8>,
 18    wasmfile_bytes: Vec<u8>,
 19    _marker: PhantomData<F>,
 20}
 21
 22impl<F: PrimeField> CircomWrapper<F> {
 23    // Creates a new instance of the CircomWrapper with the file paths.
 24    pub fn new(r1cs: PathOrBin, wasm: PathOrBin) -> Result<Self, Error> {
 25        match (r1cs, wasm) {
 26            (PathOrBin::Path(r1cs_path), PathOrBin::Path(wasm_path)) => {
 27                Self::new_from_path(r1cs_path, wasm_path)
 28            }
 29            (PathOrBin::Bin(r1cs_bin), PathOrBin::Bin(wasm_bin)) => Ok(Self {
 30                r1csfile_bytes: r1cs_bin,
 31                wasmfile_bytes: wasm_bin,
 32                _marker: PhantomData,
 33            }),
 34            _ => unreachable!("You should pass the same enum branch for both inputs"),
 35        }
 36    }
 37
 38    // Creates a new instance of the CircomWrapper with the file paths.
 39    fn new_from_path(r1cs_file_path: PathBuf, wasm_file_path: PathBuf) -> Result<Self, Error> {
 40        let mut file = File::open(r1cs_file_path)?;
 41        let metadata = File::metadata(&file)?;
 42        let mut r1csfile_bytes = vec![0; metadata.len() as usize];
 43        file.read_exact(&mut r1csfile_bytes)?;
 44
 45        let mut file = File::open(wasm_file_path)?;
 46        let metadata = File::metadata(&file)?;
 47        let mut wasmfile_bytes = vec![0; metadata.len() as usize];
 48        file.read_exact(&mut wasmfile_bytes)?;
 49
 50        Ok(CircomWrapper {
 51            r1csfile_bytes,
 52            wasmfile_bytes,
 53            _marker: PhantomData,
 54        })
 55    }
 56
 57    // Aggregated function to obtain R1CS and witness from Circom.
 58    pub fn extract_r1cs_and_witness(
 59        &self,
 60        inputs: &[(String, Vec<BigInt>)],
 61    ) -> Result<(R1CS<F>, Option<Vec<F>>), Error> {
 62        // Extracts the R1CS
 63        let r1cs_file = r1cs_reader::R1CSFile::<F>::new(Cursor::new(&self.r1csfile_bytes))?;
 64        let r1cs = r1cs_reader::R1CS::<F>::from(r1cs_file);
 65
 66        // Extracts the witness vector
 67        let witness_vec = self.extract_witness(inputs)?;
 68
 69        Ok((r1cs, Some(witness_vec)))
 70    }
 71
 72    pub fn extract_r1cs(&self) -> Result<R1CS<F>, Error> {
 73        let r1cs_file = r1cs_reader::R1CSFile::<F>::new(Cursor::new(&self.r1csfile_bytes))?;
 74        let mut r1cs = r1cs_reader::R1CS::<F>::from(r1cs_file);
 75        r1cs.wire_mapping = None;
 76        Ok(r1cs)
 77    }
 78
 79    // Extracts the witness vector as a vector of PrimeField elements.
 80    pub fn extract_witness(&self, inputs: &[(String, Vec<BigInt>)]) -> Result<Vec<F>, Error> {
 81        let witness_bigint = self.calculate_witness(inputs)?;
 82
 83        witness_bigint
 84            .iter()
 85            .map(|big_int| {
 86                self.num_bigint_to_ark_bigint(big_int)
 87                    .and_then(|ark_big_int| {
 88                        F::from_bigint(ark_big_int)
 89                            .ok_or_else(|| Error::Other("could not get F from bigint".to_string()))
 90                    })
 91            })
 92            .collect()
 93    }
 94
 95    // Calculates the witness given the Wasm filepath and inputs.
 96    pub fn calculate_witness(
 97        &self,
 98        inputs: &[(String, Vec<BigInt>)],
 99    ) -> Result<Vec<BigInt>, Error> {
100        let mut calculator = WitnessCalculator::from_binary(&self.wasmfile_bytes).map_err(|e| {
101            Error::WitnessCalculationError(format!("Failed to create WitnessCalculator: {}", e))
102        })?;
103        calculator
104            .calculate_witness(inputs.iter().cloned(), true)
105            .map_err(|e| {
106                Error::WitnessCalculationError(format!("Failed to calculate witness: {}", e))
107            })
108    }
109
110    // Converts a num_bigint::BigInt to a PrimeField::BigInt.
111    pub fn num_bigint_to_ark_bigint(&self, value: &BigInt) -> Result<F::BigInt, Error> {
112        let big_uint = value
113            .to_biguint()
114            .ok_or_else(|| Error::BigIntConversionError("BigInt is negative".to_string()))?;
115        F::BigInt::try_from(big_uint).map_err(|_| {
116            Error::BigIntConversionError("Failed to convert to PrimeField::BigInt".to_string())
117        })
118    }
119
120    // Converts a PrimeField element to a num_bigint::BigInt representation.
121    pub fn ark_primefield_to_num_bigint(&self, value: F) -> BigInt {
122        let primefield_bigint: F::BigInt = value.into_bigint();
123        let bytes = primefield_bigint.to_bytes_be();
124        BigInt::from_bytes_be(Sign::Plus, &bytes)
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use ark_bn254::Fr;
132    use ark_circom::circom::{CircomBuilder, CircomConfig};
133    use ark_circom::CircomCircuit;
134    use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystem};
135
136    //To generate .r1cs and .wasm files, run the below command in the terminal.
137    //bash ./folding-schemes/src/frontend/circom/test_folder/compile.sh
138
139    // Test the satisfication by using the CircomBuilder of circom-compat
140    #[test]
141    fn test_circombuilder_satisfied() {
142        let cfg = CircomConfig::<Fr>::new(
143            "./src/frontend/circom/test_folder/cubic_circuit_js/cubic_circuit.wasm",
144            "./src/frontend/circom/test_folder/cubic_circuit.r1cs",
145        )
146        .unwrap();
147        let mut builder = CircomBuilder::new(cfg);
148        builder.push_input("ivc_input", 3);
149
150        let circom = builder.build().unwrap();
151        let cs = ConstraintSystem::<Fr>::new_ref();
152        circom.generate_constraints(cs.clone()).unwrap();
153        assert!(cs.is_satisfied().unwrap());
154    }
155
156    // Test the satisfication by using the CircomWrapper
157    #[test]
158    fn test_extract_r1cs_and_witness() {
159        let r1cs_path = PathBuf::from("./src/frontend/circom/test_folder/cubic_circuit.r1cs");
160        let wasm_path =
161            PathBuf::from("./src/frontend/circom/test_folder/cubic_circuit_js/cubic_circuit.wasm");
162
163        let inputs = vec![("ivc_input".to_string(), vec![BigInt::from(3)])];
164        let wrapper = CircomWrapper::<Fr>::new(r1cs_path.into(), wasm_path.into()).unwrap();
165
166        let (r1cs, witness) = wrapper.extract_r1cs_and_witness(&inputs).unwrap();
167
168        let cs = ConstraintSystem::<Fr>::new_ref();
169
170        let circom_circuit = CircomCircuit {
171            r1cs,
172            witness,
173            public_inputs_indexes: vec![],
174            allocate_inputs_as_witnesses: false,
175        };
176
177        circom_circuit.generate_constraints(cs.clone()).unwrap();
178        assert!(cs.is_satisfied().unwrap());
179    }
180}