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