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