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