1use ark_ff::PrimeField;
  2pub use ark_relations::r1cs::Matrix as R1CSMatrix;
  3use ark_std::cfg_iter;
  4use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
  5
  6#[derive(Clone, Debug, Eq, PartialEq)]
  7pub struct SparseMatrix<F: PrimeField> {
  8    pub n_rows: usize,
  9    pub n_cols: usize,
 10    /// coeffs = R1CSMatrix = Vec<Vec<(F, usize)>>, which contains each row and the F is the value
 11    /// of the coefficient and the usize indicates the column position
 12    pub coeffs: R1CSMatrix<F>,
 13}
 14
 15impl<F: PrimeField> SparseMatrix<F> {
 16    pub fn to_dense(&self) -> Vec<Vec<F>> {
 17        let mut r: Vec<Vec<F>> = vec![vec![F::zero(); self.n_cols]; self.n_rows];
 18        for (row_i, row) in self.coeffs.iter().enumerate() {
 19            for &(value, col_i) in row.iter() {
 20                r[row_i][col_i] = value;
 21            }
 22        }
 23        r
 24    }
 25}
 26
 27pub fn dense_matrix_to_sparse<F: PrimeField>(m: Vec<Vec<F>>) -> SparseMatrix<F> {
 28    let mut r = SparseMatrix::<F> {
 29        n_rows: m.len(),
 30        n_cols: m[0].len(),
 31        coeffs: Vec::new(),
 32    };
 33    for m_row in m.iter() {
 34        let mut row: Vec<(F, usize)> = Vec::new();
 35        for (col_i, value) in m_row.iter().enumerate() {
 36            if !value.is_zero() {
 37                row.push((*value, col_i));
 38            }
 39        }
 40        r.coeffs.push(row);
 41    }
 42    r
 43}
 44
 45pub fn vec_add<F: PrimeField>(a: &[F], b: &[F]) -> Vec<F> {
 46    assert_eq!(a.len(), b.len());
 47    let mut r: Vec<F> = vec![F::zero(); a.len()];
 48    for i in 0..a.len() {
 49        r[i] = a[i] + b[i];
 50    }
 51    r
 52}
 53
 54pub fn vec_sub<F: PrimeField>(a: &[F], b: &[F]) -> Vec<F> {
 55    assert_eq!(a.len(), b.len());
 56    let mut r: Vec<F> = vec![F::zero(); a.len()];
 57    for i in 0..a.len() {
 58        r[i] = a[i] - b[i];
 59    }
 60    r
 61}
 62
 63pub fn vec_scalar_mul<F: PrimeField>(vec: &[F], c: &F) -> Vec<F> {
 64    let mut result = vec![F::zero(); vec.len()];
 65    for (i, a) in vec.iter().enumerate() {
 66        result[i] = *a * c;
 67    }
 68    result
 69}
 70
 71pub fn is_zero_vec<F: PrimeField>(vec: &[F]) -> bool {
 72    for e in vec {
 73        if !e.is_zero() {
 74            return false;
 75        }
 76    }
 77    true
 78}
 79
 80pub fn mat_vec_mul<F: PrimeField>(M: &Vec<Vec<F>>, z: &[F]) -> Vec<F> {
 81    assert!(!M.is_empty());
 82    assert_eq!(M[0].len(), z.len());
 83
 84    let mut r: Vec<F> = vec![F::zero(); M.len()];
 85    for (i, M_i) in M.iter().enumerate() {
 86        for (j, M_ij) in M_i.iter().enumerate() {
 87            r[i] += *M_ij * z[j];
 88        }
 89    }
 90    r
 91}
 92
 93pub fn mat_vec_mul_sparse<F: PrimeField>(matrix: &SparseMatrix<F>, vector: &[F]) -> Vec<F> {
 94    let mut res = vec![F::zero(); matrix.n_rows];
 95    for (row_i, row) in matrix.coeffs.iter().enumerate() {
 96        for &(value, col_i) in row.iter() {
 97            res[row_i] += value * vector[col_i];
 98        }
 99    }
100
101    res
102}
103
104pub fn hadamard<F: PrimeField>(a: &[F], b: &[F]) -> Vec<F> {
105    assert_eq!(a.len(), b.len());
106    cfg_iter!(a).zip(b).map(|(a, b)| *a * b).collect()
107}
108
109#[cfg(test)]
110pub mod tests {
111    use super::*;
112    use ark_pallas::Fr;
113
114    pub fn to_F_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> SparseMatrix<F> {
115        dense_matrix_to_sparse(to_F_dense_matrix(M))
116    }
117    pub fn to_F_dense_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> Vec<Vec<F>> {
118        let mut R: Vec<Vec<F>> = vec![Vec::new(); M.len()];
119        for i in 0..M.len() {
120            R[i] = vec![F::zero(); M[i].len()];
121            for j in 0..M[i].len() {
122                R[i][j] = F::from(M[i][j] as u64);
123            }
124        }
125        R
126    }
127    pub fn to_F_vec<F: PrimeField>(z: Vec<usize>) -> Vec<F> {
128        let mut r: Vec<F> = vec![F::zero(); z.len()];
129        for i in 0..z.len() {
130            r[i] = F::from(z[i] as u64);
131        }
132        r
133    }
134
135    #[test]
136    fn test_dense_sparse_conversions() {
137        let A = to_F_dense_matrix::<Fr>(vec![
138            vec![0, 1, 0, 0, 0, 0],
139            vec![0, 0, 0, 1, 0, 0],
140            vec![0, 1, 0, 0, 1, 0],
141            vec![5, 0, 0, 0, 0, 1],
142        ]);
143        let A_sparse = dense_matrix_to_sparse(A.clone());
144        assert_eq!(A_sparse.to_dense(), A);
145    }
146
147    // test mat_vec_mul & mat_vec_mul_sparse
148    #[test]
149    fn test_mat_vec_mul() {
150        let A = to_F_matrix::<Fr>(vec![
151            vec![0, 1, 0, 0, 0, 0],
152            vec![0, 0, 0, 1, 0, 0],
153            vec![0, 1, 0, 0, 1, 0],
154            vec![5, 0, 0, 0, 0, 1],
155        ])
156        .to_dense();
157        let z = to_F_vec(vec![1, 3, 35, 9, 27, 30]);
158        assert_eq!(mat_vec_mul(&A, &z), to_F_vec(vec![3, 9, 30, 35]));
159        assert_eq!(
160            mat_vec_mul_sparse(&dense_matrix_to_sparse(A), &z),
161            to_F_vec(vec![3, 9, 30, 35])
162        );
163
164        let A = to_F_matrix::<Fr>(vec![vec![2, 3, 4, 5], vec![4, 8, 12, 14], vec![9, 8, 7, 6]]);
165        let v = to_F_vec(vec![19, 55, 50, 3]);
166
167        assert_eq!(
168            mat_vec_mul(&A.to_dense(), &v),
169            to_F_vec(vec![418, 1158, 979])
170        );
171        assert_eq!(mat_vec_mul_sparse(&A, &v), to_F_vec(vec![418, 1158, 979]));
172    }
173
174    #[test]
175    fn test_hadamard_product() {
176        let a = to_F_vec::<Fr>(vec![1, 2, 3, 4, 5, 6]);
177        let b = to_F_vec(vec![7, 8, 9, 10, 11, 12]);
178        assert_eq!(hadamard(&a, &b), to_F_vec(vec![7, 16, 27, 40, 55, 72]));
179    }
180    #[test]
181    fn test_vec_add() {
182        let a: Vec<Fr> = to_F_vec::<Fr>(vec![1, 2, 3, 4, 5, 6]);
183        let b: Vec<Fr> = to_F_vec(vec![7, 8, 9, 10, 11, 12]);
184        assert_eq!(vec_add(&a, &b), to_F_vec(vec![8, 10, 12, 14, 16, 18]));
185    }
186}