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