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