1use ark_ec::CurveGroup;
  2use ark_ff::{Field, PrimeField};
  3use ark_r1cs_std::{
  4    alloc::{AllocVar, AllocationMode},
  5    eq::EqGadget,
  6    fields::{fp::FpVar, FieldVar},
  7};
  8use ark_relations::r1cs::{Namespace, SynthesisError};
  9use core::{borrow::Borrow, marker::PhantomData};
 10
 11use crate::ccs::r1cs::RelaxedR1CS;
 12use crate::utils::vec::SparseMatrix;
 13use crate::Error;
 14
 15pub type ConstraintF<C> = <<C as CurveGroup>::BaseField as Field>::BasePrimeField;
 16
 17#[derive(Debug, Clone)]
 18pub struct RelaxedR1CSGadget<F: PrimeField> {
 19    _f: PhantomData<F>,
 20}
 21impl<F: PrimeField> RelaxedR1CSGadget<F> {
 22    /// performs the RelaxedR1CS check (Az∘Bz==uCz+E)
 23    pub fn check(rel_r1cs: RelaxedR1CSVar<F>, z: Vec<FpVar<F>>) -> Result<(), Error> {
 24        let Az = mat_vec_mul_sparse(rel_r1cs.A, z.clone());
 25        let Bz = mat_vec_mul_sparse(rel_r1cs.B, z.clone());
 26        let Cz = mat_vec_mul_sparse(rel_r1cs.C, z.clone());
 27        let uCz = vec_scalar_mul(&Cz, &rel_r1cs.u);
 28        let uCzE = vec_add(&uCz, &rel_r1cs.E)?;
 29        let AzBz = hadamard(&Az, &Bz)?;
 30        for i in 0..AzBz.len() {
 31            AzBz[i].enforce_equal(&uCzE[i].clone())?;
 32        }
 33        Ok(())
 34    }
 35}
 36
 37fn mat_vec_mul_sparse<F: PrimeField>(m: SparseMatrixVar<F>, v: Vec<FpVar<F>>) -> Vec<FpVar<F>> {
 38    let mut res = vec![FpVar::<F>::zero(); m.n_rows];
 39    for (row_i, row) in m.coeffs.iter().enumerate() {
 40        for (value, col_i) in row.iter() {
 41            res[row_i] += value * v[*col_i].clone();
 42        }
 43    }
 44    res
 45}
 46pub fn vec_add<F: PrimeField>(
 47    a: &Vec<FpVar<F>>,
 48    b: &Vec<FpVar<F>>,
 49) -> Result<Vec<FpVar<F>>, Error> {
 50    if a.len() != b.len() {
 51        return Err(Error::NotSameLength);
 52    }
 53    let mut r: Vec<FpVar<F>> = vec![FpVar::<F>::zero(); a.len()];
 54    for i in 0..a.len() {
 55        r[i] = a[i].clone() + b[i].clone();
 56    }
 57    Ok(r)
 58}
 59pub fn vec_scalar_mul<F: PrimeField>(vec: &Vec<FpVar<F>>, c: &FpVar<F>) -> Vec<FpVar<F>> {
 60    let mut result = vec![FpVar::<F>::zero(); vec.len()];
 61    for (i, a) in vec.iter().enumerate() {
 62        result[i] = a.clone() * c;
 63    }
 64    result
 65}
 66pub fn hadamard<F: PrimeField>(
 67    a: &Vec<FpVar<F>>,
 68    b: &Vec<FpVar<F>>,
 69) -> Result<Vec<FpVar<F>>, Error> {
 70    if a.len() != b.len() {
 71        return Err(Error::NotSameLength);
 72    }
 73    let mut r: Vec<FpVar<F>> = vec![FpVar::<F>::zero(); a.len()];
 74    for i in 0..a.len() {
 75        r[i] = a[i].clone() * b[i].clone();
 76    }
 77    Ok(r)
 78}
 79
 80#[derive(Debug, Clone)]
 81pub struct SparseMatrixVar<F: PrimeField> {
 82    pub n_rows: usize,
 83    pub n_cols: usize,
 84    // same format as the native SparseMatrix (which follows ark_relations::r1cs::Matrix format
 85    pub coeffs: Vec<Vec<(FpVar<F>, usize)>>,
 86}
 87
 88impl<F> AllocVar<SparseMatrix<F>, F> for SparseMatrixVar<F>
 89where
 90    F: PrimeField,
 91{
 92    fn new_variable<T: Borrow<SparseMatrix<F>>>(
 93        cs: impl Into<Namespace<F>>,
 94        f: impl FnOnce() -> Result<T, SynthesisError>,
 95        mode: AllocationMode,
 96    ) -> Result<Self, SynthesisError> {
 97        f().and_then(|val| {
 98            let cs = cs.into();
 99
100            let mut coeffs: Vec<Vec<(FpVar<F>, usize)>> = Vec::new();
101            for row in val.borrow().coeffs.iter() {
102                let mut rowVar: Vec<(FpVar<F>, usize)> = Vec::new();
103                for &(value, col_i) in row.iter() {
104                    let coeffVar = FpVar::<F>::new_variable(cs.clone(), || Ok(value), mode)?;
105                    rowVar.push((coeffVar, col_i));
106                }
107                coeffs.push(rowVar);
108            }
109
110            Ok(Self {
111                n_rows: val.borrow().n_rows,
112                n_cols: val.borrow().n_cols,
113                coeffs,
114            })
115        })
116    }
117}
118
119#[derive(Debug, Clone)]
120pub struct RelaxedR1CSVar<F: PrimeField> {
121    pub A: SparseMatrixVar<F>,
122    pub B: SparseMatrixVar<F>,
123    pub C: SparseMatrixVar<F>,
124    pub u: FpVar<F>,
125    pub E: Vec<FpVar<F>>,
126}
127
128impl<F> AllocVar<RelaxedR1CS<F>, F> for RelaxedR1CSVar<F>
129where
130    F: PrimeField,
131{
132    fn new_variable<T: Borrow<RelaxedR1CS<F>>>(
133        cs: impl Into<Namespace<F>>,
134        f: impl FnOnce() -> Result<T, SynthesisError>,
135        mode: AllocationMode,
136    ) -> Result<Self, SynthesisError> {
137        f().and_then(|val| {
138            let cs = cs.into();
139
140            let A = SparseMatrixVar::<F>::new_constant(cs.clone(), &val.borrow().A)?;
141            let B = SparseMatrixVar::<F>::new_constant(cs.clone(), &val.borrow().B)?;
142            let C = SparseMatrixVar::<F>::new_constant(cs.clone(), &val.borrow().C)?;
143            let E = Vec::<FpVar<F>>::new_variable(cs.clone(), || Ok(val.borrow().E.clone()), mode)?;
144            let u = FpVar::<F>::new_variable(cs.clone(), || Ok(val.borrow().u), mode)?;
145
146            Ok(Self { A, B, C, E, u })
147        })
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use ark_crypto_primitives::crh::{
155        sha256::{
156            constraints::{Sha256Gadget, UnitVar},
157            Sha256,
158        },
159        CRHScheme, CRHSchemeGadget,
160    };
161    use ark_ff::BigInteger;
162    use ark_pallas::Fr;
163    use ark_r1cs_std::{alloc::AllocVar, bits::uint8::UInt8};
164    use ark_relations::r1cs::{
165        ConstraintSynthesizer, ConstraintSystem, ConstraintSystemRef, SynthesisError,
166    };
167    use ark_std::One;
168
169    use crate::ccs::r1cs::{
170        tests::{get_test_r1cs, get_test_z},
171        R1CS,
172    };
173    use crate::frontend::arkworks::{extract_r1cs_and_z, tests::TestCircuit};
174
175    #[test]
176    fn test_relaxed_r1cs_small_gadget_handcrafted() {
177        let r1cs: R1CS<Fr> = get_test_r1cs();
178        let rel_r1cs = r1cs.relax();
179        let z = get_test_z(3);
180
181        let cs = ConstraintSystem::<Fr>::new_ref();
182
183        let zVar = Vec::<FpVar<Fr>>::new_witness(cs.clone(), || Ok(z)).unwrap();
184        let rel_r1csVar = RelaxedR1CSVar::<Fr>::new_witness(cs.clone(), || Ok(rel_r1cs)).unwrap();
185
186        RelaxedR1CSGadget::<Fr>::check(rel_r1csVar, zVar).unwrap();
187        assert!(cs.is_satisfied().unwrap());
188        dbg!(cs.num_constraints());
189    }
190
191    // gets as input a circuit that implements the ConstraintSynthesizer trait, and that has been
192    // initialized.
193    fn test_relaxed_r1cs_gadget<CS: ConstraintSynthesizer<Fr>>(circuit: CS) {
194        let cs = ConstraintSystem::<Fr>::new_ref();
195
196        circuit.generate_constraints(cs.clone()).unwrap();
197        cs.finalize();
198        assert!(cs.is_satisfied().unwrap());
199
200        // num constraints of the original circuit
201        dbg!(cs.num_constraints());
202
203        let cs = cs.into_inner().unwrap();
204
205        let (r1cs, z) = extract_r1cs_and_z::<Fr>(&cs);
206        r1cs.check_relation(&z).unwrap();
207
208        let relaxed_r1cs = r1cs.relax();
209        relaxed_r1cs.check_relation(&z).unwrap();
210
211        // set new CS for the circuit that checks the RelaxedR1CS of our original circuit
212        let cs = ConstraintSystem::<Fr>::new_ref();
213        // prepare the inputs for our circuit
214        let zVar = Vec::<FpVar<Fr>>::new_witness(cs.clone(), || Ok(z)).unwrap();
215        let rel_r1csVar =
216            RelaxedR1CSVar::<Fr>::new_witness(cs.clone(), || Ok(relaxed_r1cs)).unwrap();
217
218        RelaxedR1CSGadget::<Fr>::check(rel_r1csVar, zVar).unwrap();
219        assert!(cs.is_satisfied().unwrap());
220
221        // num constraints of the circuit that checks the RelaxedR1CS of the original circuit
222        dbg!(cs.num_constraints());
223    }
224
225    #[test]
226    fn test_relaxed_r1cs_small_gadget_arkworks() {
227        let x = Fr::from(5_u32);
228        let y = x * x * x + x + Fr::from(5_u32);
229        let circuit = TestCircuit::<Fr> { x, y };
230        test_relaxed_r1cs_gadget(circuit);
231    }
232
233    struct Sha256TestCircuit<F: PrimeField> {
234        _f: PhantomData<F>,
235        pub x: Vec<u8>,
236        pub y: Vec<u8>,
237    }
238    impl<F: PrimeField> ConstraintSynthesizer<F> for Sha256TestCircuit<F> {
239        fn generate_constraints(self, cs: ConstraintSystemRef<F>) -> Result<(), SynthesisError> {
240            let x = Vec::<UInt8<F>>::new_witness(cs.clone(), || Ok(self.x))?;
241            let y = Vec::<UInt8<F>>::new_input(cs.clone(), || Ok(self.y))?;
242
243            let unitVar = UnitVar::default();
244            let comp_y = <Sha256Gadget<F> as CRHSchemeGadget<Sha256, F>>::evaluate(&unitVar, &x)?;
245            comp_y.0.enforce_equal(&y)?;
246            Ok(())
247        }
248    }
249    #[test]
250    fn test_relaxed_r1cs_medium_gadget_arkworks() {
251        let x = Fr::from(5_u32).into_bigint().to_bytes_le();
252        let y = <Sha256 as CRHScheme>::evaluate(&(), x.clone()).unwrap();
253
254        let circuit = Sha256TestCircuit::<Fr> {
255            _f: PhantomData,
256            x,
257            y,
258        };
259        test_relaxed_r1cs_gadget(circuit);
260    }
261
262    // circuit that has the numer of constraints specified in the `n_constraints` parameter. Note
263    // that the generated circuit will have very sparse matrices, so the resulting constraints
264    // number of the RelaxedR1CS gadget must take that into account.
265    struct CustomTestCircuit<F: PrimeField> {
266        _f: PhantomData<F>,
267        pub n_constraints: usize,
268        pub x: F,
269        pub y: F,
270    }
271    impl<F: PrimeField> ConstraintSynthesizer<F> for CustomTestCircuit<F> {
272        fn generate_constraints(self, cs: ConstraintSystemRef<F>) -> Result<(), SynthesisError> {
273            let x = FpVar::<F>::new_witness(cs.clone(), || Ok(self.x))?;
274            let y = FpVar::<F>::new_input(cs.clone(), || Ok(self.y))?;
275
276            let mut comp_y = FpVar::<F>::new_witness(cs.clone(), || Ok(F::one()))?;
277            for _ in 0..self.n_constraints - 1 {
278                comp_y *= x.clone();
279            }
280
281            comp_y.enforce_equal(&y)?;
282            Ok(())
283        }
284    }
285    #[test]
286    fn test_relaxed_r1cs_custom_circuit() {
287        let n_constraints = 10_000;
288        let x = Fr::from(5_u32);
289        let mut y = Fr::one();
290        for _ in 0..n_constraints - 1 {
291            y *= x;
292        }
293
294        let circuit = CustomTestCircuit::<Fr> {
295            _f: PhantomData,
296            n_constraints,
297            x,
298            y,
299        };
300        test_relaxed_r1cs_gadget(circuit);
301    }
302}