1use ark_ff::PrimeField;
  2use ark_poly::{
  3    univariate::DensePolynomial, EvaluationDomain, Evaluations, GeneralEvaluationDomain,
  4};
  5pub use ark_relations::r1cs::Matrix as R1CSMatrix;
  6use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
  7use ark_std::cfg_iter;
  8use ark_std::rand::Rng;
  9use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
 10
 11use crate::Error;
 12
 13#[derive(Clone, Debug, Eq, PartialEq, CanonicalSerialize, CanonicalDeserialize)]
 14pub struct SparseMatrix<F: PrimeField> {
 15    pub n_rows: usize,
 16    pub n_cols: usize,
 17    /// coeffs = R1CSMatrix = Vec<Vec<(F, usize)>>, which contains each row and the F is the value
 18    /// of the coefficient and the usize indicates the column position
 19    pub coeffs: R1CSMatrix<F>,
 20}
 21
 22impl<F: PrimeField> SparseMatrix<F> {
 23    pub fn rand<R: Rng>(rng: &mut R, n_rows: usize, n_cols: usize) -> Self {
 24        const ZERO_VAL_PROBABILITY: f64 = 0.8f64;
 25
 26        let dense = (0..n_rows)
 27            .map(|_| {
 28                (0..n_cols)
 29                    .map(|_| {
 30                        if !rng.gen_bool(ZERO_VAL_PROBABILITY) {
 31                            return F::rand(rng);
 32                        }
 33                        F::zero()
 34                    })
 35                    .collect::<Vec<F>>()
 36            })
 37            .collect::<Vec<Vec<F>>>();
 38        dense_matrix_to_sparse(dense)
 39    }
 40    pub fn to_dense(&self) -> Vec<Vec<F>> {
 41        let mut r: Vec<Vec<F>> = vec![vec![F::zero(); self.n_cols]; self.n_rows];
 42        for (row_i, row) in self.coeffs.iter().enumerate() {
 43            for &(value, col_i) in row.iter() {
 44                r[row_i][col_i] = value;
 45            }
 46        }
 47        r
 48    }
 49}
 50
 51pub fn dense_matrix_to_sparse<F: PrimeField>(m: Vec<Vec<F>>) -> SparseMatrix<F> {
 52    let mut r = SparseMatrix::<F> {
 53        n_rows: m.len(),
 54        n_cols: m[0].len(),
 55        coeffs: Vec::new(),
 56    };
 57    for m_row in m.iter() {
 58        let mut row: Vec<(F, usize)> = Vec::new();
 59        for (col_i, value) in m_row.iter().enumerate() {
 60            if !value.is_zero() {
 61                row.push((*value, col_i));
 62            }
 63        }
 64        r.coeffs.push(row);
 65    }
 66    r
 67}
 68
 69pub fn vec_add<F: PrimeField>(a: &[F], b: &[F]) -> Result<Vec<F>, Error> {
 70    if a.len() != b.len() {
 71        return Err(Error::NotSameLength(
 72            "a.len()".to_string(),
 73            a.len(),
 74            "b.len()".to_string(),
 75            b.len(),
 76        ));
 77    }
 78    Ok(a.iter().zip(b.iter()).map(|(x, y)| *x + y).collect())
 79}
 80
 81pub fn vec_sub<F: PrimeField>(a: &[F], b: &[F]) -> Result<Vec<F>, Error> {
 82    if a.len() != b.len() {
 83        return Err(Error::NotSameLength(
 84            "a.len()".to_string(),
 85            a.len(),
 86            "b.len()".to_string(),
 87            b.len(),
 88        ));
 89    }
 90    Ok(a.iter().zip(b.iter()).map(|(x, y)| *x - y).collect())
 91}
 92
 93pub fn vec_scalar_mul<F: PrimeField>(vec: &[F], c: &F) -> Vec<F> {
 94    vec.iter().map(|a| *a * c).collect()
 95}
 96
 97pub fn is_zero_vec<F: PrimeField>(vec: &[F]) -> bool {
 98    vec.iter().all(|a| a.is_zero())
 99}
