1//! This file implements the struct for an Tuple of Ring Rq elements and its
 2//! operations, which are performed element-wise.
 3
 4use anyhow::Result;
 5use itertools::zip_eq;
 6use rand::{distributions::Distribution, Rng};
 7use rand_distr::{Normal, Uniform};
 8use std::iter::Sum;
 9use std::{
10    array,
11    ops::{Add, Mul, Sub},
12};
13
14use crate::Ring;
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>(pub Vec<R>);
20
21impl<R: Ring, const K: usize> TR<R, K> {
22    pub fn zero() -> Self {
23        Self((0..K).into_iter().map(|_| R::zero()).collect())
24    }
25    pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
26        Self(
27            (0..K)
28                .into_iter()
29                .map(|_| R::rand(&mut rng, &dist))
30                .collect(),
31        )
32    }
33    // returns the decomposition of each polynomial element
34    pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
35        unimplemented!()
36        // self.0.iter().map(|r| r.decompose(beta, l)).collect() // this is Vec<Vec<Vec<R::C>>>
37    }
38}
39
40impl<R: Ring, const K: usize> Add<TR<R, K>> for TR<R, K> {
41    type Output = Self;
42    fn add(self, other: Self) -> Self {
43        Self(
44            zip_eq(self.0, other.0)
45                .map(|(s, o)| s + o)
46                .collect::<Vec<_>>(),
47        )
48    }
49}
50
51impl<R: Ring, const K: usize> Sub<TR<R, K>> for TR<R, K> {
52    type Output = Self;
53    fn sub(self, other: Self) -> Self {
54        Self(zip_eq(self.0, other.0).map(|(s, o)| s - o).collect())
55    }
56}
57
58/// for (TR,TR), the Mul operation is defined as the dot product:
59/// for A, B \in R^k, result = Σ A_i * B_i \in R
60impl<R: Ring, const K: usize> Mul<TR<R, K>> for TR<R, K> {
61    type Output = R;
62    fn mul(self, other: Self) -> R {
63        zip_eq(self.0, other.0).map(|(s, o)| s * o).sum()
64    }
65}
66impl<R: Ring, const K: usize> Mul<&TR<R, K>> for &TR<R, K> {
67    type Output = R;
68    fn mul(self, other: &TR<R, K>) -> R {
69        zip_eq(self.0.clone(), other.0.clone())
70            .map(|(s, o)| s * o)
71            .sum()
72    }
73}
74
75/// for (TR, R), the Mul operation is defined as each element of TR is
76/// multiplied by R
77impl<R: Ring, const K: usize> Mul<R> for TR<R, K> {
78    type Output = TR<R, K>;
79    fn mul(self, other: R) -> TR<R, K> {
80        Self(self.0.iter().map(|s| s.clone() * other.clone()).collect())
81    }
82}
83impl<R: Ring, const K: usize> Mul<&R> for &TR<R, K> {
84    type Output = TR<R, K>;
85    fn mul(self, other: &R) -> TR<R, K> {
86        TR::<R, K>(self.0.iter().map(|s| s.clone() * other.clone()).collect())
87    }
88}