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