1use ark_ff::PrimeField;
  2use ark_relations::r1cs::ConstraintSystem;
  3use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
  4use ark_std::rand::Rng;
  5
  6use super::Arith;
  7use crate::utils::vec::{hadamard, mat_vec_mul, vec_add, vec_scalar_mul, SparseMatrix};
  8use crate::Error;
  9
 10#[derive(Debug, Clone, Eq, PartialEq, CanonicalSerialize, CanonicalDeserialize)]
 11pub struct R1CS<F: PrimeField> {
 12    pub l: usize, // io len
 13    pub A: SparseMatrix<F>,
 14    pub B: SparseMatrix<F>,
 15    pub C: SparseMatrix<F>,
 16}
 17
 18impl<F: PrimeField> Arith<F> for R1CS<F> {
 19    /// check that a R1CS structure is satisfied by a z vector. Only for testing.
 20    fn check_relation(&self, z: &[F]) -> Result<(), Error> {
 21        let Az = mat_vec_mul(&self.A, z)?;
 22        let Bz = mat_vec_mul(&self.B, z)?;
 23        let Cz = mat_vec_mul(&self.C, z)?;
 24        let AzBz = hadamard(&Az, &Bz)?;
 25        if AzBz != Cz {
 26            return Err(Error::NotSatisfied);
 27        }
 28        Ok(())
 29    }
 30
 31    fn params_to_le_bytes(&self) -> Vec<u8> {
 32        [
 33            self.l.to_le_bytes(),
 34            self.A.n_rows.to_le_bytes(),
 35            self.A.n_cols.to_le_bytes(),
 36        ]
 37        .concat()
 38    }
 39}
 40
 41impl<F: PrimeField> R1CS<F> {
 42    pub fn rand<R: Rng>(rng: &mut R, n_rows: usize, n_cols: usize) -> Self {
 43        Self {
 44            l: 1,
 45            A: SparseMatrix::rand(rng, n_rows, n_cols),
 46            B: SparseMatrix::rand(rng, n_rows, n_cols),
 47            C: SparseMatrix::rand(rng, n_rows, n_cols),
 48        }
 49    }
 50
 51    /// returns a tuple containing (w, x) (witness and public inputs respectively)
 52    pub fn split_z(&self, z: &[F]) -> (Vec<F>, Vec<F>) {
 53        (z[self.l + 1..].to_vec(), z[1..self.l + 1].to_vec())
 54    }
 55
 56    /// converts the R1CS instance into a RelaxedR1CS as described in
 57    /// [Nova](https://eprint.iacr.org/2021/370.pdf) section 4.1.
 58    pub fn relax(self) -> RelaxedR1CS<F> {
 59        RelaxedR1CS::<F> {
 60            l: self.l,
 61            E: vec![F::zero(); self.A.n_rows],
 62            A: self.A,
 63            B: self.B,
 64            C: self.C,
 65            u: F::one(),
 66        }
 67    }
 68}
 69
 70#[derive(Debug, Clone, Eq, PartialEq)]
 71pub struct RelaxedR1CS<F: PrimeField> {
 72    pub l: usize, // io len
 73    pub A: SparseMatrix<F>,
 74    pub B: SparseMatrix<F>,
 75    pub C: SparseMatrix<F>,
 76    pub u: F,
 77    pub E: Vec<F>,
 78}
 79
 80impl<F: PrimeField> RelaxedR1CS<F> {
 81    /// check that a RelaxedR1CS structure is satisfied by a z vector. Only for testing.
 82    pub fn check_relation(&self, z: &[F]) -> Result<(), Error> {
 83        let Az = mat_vec_mul(&self.A, z)?;
 84        let Bz = mat_vec_mul(&self.B, z)?;
 85        let Cz = mat_vec_mul(&self.C, z)?;
 86        let uCz = vec_scalar_mul(&Cz, &self.u);
 87        let uCzE = vec_add(&uCz, &self.E)?;
 88        let AzBz = hadamard(&Az, &Bz)?;
 89        if AzBz != uCzE {
 90            return Err(Error::NotSatisfied);
 91        }
 92
 93        Ok(())
 94    }
 95}
 96
 97/// extracts arkworks ConstraintSystem matrices into crate::utils::vec::SparseMatrix format as R1CS
 98/// struct.
 99pub fn extract_r1cs<F: PrimeField>(cs: &ConstraintSystem<F>) -> R1CS<F> {
100    let m = cs.to_matrices().unwrap();
101
102    let n_rows = cs.num_constraints;
103    let n_cols = cs.num_instance_variables + cs.num_witness_variables; // cs.num_instance_variables already counts the 1
104
105    let A = SparseMatrix::<F> {
106        n_rows,
107        n_cols,
108        coeffs: m.a,
109    };
110    let B = SparseMatrix::<F> {
111        n_rows,
112        n_cols,
113        coeffs: m.b,
114    };
115    let C = SparseMatrix::<F> {
116        n_rows,
117        n_cols,
118        coeffs: m.c,
119    };
120
121    R1CS::<F> {
122        l: cs.num_instance_variables - 1, // -1 to subtract the first '1'
123        A,
124        B,
125        C,
126    }
127}
128
129/// extracts the witness and the public inputs from arkworks ConstraintSystem.
130pub fn extract_w_x<F: PrimeField>(cs: &ConstraintSystem<F>) -> (Vec<F>, Vec<F>) {
131    (
132        cs.witness_assignment.clone(),
133        // skip the first element which is '1'
134        cs.instance_assignment[1..].to_vec(),
135    )
136}
137
138#[cfg(test)]
139pub mod tests {
140    use super::*;
141    use crate::utils::vec::tests::{to_F_matrix, to_F_vec};
142
143    use ark_pallas::Fr;
144
145    pub fn get_test_r1cs<F: PrimeField>() -> R1CS<F> {
146        // R1CS for: x^3 + x + 5 = y (example from article
147        // https://www.vitalik.ca/general/2016/12/10/qap.html )
148        let A = to_F_matrix::<F>(vec![
149            vec![0, 1, 0, 0, 0, 0],
150            vec![0, 0, 0, 1, 0, 0],
151            vec![0, 1, 0, 0, 1, 0],
152            vec![5, 0, 0, 0, 0, 1],
153        ]);
154        let B = to_F_matrix::<F>(vec![
155            vec![0, 1, 0, 0, 0, 0],
156            vec![0, 1, 0, 0, 0, 0],
157            vec![1, 0, 0, 0, 0, 0],
158            vec![1, 0, 0, 0, 0, 0],
159        ]);
160        let C = to_F_matrix::<F>(vec![
161            vec![0, 0, 0, 1, 0, 0],
162            vec![0, 0, 0, 0, 1, 0],
163            vec![0, 0, 0, 0, 0, 1],
164            vec![0, 0, 1, 0, 0, 0],
165        ]);
166
167        R1CS::<F> { l: 1, A, B, C }
168    }
169
170    pub fn get_test_z<F: PrimeField>(input: usize) -> Vec<F> {
171        // z = (1, io, w)
172        to_F_vec(vec![
173            1,
174            input,                             // io
175            input * input * input + input + 5, // x^3 + x + 5
176            input * input,                     // x^2
177            input * input * input,             // x^2 * x
178            input * input * input + input,     // x^3 + x
179        ])
180    }
181
182    pub fn get_test_z_split<F: PrimeField>(input: usize) -> (F, Vec<F>, Vec<F>) {
183        // z = (1, io, w)
184        (
185            F::one(),
186            to_F_vec(vec![
187                input, // io
188            ]),
189            to_F_vec(vec![
190                input * input * input + input + 5, // x^3 + x + 5
191                input * input,                     // x^2
192                input * input * input,             // x^2 * x
193                input * input * input + input,     // x^3 + x
194            ]),
195        )
196    }
197
198    #[test]
199    fn test_check_relation() {
200        let r1cs = get_test_r1cs::<Fr>();
201        let z = get_test_z(5);
202        r1cs.check_relation(&z).unwrap();
203        r1cs.relax().check_relation(&z).unwrap();
204    }
205}