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>(pub Vec<R>);
17
18impl<R: Ring, const K: usize> TR<R, K> {
19    pub fn zero() -> Self {
20        Self((0..K).into_iter().map(|_| R::zero()).collect())
21    }
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    // returns the decomposition of each polynomial element
31    pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
32        unimplemented!()
33        // self.0.iter().map(|r| r.decompose(beta, l)).collect() // this is Vec<Vec<Vec<R::C>>>
34    }
35}
36
37impl<R: Ring, const K: usize> ops::Add<TR<R, K>> for TR<R, K> {
38    type Output = Self;
39    fn add(self, other: Self) -> Self {
40        Self(
41            zip_eq(self.0, other.0)
42                .map(|(s, o)| s + o)
43                .collect::<Vec<_>>(),
44        )
45    }
46}
47
48impl<R: Ring, const K: usize> ops::Sub<TR<R, K>> for TR<R, K> {
49    type Output = Self;
50    fn sub(self, other: Self) -> Self {
51        Self(zip_eq(self.0, other.0).map(|(s, o)| s - o).collect())
52    }
53}
54
55/// for (TR,TR), the Mul operation is defined as:
56/// for A, B \in R^k, result = Σ A_i * B_i \in R
57impl<R: Ring, const K: usize> ops::Mul<TR<R, K>> for TR<R, K> {
58    type Output = R;
59    fn mul(self, other: Self) -> R {
60        zip_eq(self.0, other.0).map(|(s, o)| s * o).sum()
61    }
62}
63impl<R: Ring, const K: usize> ops::Mul<&TR<R, K>> for &TR<R, K> {
64    type Output = R;
65    fn mul(self, other: &TR<R, K>) -> R {
66        zip_eq(self.0.clone(), other.0.clone())
67            .map(|(s, o)| s * o)
68            .sum()
69    }
70}
71
72/// for (TR, R), the Mul operation is defined as each element of TR is
73/// multiplied by R
74impl<R: Ring, const K: usize> ops::Mul<R> for TR<R, K> {
75    type Output = TR<R, K>;
76    fn mul(self, other: R) -> TR<R, K> {
77        Self(self.0.iter().map(|s| s.clone() * other.clone()).collect())
78    }
79}
80impl<R: Ring, const K: usize> ops::Mul<&R> for &TR<R, K> {
81    type Output = TR<R, K>;
82    fn mul(self, other: &R) -> TR<R, K> {
83        TR::<R, K>(self.0.iter().map(|s| s.clone() * other.clone()).collect())
84    }
85}