1//! This file implements the struct for an Tuple of Ring Rq elements and its
 2//! operations.
 3
 4use anyhow::Result;
 5use itertools::zip_eq;
 6use rand::{distributions::Distribution, Rng};
 7use rand_distr::{Normal, Uniform};
 8use std::iter::Sum;
 9use std::{array, ops};
10
11use crate::Ring;
12
13/// Tuple of K Ring (Rq) elements. We use Vec<R> to allocate it in the heap,
14/// since if using a fixed-size array it would overflow the stack.
15#[derive(Clone, Debug)]
16pub struct TR<R: Ring, const K: usize>(Vec<R>);
17
18impl<R: Ring, const K: usize> TR<R, K> {
19    pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
20        Self(
21            (0..K)
22                .into_iter()
23                .map(|_| R::rand(&mut rng, &dist))
24                .collect(),
25        )
26    }
27}
28
29impl<R: Ring, const K: usize> ops::Add<TR<R, K>> for TR<R, K> {
30    type Output = Self;
31    fn add(self, other: Self) -> Self {
32        Self(
33            zip_eq(self.0, other.0)
34                .map(|(s, o)| s + o)
35                .collect::<Vec<_>>(),
36        )
37    }
38}
39
40impl<R: Ring, const K: usize> ops::Sub<TR<R, K>> for TR<R, K> {
41    type Output = Self;
42    fn sub(self, other: Self) -> Self {
43        Self(zip_eq(self.0, other.0).map(|(s, o)| s - o).collect())
44    }
45}
46
47/// for (TR,TR), the Mul operation is defined as:
48/// for A, B \in R^k, result = Σ A_i * B_i \in R
49impl<R: Ring, const K: usize> ops::Mul<TR<R, K>> for TR<R, K> {
50    type Output = R;
51    fn mul(self, other: Self) -> R {
52        zip_eq(self.0, other.0).map(|(s, o)| s * o).sum()
53    }
54}
55impl<R: Ring, const K: usize> ops::Mul<&TR<R, K>> for &TR<R, K> {
56    type Output = R;
57    fn mul(self, other: &TR<R, K>) -> R {
58        zip_eq(self.0.clone(), other.0.clone())
59            .map(|(s, o)| s * o)
60            .sum()
61    }
62}
63
64/// for (TR, R), the Mul operation is defined as each element of TR is
65/// multiplied by R
66impl<R: Ring, const K: usize> ops::Mul<R> for TR<R, K> {
67    type Output = TR<R, K>;
68    fn mul(self, other: R) -> TR<R, K> {
69        Self(self.0.iter().map(|s| s.clone() * other.clone()).collect())
70    }
71}
72impl<R: Ring, const K: usize> ops::Mul<&R> for &TR<R, K> {
73    type Output = TR<R, K>;
74    fn mul(self, other: &R) -> TR<R, K> {
75        TR::<R, K>(self.0.iter().map(|s| s.clone() * other.clone()).collect())
76    }
77}