sonobe/src/utils/vec.rs
src/utils/vec.rsRaw
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 }
51 let mut r: Vec<F> = vec![F::zero(); a.len()];
52 for i in 0..a.len() {
53 r[i] = a[i] + b[i];
54 }
55 Ok(r)
56}
57
58pub fn vec_sub<F: PrimeField>(a: &[F], b: &[F]) -> Result<Vec<F>, Error> {
59 if a.len() != b.len() {
60 return Err(Error::NotSameLength);
61 }
62 let mut r: Vec<F> = vec![F::zero(); a.len()];
63 for i in 0..a.len() {
64 r[i] = a[i] - b[i];
65 }
66 Ok(r)
67}
68
69pub fn vec_scalar_mul<F: PrimeField>(vec: &[F], c: &F) -> Vec<F> {
70 let mut result = vec![F::zero(); vec.len()];
71 for (i, a) in vec.iter().enumerate() {
72 result[i] = *a * c;
73 }
74 result
75}
76
77pub fn is_zero_vec<F: PrimeField>(vec: &[F]) -> bool {
78 for e in vec {
79 if !e.is_zero() {
80 return false;
81 }
82 }
83 true
84}
85
86pub fn mat_vec_mul<F: PrimeField>(M: &Vec<Vec<F>>, z: &[F]) -> Result<Vec<F>, Error> {
87 if M.is_empty() {
88 return Err(Error::Empty);
89 }
90 if M[0].len() != z.len() {
91 return Err(Error::NotSameLength);
92 }
93
94 let mut r: Vec<F> = vec![F::zero(); M.len()];
95 for (i, M_i) in M.iter().enumerate() {
96 for (j, M_ij) in M_i.iter().enumerate() {
97 r[i] += *M_ij * z[j];
98 }
99 }
100 Ok(r)
101}
102
103pub fn mat_vec_mul_sparse<F: PrimeField>(matrix: &SparseMatrix<F>, vector: &[F]) -> Vec<F> {
104 let mut res = vec![F::zero(); matrix.n_rows];
105 for (row_i, row) in matrix.coeffs.iter().enumerate() {
106 for &(value, col_i) in row.iter() {
107 res[row_i] += value * vector[col_i];
108 }
109 }
110
111 res
112}
113
114pub fn hadamard<F: PrimeField>(a: &[F], b: &[F]) -> Result<Vec<F>, Error> {
115 if a.len() != b.len() {
116 return Err(Error::NotSameLength);
117 }
118 Ok(cfg_iter!(a).zip(b).map(|(a, b)| *a * b).collect())
119}
120
121#[cfg(test)]
122pub mod tests {
123 use super::*;
124 use ark_pallas::Fr;
125
126 pub fn to_F_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> SparseMatrix<F> {
127 dense_matrix_to_sparse(to_F_dense_matrix(M))
128 }
129 pub fn to_F_dense_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> Vec<Vec<F>> {
130 let mut R: Vec<Vec<F>> = vec![Vec::new(); M.len()];
131 for i in 0..M.len() {
132 R[i] = vec![F::zero(); M[i].len()];
133 for j in 0..M[i].len() {
134 R[i][j] = F::from(M[i][j] as u64);
135 }
136 }
137 R
138 }
139 pub fn to_F_vec<F: PrimeField>(z: Vec<usize>) -> Vec<F> {
140 let mut r: Vec<F> = vec![F::zero(); z.len()];
141 for i in 0..z.len() {
142 r[i] = F::from(z[i] as u64);
143 }
144 r
145 }
146
147 #[test]
148 fn test_dense_sparse_conversions() {
149 let A = to_F_dense_matrix::<Fr>(vec![
150 vec![0, 1, 0, 0, 0, 0],
151 vec![0, 0, 0, 1, 0, 0],
152 vec![0, 1, 0, 0, 1, 0],
153 vec![5, 0, 0, 0, 0, 1],
154 ]);
155 let A_sparse = dense_matrix_to_sparse(A.clone());
156 assert_eq!(A_sparse.to_dense(), A);
157 }
158
159 // test mat_vec_mul & mat_vec_mul_sparse
160 #[test]
161 fn test_mat_vec_mul() {
162 let A = to_F_matrix::<Fr>(vec![
163 vec![0, 1, 0, 0, 0, 0],
164 vec![0, 0, 0, 1, 0, 0],
165 vec![0, 1, 0, 0, 1, 0],
166 vec![5, 0, 0, 0, 0, 1],
167 ])
168 .to_dense();
169 let z = to_F_vec(vec![1, 3, 35, 9, 27, 30]);
170 assert_eq!(mat_vec_mul(&A, &z).unwrap(), to_F_vec(vec![3, 9, 30, 35]));
171 assert_eq!(
172 mat_vec_mul_sparse(&dense_matrix_to_sparse(A), &z),
173 to_F_vec(vec![3, 9, 30, 35])
174 );
175
176 let A = to_F_matrix::<Fr>(vec![vec![2, 3, 4, 5], vec![4, 8, 12, 14], vec![9, 8, 7, 6]]);
177 let v = to_F_vec(vec![19, 55, 50, 3]);
178
179 assert_eq!(
180 mat_vec_mul(&A.to_dense(), &v).unwrap(),
181 to_F_vec(vec![418, 1158, 979])
182 );
183 assert_eq!(mat_vec_mul_sparse(&A, &v), to_F_vec(vec![418, 1158, 979]));
184 }
185
186 #[test]
187 fn test_hadamard_product() {
188 let a = to_F_vec::<Fr>(vec![1, 2, 3, 4, 5, 6]);
189 let b = to_F_vec(vec![7, 8, 9, 10, 11, 12]);
190 assert_eq!(
191 hadamard(&a, &b).unwrap(),
192 to_F_vec(vec![7, 16, 27, 40, 55, 72])
193 );
194 }
195 #[test]
196 fn test_vec_add() {
197 let a: Vec<Fr> = to_F_vec::<Fr>(vec![1, 2, 3, 4, 5, 6]);
198 let b: Vec<Fr> = to_F_vec(vec![7, 8, 9, 10, 11, 12]);
199 assert_eq!(
200 vec_add(&a, &b).unwrap(),
201 to_F_vec(vec![8, 10, 12, 14, 16, 18])
202 );
203 }
204}