1#![allow(non_snake_case)]
  2#![allow(unused_doc_comments)]
  3#![allow(dead_code)]
  4
  5use ark_ff::PrimeField;
  6use ark_r1cs_std::fields::nonnative::NonNativeFieldVar;
  7use ark_r1cs_std::{alloc::AllocVar, eq::EqGadget, fields::FieldVar, R1CSVar};
  8use ark_relations::r1cs::{ConstraintSynthesizer, ConstraintSystemRef, SynthesisError};
  9use core::marker::PhantomData;
 10use std::ops::Mul;
 11
 12mod utils;
 13use utils::*;
 14
 15/// - F stands for the field that we represent
 16/// - CF stands for the ConstraintField over which we do the operations
 17
 18/// Implements the A * z matrix-vector-product by fixing the combinations of 'z'.
 19fn handcrafted_A_by_z<F: PrimeField, CF: PrimeField>(
 20    cs: ConstraintSystemRef<CF>,
 21    z: Vec<NonNativeFieldVar<F, CF>>,
 22) -> Result<Vec<NonNativeFieldVar<F, CF>>, SynthesisError> {
 23    let five = NonNativeFieldVar::<F, CF>::new_constant(cs.clone(), F::from(5u32))?;
 24    // directly hand-craft the output vector containing the operations in-place:
 25    Ok(vec![
 26        z[1].clone() + five.clone() * z[4].clone(),
 27        z[1].clone() + z[3].clone(),
 28        z[1].clone() + z[4].clone(),
 29        five * z[0].clone() + z[4].clone() + z[5].clone(),
 30    ]
 31    .clone())
 32}
 33
 34/// Implements the A * z matrix-vector-product by doing the sparse matrix by vector algorithm, and
 35/// assuming that the elements of the matrix A are constants of the system.
 36pub fn mat_vec_mul_sparse_gadget<F: PrimeField, CF: PrimeField>(
 37    m: SparseMatrixVar<F, CF>,
 38    v: Vec<NonNativeFieldVar<F, CF>>,
 39) -> Vec<NonNativeFieldVar<F, CF>> {
 40    let mut res = vec![NonNativeFieldVar::<F, CF>::zero(); m.n_rows];
 41    for (row_i, row) in m.coeffs.iter().enumerate() {
 42        for (value, col_i) in row.iter() {
 43            if value.value().unwrap() == F::one() {
 44                res[row_i] += v[*col_i].clone(); // when value==1, no need to multiply by it
 45                continue;
 46            }
 47            res[row_i] += value.clone().mul(&v[*col_i].clone());
 48        }
 49    }
 50    res
 51}
 52
 53/// Circuit that takes as constants the sparse matrix A, and as inputs the vectors z and y. It
 54/// computes the matrix by vector product between A and z, and checks that is equal to y
 55/// (ie. y == A*z)
 56struct MatrixVectorCircuit<F: PrimeField, CF: PrimeField> {
 57    _cf: PhantomData<CF>,
 58    pub A: SparseMatrix<F>,
 59    pub z: Vec<F>,
 60    pub y: Vec<F>,
 61}
 62impl<F: PrimeField, CF: PrimeField> ConstraintSynthesizer<CF> for MatrixVectorCircuit<F, CF> {
 63    fn generate_constraints(self, cs: ConstraintSystemRef<CF>) -> Result<(), SynthesisError> {
 64        // set A as circuit constants
 65        let A = SparseMatrixVar::<F, CF>::new_constant(cs.clone(), self.A)?;
 66        // set z and y as witness (private inputs)
 67        let z: Vec<NonNativeFieldVar<F, CF>> = Vec::new_witness(cs.clone(), || Ok(self.z.clone()))?;
 68        let y: Vec<NonNativeFieldVar<F, CF>> = Vec::new_witness(cs.clone(), || Ok(self.y.clone()))?;
 69
 70        /// The next two lines are the ones that can be swapped to see the number of constraints
 71        /// taken by the two approaches:
 72        let Az = mat_vec_mul_sparse_gadget(A, z);
 73        // let Az = handcrafted_A_by_z(cs, z)?;
 74
 75        Az.enforce_equal(&y)?;
 76        Ok(())
 77    }
 78}
 79
 80#[cfg(test)]
 81mod tests {
 82    use super::*;
 83    use ark_pallas::{Fq, Fr};
 84    use ark_relations::r1cs::ConstraintSystem;
 85
 86    #[test]
 87    fn test_relaxed_r1cs_nonnative_matrix_vector_product() {
 88        let A = to_F_matrix::<Fq>(vec![
 89            vec![0, 1, 0, 0, 5, 0],
 90            vec![0, 1, 0, 1, 0, 0],
 91            vec![0, 1, 0, 0, 1, 0],
 92            vec![5, 0, 0, 0, 1, 1],
 93        ]);
 94        let z = to_F_vec(vec![1, 123, 35, 53, 80, 30]);
 95        let y = mat_vec_mul_sparse(&A, &z); // y = A*z
 96        println!("Matrix of size {} x {}", A.n_rows, A.n_cols);
 97        println!("Vector of size {}", z.len());
 98
 99        println!(
100            "Build the circuit that computes the matrix-vector-product over a non-native field"
101        );
102        let cs = ConstraintSystem::<Fr>::new_ref();
103        let circuit = MatrixVectorCircuit::<Fq, Fr> {
104            _cf: PhantomData,
105            A,
106            z,
107            y,
108        };
109        circuit.generate_constraints(cs.clone()).unwrap();
110        println!("Number of constraints: {}", cs.num_constraints());
111        assert!(cs.is_satisfied().unwrap());
112    }
113}