1use ark_ff::fields::PrimeField;
  2use core::ops::Add;
  3
  4pub fn vector_elem_product<F: PrimeField>(a: &Vec<F>, e: &F) -> Vec<F> {
  5    // maybe move this method to operator a * e
  6    let mut r: Vec<F> = vec![F::zero(); a.len()];
  7    for i in 0..a.len() {
  8        r[i] = a[i] * e;
  9    }
 10    r
 11}
 12pub fn matrix_vector_product<F: PrimeField>(M: &Vec<Vec<F>>, z: &Vec<F>) -> Vec<F> {
 13    // TODO assert len
 14    let mut r: Vec<F> = vec![F::zero(); M.len()];
 15    for i in 0..M.len() {
 16        for j in 0..M[i].len() {
 17            r[i] += M[i][j] * z[j];
 18        }
 19    }
 20    r
 21}
 22pub fn hadamard_product<F: PrimeField>(a: Vec<F>, b: Vec<F>) -> Vec<F> {
 23    // maybe move this method to operator a * b
 24    // TODO assert equals len
 25    let mut r: Vec<F> = vec![F::zero(); a.len()];
 26    for i in 0..a.len() {
 27        r[i] = a[i] * b[i];
 28    }
 29    r
 30}
 31
 32pub fn vec_add<F: PrimeField>(a: Vec<F>, b: Vec<F>) -> Vec<F> {
 33    let mut r: Vec<F> = vec![F::zero(); a.len()];
 34    for i in 0..a.len() {
 35        r[i] = a[i] + b[i];
 36    }
 37    r
 38}
 39// TODO instead of vec_add do:
 40// impl<'a, 'b, F> Add<&'b Vec<F>> for &'a Vec<F> {
 41//     type Output = Vec<F>;
 42//     fn add(self, rhs: &'b Vec<F>) -> Vec<F> {
 43//         let mut r: Vec<F> = vec![F::zero(); self.len()];
 44//         for i in 0..self.len() {
 45//             r[i] = self[i] + rhs[i];
 46//         }
 47//         r
 48//     }
 49// }
 50pub fn vec_sub<F: PrimeField>(a: Vec<F>, b: Vec<F>) -> Vec<F> {
 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    r
 56}
 57
 58#[cfg(test)]
 59mod tests {
 60    use super::*;
 61    use ark_bn254::{g1::G1Affine, Fr};
 62    use ark_ec::CurveGroup;
 63    use ark_std::{One, Zero};
 64    use std::ops::Mul;
 65
 66    fn to_F_matrix<F: PrimeField>(M: Vec<Vec<usize>>) -> Vec<Vec<F>> {
 67        let mut R: Vec<Vec<F>> = vec![Vec::new(); M.len()];
 68        for i in 0..M.len() {
 69            R[i] = vec![F::zero(); M[i].len()];
 70            for j in 0..M[i].len() {
 71                R[i][j] = F::from(M[i][j] as u64);
 72            }
 73        }
 74        R
 75    }
 76    fn to_F_vec<F: PrimeField>(z: Vec<usize>) -> Vec<F> {
 77        let mut r: Vec<F> = vec![F::zero(); z.len()];
 78        for i in 0..z.len() {
 79            r[i] = F::from(z[i] as u64);
 80        }
 81        r
 82    }
 83
 84    #[test]
 85    fn test_matrix_vector_product() {
 86        let A = to_F_matrix::<Fr>(vec![
 87            vec![0, 1, 0, 0, 0, 0],
 88            vec![0, 0, 0, 1, 0, 0],
 89            vec![0, 1, 0, 0, 1, 0],
 90            vec![5, 0, 0, 0, 0, 1],
 91        ]);
 92        let z = to_F_vec(vec![1, 3, 35, 9, 27, 30]);
 93        let Az = matrix_vector_product(&A, &z);
 94        assert_eq!(Az, to_F_vec(vec![3, 9, 30, 35]));
 95    }
 96
 97    #[test]
 98    fn test_hadamard_product() {
 99        let a = to_F_vec::<Fr>(vec![1, 2, 3, 4, 5, 6]);
100        let b = to_F_vec(vec![7, 8, 9, 10, 11, 12]);
101        assert_eq!(
102            hadamard_product(a, b),
103            to_F_vec(vec![7, 16, 27, 40, 55, 72])
104        );
105    }
106
107    #[test]
108    fn test_ABC_hadamard() {
109        let A = to_F_matrix::<Fr>(vec![
110            vec![0, 1, 0, 0, 0, 0],
111            vec![0, 0, 0, 1, 0, 0],
112            vec![0, 1, 0, 0, 1, 0],
113            vec![5, 0, 0, 0, 0, 1],
114        ]);
115        let B = to_F_matrix(vec![
116            vec![0, 1, 0, 0, 0, 0],
117            vec![0, 1, 0, 0, 0, 0],
118            vec![1, 0, 0, 0, 0, 0],
119            vec![1, 0, 0, 0, 0, 0],
120        ]);
121        let C = to_F_matrix(vec![
122            vec![0, 0, 0, 1, 0, 0],
123            vec![0, 0, 0, 0, 1, 0],
124            vec![0, 0, 0, 0, 0, 1],
125            vec![0, 0, 1, 0, 0, 0],
126        ]);
127        let z = to_F_vec(vec![1, 3, 35, 9, 27, 30]);
128
129        let Az = matrix_vector_product(&A, &z);
130        let Bz = matrix_vector_product(&B, &z);
131        let Cz = matrix_vector_product(&C, &z);
132        assert_eq!(hadamard_product(Az, Bz), Cz);
133    }
134}