1use crate::utils::inputs::{random_f_bigint, serialize_phrase, serialize_username};
  2use ark_circom::{
  3    circom::{r1cs_reader, R1CS},
  4    WitnessCalculator,
  5};
  6use ark_ff::{BigInteger, PrimeField};
  7use color_eyre::Result;
  8use num_bigint::{BigInt, Sign};
  9use sonobe::Error as SonobeError;
 10use std::{fs::File, io::BufReader, marker::PhantomData, path::PathBuf};
 11
 12#[derive(Clone, Debug)]
 13pub struct CircomPrivateInput {
 14    pub phrase: Option<String>,
 15    pub usernames: [Option<String>; 2],
 16    pub auth_secrets: [Option<BigInt>; 2],
 17    pub chaff: bool,
 18}
 19
 20impl CircomPrivateInput {
 21
 22    /**
 23     * Creates empty inputs
 24     * 
 25     * @param chaff - if true, should compute random vars for chaff in circuit
 26     */
 27    pub fn empty(chaff: bool) -> Self {
 28        Self {
 29            phrase: None,
 30            usernames: [None, None],
 31            auth_secrets: [None, None],
 32            chaff,
 33        }
 34    }
 35
 36    pub fn uninitialized(&self) -> bool {
 37        let not_chaff =self.phrase.is_none()
 38            && self.usernames.iter().all(|u| u.is_none())
 39            && self.auth_secrets.iter().all(|a| a.is_none());
 40        not_chaff && !self.chaff
 41    }
 42}
 43
 44// Wrapper for circom functionalities (extract R1CS and witness)
 45#[derive(Clone, Debug)]
 46pub struct CircomWrapper<F: PrimeField> {
 47    r1cs_path: PathBuf,
 48    wc_path: PathBuf,
 49    _marker: PhantomData<F>,
 50}
 51
 52impl<F: PrimeField> CircomWrapper<F> {
 53    // creates a new instance of the wrapper with filepaths
 54    pub fn new(r1cs_path: PathBuf, wc_path: PathBuf) -> Self {
 55        Self {
 56            r1cs_path,
 57            wc_path,
 58            _marker: PhantomData,
 59        }
 60    }
 61
 62    /**
 63     * Marshals the private inputs into the format expected by circom
 64     *
 65     * @param inputs - the private inputs
 66     * @return - the marshalled inputs
 67     */
 68    pub fn marshal_private_inputs(inputs: &CircomPrivateInput) -> [(String, Vec<BigInt>); 3] {
 69        // handle phrase presence (if not present infer chaff)
 70        let phrase = match &inputs.phrase {
 71            Some(phrase) => serialize_phrase(&phrase).unwrap().to_vec(),
 72            None => (0..6)
 73                .map(|_| random_f_bigint::<F>())
 74                .collect::<Vec<BigInt>>(),
 75        };
 76
 77        // determine inputs: first step ([0] = None), Nth step ([1] = Some), and chaff ([2] = None)
 78        // marshal usernames
 79        let usernames = match inputs.usernames[0] {
 80            Some(_) => inputs
 81                .usernames
 82                .iter()
 83                .map(|u| serialize_username(&u.clone().unwrap()).unwrap())
 84                .collect::<Vec<BigInt>>(),
 85            None => match &inputs.usernames[1] {
 86                Some(username) => vec![BigInt::from(0), serialize_username(&username).unwrap()],
 87                None => vec![random_f_bigint::<F>(), random_f_bigint::<F>()],
 88            },
 89        };
 90
 91        // marshal auth secrets
 92        let auth_sec = match inputs.auth_secrets[0] {
 93            Some(_) => inputs
 94                .auth_secrets
 95                .iter()
 96                .map(|a| a.clone().unwrap())
 97                .collect::<Vec<BigInt>>(),
 98            None => match &inputs.auth_secrets[1] {
 99                Some(auth_secret) => vec![BigInt::from(0), auth_secret.clone()],
100                None => vec![random_f_bigint::<F>(), random_f_bigint::<F>()],
101            },
102        };
103
104        // label the inputs for circom
105        [
106            ("phrase".to_string(), phrase),
107            ("usernames".to_string(), usernames),
108            ("auth_secrets".to_string(), auth_sec),
109        ]
110    }
111
112    // aggregated function to obtain r1cs and witness from circom
113    pub fn extract_r1cs_and_witness(
114        &self,
115        inputs: &[(String, Vec<BigInt>)],
116    ) -> Result<(R1CS<F>, Option<Vec<F>>), SonobeError> {
117        // extract R1CS
118        let file = File::open(&self.r1cs_path)?;
119        let reader = BufReader::new(file);
120        let r1cs_file = r1cs_reader::R1CSFile::<F>::new(reader)?;
121        let r1cs = r1cs_reader::R1CS::<F>::from(r1cs_file);
122
123        // extract witness vector
124        let witness_vec = self.extract_witness(inputs)?;
125
126        Ok((r1cs, Some(witness_vec)))
127    }
128
129    pub fn extract_witness(&self, inputs: &[(String, Vec<BigInt>)]) -> Result<Vec<F>, SonobeError> {
130        let witness_bigint = self.calculate_witness(inputs)?;
131        witness_bigint
132            .iter()
133            .map(|bigint| {
134                Self::num_bigint_to_ark_bigint(bigint)
135                    .and_then(|ark_bigint| {
136                        F::from_bigint(ark_bigint).ok_or_else(|| {
137                            SonobeError::Other("Could not get F from bigint".to_string())
138                        })
139                    })
140            })
141            .collect()
142    }
143
144    pub fn calculate_witness(
145        &self,
146        inputs: &[(String, Vec<BigInt>)],
147    ) -> Result<Vec<BigInt>, SonobeError> {
148        let mut calculator = WitnessCalculator::new(&self.wc_path).map_err(|e| {
149            SonobeError::WitnessCalculationError(format!(
150                "Failed to create WitnessCalculator: {}",
151                e
152            ))
153        })?;
154        calculator
155            .calculate_witness(inputs.iter().cloned(), true)
156            .map_err(|e| {
157                SonobeError::WitnessCalculationError(format!("Failed to calculate witness: {}", e))
158            })
159    }
160
161    pub fn num_bigint_to_ark_bigint(value: &BigInt) -> Result<F::BigInt, SonobeError> {
162        let big_uint = value
163            .to_biguint()
164            .ok_or_else(|| SonobeError::BigIntConversionError("BigInt is negative".to_string()))?;
165        F::BigInt::try_from(big_uint).map_err(|_| {
166            SonobeError::BigIntConversionError(
167                "Failed to convert to Primefield::BigInt".to_string(),
168            )
169        })
170    }
171
172    pub fn ark_primefield_to_num_bigint(value: F) -> BigInt {
173        let primefield_bigint: F::BigInt = value.into_bigint();
174        let bytes = primefield_bigint.to_bytes_be();
175        BigInt::from_bytes_be(Sign::Plus, &bytes)
176    }
177}