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