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