100
101pub fn mat_vec_mul_dense<F: PrimeField>(M: &[Vec<F>], z: &[F]) -> Result<Vec<F>, Error> {
102    if M.is_empty() {
103        return Err(Error::Empty);
104    }
105    if M[0].len() != z.len() {
106        return Err(Error::NotSameLength(
107            "M[0].len()".to_string(),
108            M[0].len(),
109            "z.len()".to_string(),
110            z.len(),
111        ));
112    }
113
114    let mut r: Vec<F> = vec![F::zero(); M.len()];
115    for (i, M_i) in M.iter().enumerate() {
116        for (j, M_ij) in M_i.iter().enumerate() {
117            r[i] += *M_ij * z[j];
118        }
119    }
120    Ok(r)
121}
122
123pub fn mat_vec_mul<F: PrimeField>(M: &SparseMatrix<F>, z: &[F]) -> Result<Vec<F>, Error> {
124    if M.n_cols != z.len() {
125        return Err(Error::NotSameLength(
126            "M.n_cols".to_string(),
127            M.n_cols,
128            "z.len()".to_string(),
129            z.len(),
130        ));
131    }
132    let mut res = vec![F::zero(); M.n_rows];
133    for (row_i, row) in M.coeffs.iter().enumerate() {
134        for &(value, col_i) in row.iter() {
135            res[row_i] += value * z[col_i];
136        }
137    }
138    Ok(res)
139}
140
141pub fn mat_from_str_mat<F: PrimeField>(str_mat: Vec<Vec<&str>>) -> Result<Vec<Vec<F>>, Error> {
142    str_mat
143        .into_iter()
144        .map(|row| {
145            row.into_iter()
146                .map(|s| {
147                    F::from_str(s).map_err(|_| Error::Other("Invalid decimal string".to_string()))
148                })
149                .collect()
150        })
151        .collect()
152}
153
154pub fn hadamard<F: PrimeField>(a: &[F], b: &[F]) -> Result<Vec<F>, Error> {
155    if a.len() != b.len() {
156        return Err(Error::NotSameLength(
157            "a.len()".to_string(),
158            a.len(),
159            "b.len()".to_string(),
160            b.len(),
161        ));
162    }
163    Ok(cfg_iter!(a).zip(b).map(|(a, b)| *a * b).collect())
164}
165
166/// returns the interpolated polynomial of degree=v.len().next_power_of_two(), which passes through all
167/// the given elements of v.
168pub fn poly_from_vec<F: PrimeField>(v: Vec<F>) -> Result<DensePolynomial<F>, Error> {
169    let D = GeneralEvaluationDomain::<F>::new(v.len()).ok_or(Error::NewDomainFail)?;
170    Ok(Evaluations::from_vec_and_domain(v, D).interpolate())
171}
172
173#[cfg(test)]
174pub mod tests {
175    use super::*;
176    use ark_pallas::Fr;
177
178    pub fn to_F_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> SparseMatrix<F> {
179        dense_matrix_to_sparse(to_F_dense_matrix(M))
180    }
181    pub fn to_F_dense_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> Vec<Vec<F>> {
182        M.iter()
183            .map(|m| m.iter().map(|r| F::from(*r as u64)).collect())
184            .collect()
185    }
186    pub fn to_F_vec<F: PrimeField>(z: Vec<usize>) -> Vec<F> {
187        z.iter().map(|c| F::from(*c as u64)).collect()
188    }
189
190    #[test]
191    fn test_dense_sparse_conversions() {
192        let A = to_F_dense_matrix::<Fr>(vec![
193            vec![0, 1, 0, 0, 0, 0],
194            vec![0, 0, 0, 1, 0, 0],
195            vec![0, 1, 0, 0, 1, 0],
196            vec![5, 0, 0, 0, 0, 1],
197        ]);
198        let A_sparse = dense_matrix_to_sparse(A.clone());
199        assert_eq!(A_sparse.to_dense(), A);
200    }
201
202    // test mat_vec_mul & mat_vec_mul_sparse
203    #[test]
204    fn test_mat_vec_mul() {
205        let A = to_F_matrix::<Fr>(vec![
206            vec![0, 1, 0, 0, 0, 0],
207            vec![0, 0, 0, 1, 0, 0],
208            vec![0, 1, 0, 0, 1, 0],
209            vec![5, 0, 0, 0, 0, 1],
210        ])
211        .to_dense();
212        let z = to_F_vec(vec![1, 3, 35, 9, 27, 30]);
213        assert_eq!(
214            mat_vec_mul_dense(&A, &z).unwrap(),
215            to_F_vec(vec![3, 9, 30, 35])
216        );
217        assert_eq!(
218            mat_vec_mul(&dense_matrix_to_sparse(A), &z).unwrap(),
219            to_F_vec(vec![3, 9, 30, 35])
220        );
221
222        let A = to_F_matrix::<Fr>(vec![vec![2, 3, 4, 5], vec![4, 8, 12, 14], vec![9, 8, 7, 6]]);
223        let v = to_F_vec(vec![19, 55, 50, 3]);
224
225        assert_eq!(
226            mat_vec_mul_dense(&A.to_dense(), &v).unwrap(),
227            to_F_vec(vec![418, 1158, 979])
228        );
229        assert_eq!(mat_vec_mul(&A, &v).unwrap(), to_F_vec(vec![418, 1158, 979]));
230    }
231
232    #[test]
233    fn test_hadamard_product() {
234        let a = to_F_vec::<Fr>(vec![1, 2, 3, 4, 5, 6]);
235        let b = to_F_vec(vec![7, 8, 9, 10, 11, 12]);
236        assert_eq!(
237            hadamard(&a, &b).unwrap(),
238            to_F_vec(vec![7, 16, 27, 40, 55, 72])
239        );
240    }
241
242    #[test]
243    fn test_vec_add() {
244        let a: Vec<Fr> = to_F_vec::<Fr>(vec![1, 2, 3, 4, 5, 6]);
245        let b: Vec<Fr> = to_F_vec(vec![7, 8, 9, 10, 11, 12]);
246        assert_eq!(
247            vec_add(&a, &b).unwrap(),
248            to_F_vec(vec![8, 10, 12, 14, 16, 18])
249        );
250    }
251}