1use ark_ec::CurveGroup;
  2use ark_ff::fields::PrimeField;
  3use core::ops::{Add, Sub};
  4use std::fmt;
  5
  6pub fn vector_elem_product<F: PrimeField>(a: &Vec<F>, e: &F) -> Vec<F> {
  7    // maybe move this method to operator a * e
  8    let mut r: Vec<F> = vec![F::zero(); a.len()];
  9    for i in 0..a.len() {
 10        r[i] = a[i] * e;
 11    }
 12    r
 13}
 14
 15#[allow(clippy::needless_range_loop)]
 16pub fn matrix_vector_product<F: PrimeField>(M: &Vec<Vec<F>>, z: &[F]) -> Vec<F> {
 17    // TODO assert len
 18    let mut r: Vec<F> = vec![F::zero(); M.len()];
 19    for i in 0..M.len() {
 20        for j in 0..M[i].len() {
 21            r[i] += M[i][j] * z[j];
 22        }
 23    }
 24    r
 25}
 26pub fn hadamard_product<F: PrimeField>(a: Vec<F>, b: Vec<F>) -> Vec<F> {
 27    // maybe move this method to operator a * b
 28    // TODO assert equals len
 29    let mut r: Vec<F> = vec![F::zero(); a.len()];
 30    for i in 0..a.len() {
 31        r[i] = a[i] * b[i];
 32    }
 33    r
 34}
 35
 36// rlin: random linear combination
 37// pub fn rlin<F: PrimeField>(a: Vec<F>, b: Vec<F>, r: F) -> Vec<F> {
 38//     vec_add(a, vector_elem_product(&b, &r)) // WIP probably group loops
 39// }
 40
 41pub fn naive_msm<C: CurveGroup>(s: &Vec<C::ScalarField>, p: &[C]) -> C {
 42    // TODO check lengths, or at least check s.len()>= p.len()
 43
 44    let mut r = p[0].mul(s[0]);
 45    for i in 1..s.len() {
 46        r = p[i].mul(s[i]);
 47    }
 48    r
 49}
 50
 51pub fn vec_add<F: PrimeField>(a: &Vec<F>, b: &[F]) -> Vec<F> {
 52    let mut r: Vec<F> = vec![F::zero(); a.len()];
 53    for i in 0..a.len() {
 54        r[i] = a[i] + b[i];
 55    }
 56    r
 57}
 58pub fn vec_sub<F: PrimeField>(a: Vec<F>, b: Vec<F>) -> Vec<F> {
 59    let mut r: Vec<F> = vec![F::zero(); a.len()];
 60    for i in 0..a.len() {
 61        r[i] = a[i] - b[i];
 62    }
 63    r
 64}
 65
 66// instead of vec_{add/sub} can use Ve wrapper which has '+', '-' operators
 67#[derive(Debug, Clone, PartialEq)]
 68pub struct Ve<F: PrimeField>(pub Vec<F>);
 69
 70impl<F: PrimeField> fmt::Display for Ve<F> {
 71    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 72        for (i, e) in self.0.iter().enumerate() {
 73            if i == self.0.len() - 1 {
 74                write!(f, "{}", e)?;
 75                break;
 76            }
 77            write!(f, "{}, ", e)?;
 78        }
 79        Ok(())
 80    }
 81}
 82
 83impl<F: PrimeField> Add<Ve<F>> for Ve<F> {
 84    type Output = Ve<F>;
 85    fn add(self, rhs_vec: Self) -> Ve<F> {
 86        let lhs = self.0.clone();
 87        let rhs = rhs_vec.0;
 88        let mut r: Vec<F> = vec![F::zero(); lhs.len()];
 89        for i in 0..self.0.len() {
 90            r[i] = lhs[i] + rhs[i];
 91        }
 92        Ve(r)
 93    }
 94}
 95impl<F: PrimeField> Sub<Ve<F>> for Ve<F> {
 96    type Output = Ve<F>;
 97    fn sub(self, rhs_vec: Self) -> Ve<F> {
 98        let lhs = self.0.clone();
 99        let rhs = rhs_vec.0;
100        let mut r: Vec<F> = vec![F::zero(); lhs.len()];
101        for i in 0..self.0.len() {
102            r[i] = lhs[i] - rhs[i];
103        }
104        Ve(r)
105    }
106}
107
108pub fn to_F_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> Vec<Vec<F>> {
109    let mut R: Vec<Vec<F>> = vec![Vec::new(); M.len()];
110    for i in 0..M.len() {
111        R[i] = vec![F::zero(); M[i].len()];
112        for j in 0..M[i].len() {
113            R[i][j] = F::from(M[i][j] as u64);
114        }
115    }
116    R
117}
118pub fn to_F_vec<F: PrimeField>(z: Vec<usize>) -> Vec<F> {
119    let mut r: Vec<F> = vec![F::zero(); z.len()];
120    for i in 0..z.len() {
121        r[i] = F::from(z[i] as u64);
122    }
123    r
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use ark_mnt4_298::Fr;
130
131    #[test]
132    fn test_matrix_vector_product() {
133        let A = to_F_matrix::<Fr>(vec![
134            vec![0, 1, 0, 0, 0, 0],
135            vec![0, 0, 0, 1, 0, 0],
136            vec![0, 1, 0, 0, 1, 0],
137            vec![5, 0, 0, 0, 0, 1],
138        ]);
139        let z = to_F_vec(vec![1, 3, 35, 9, 27, 30]);
140        let Az = matrix_vector_product(&A, &z);
141        assert_eq!(Az, to_F_vec(vec![3, 9, 30, 35]));
142    }
143
144    #[test]
145    fn test_hadamard_product() {
146        let a = to_F_vec::<Fr>(vec![1, 2, 3, 4, 5, 6]);
147        let b = to_F_vec(vec![7, 8, 9, 10, 11, 12]);
148        assert_eq!(
149            hadamard_product(a, b),
150            to_F_vec(vec![7, 16, 27, 40, 55, 72])
151        );
152    }
153    #[test]
154    fn test_vec_add() {
155        let a: Vec<Fr> = to_F_vec::<Fr>(vec![1, 2, 3, 4, 5, 6]);
156        let b: Vec<Fr> = to_F_vec(vec![7, 8, 9, 10, 11, 12]);
157        assert_eq!(vec_add(&a, &b), (Ve(a) + Ve(b)).0);
158    }
159
160    #[test]
161    fn test_ABC_hadamard() {
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        let B = to_F_matrix(vec![
169            vec![0, 1, 0, 0, 0, 0],
170            vec![0, 1, 0, 0, 0, 0],
171            vec![1, 0, 0, 0, 0, 0],
172            vec![1, 0, 0, 0, 0, 0],
173        ]);
174        let C = to_F_matrix(vec![
175            vec![0, 0, 0, 1, 0, 0],
176            vec![0, 0, 0, 0, 1, 0],
177            vec![0, 0, 0, 0, 0, 1],
178            vec![0, 0, 1, 0, 0, 0],
179        ]);
180        let z = to_F_vec(vec![1, 3, 35, 9, 27, 30]);
181
182        let Az = matrix_vector_product(&A, &z);
183        let Bz = matrix_vector_product(&B, &z);
184        let Cz = matrix_vector_product(&C, &z);
185        assert_eq!(hadamard_product(Az, Bz), Cz);
186    }
187}