1use ark_ff::PrimeField;
  2use ark_r1cs_std::{
  3    alloc::{AllocVar, AllocationMode},
  4    fields::nonnative::NonNativeFieldVar,
  5};
  6use ark_relations::r1cs::{Matrix as R1CSMatrix, Namespace, SynthesisError};
  7use core::{borrow::Borrow, marker::PhantomData};
  8
  9pub struct SparseMatrix<F: PrimeField> {
 10    pub n_rows: usize,
 11    pub n_cols: usize,
 12    /// coeffs = R1CSMatrix = Vec<Vec<(F, usize)>>, which contains each row and the F is the value
 13    /// of the coefficient and the usize indicates the column position
 14    pub coeffs: R1CSMatrix<F>,
 15}
 16
 17#[derive(Debug, Clone)]
 18pub struct SparseMatrixVar<F: PrimeField, CF: PrimeField> {
 19    _f: PhantomData<F>,
 20    _cf: PhantomData<CF>,
 21    pub n_rows: usize,
 22    pub n_cols: usize,
 23    // same format as the native SparseMatrix (which follows ark_relations::r1cs::Matrix format
 24    pub coeffs: Vec<Vec<(NonNativeFieldVar<F, CF>, usize)>>,
 25}
 26
 27impl<F, CF> AllocVar<SparseMatrix<F>, CF> for SparseMatrixVar<F, CF>
 28where
 29    F: PrimeField,
 30    CF: PrimeField,
 31{
 32    fn new_variable<T: Borrow<SparseMatrix<F>>>(
 33        cs: impl Into<Namespace<CF>>,
 34        f: impl FnOnce() -> Result<T, SynthesisError>,
 35        mode: AllocationMode,
 36    ) -> Result<Self, SynthesisError> {
 37        f().and_then(|val| {
 38            let cs = cs.into();
 39
 40            let mut coeffs: Vec<Vec<(NonNativeFieldVar<F, CF>, usize)>> = Vec::new();
 41            for row in val.borrow().coeffs.iter() {
 42                let mut rowVar: Vec<(NonNativeFieldVar<F, CF>, usize)> = Vec::new();
 43                for &(value, col_i) in row.iter() {
 44                    let coeffVar =
 45                        NonNativeFieldVar::<F, CF>::new_variable(cs.clone(), || Ok(value), mode)?;
 46                    rowVar.push((coeffVar, col_i));
 47                }
 48                coeffs.push(rowVar);
 49            }
 50
 51            Ok(Self {
 52                _f: PhantomData,
 53                _cf: PhantomData,
 54                n_rows: val.borrow().n_rows,
 55                n_cols: val.borrow().n_cols,
 56                coeffs,
 57            })
 58        })
 59    }
 60}
 61
 62pub fn mat_vec_mul_sparse<F: PrimeField>(M: &SparseMatrix<F>, z: &[F]) -> Vec<F> {
 63    assert_eq!(M.n_cols, z.len());
 64    let mut res = vec![F::zero(); M.n_rows];
 65    for (row_i, row) in M.coeffs.iter().enumerate() {
 66        for &(value, col_i) in row.iter() {
 67            res[row_i] += value * z[col_i];
 68        }
 69    }
 70    res
 71}
 72pub fn dense_matrix_to_sparse<F: PrimeField>(m: Vec<Vec<F>>) -> SparseMatrix<F> {
 73    let mut r = SparseMatrix::<F> {
 74        n_rows: m.len(),
 75        n_cols: m[0].len(),
 76        coeffs: Vec::new(),
 77    };
 78    for m_row in m.iter() {
 79        let mut row: Vec<(F, usize)> = Vec::new();
 80        for (col_i, value) in m_row.iter().enumerate() {
 81            if !value.is_zero() {
 82                row.push((*value, col_i));
 83            }
 84        }
 85        r.coeffs.push(row);
 86    }
 87    r
 88}
 89
 90// just some helpers to define matrices and vectors by hand
 91pub fn to_F_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> SparseMatrix<F> {
 92    dense_matrix_to_sparse(to_F_dense_matrix(M))
 93}
 94pub fn to_F_dense_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> Vec<Vec<F>> {
 95    M.iter()
 96        .map(|m| m.iter().map(|r| F::from(*r as u64)).collect())
 97        .collect()
 98}
 99pub fn to_F_vec<F: PrimeField>(z: Vec<usize>) -> Vec<F> {
100    z.iter().map(|c| F::from(*c as u64)).collect()
101}