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