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, Neg, 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> TR<R, K> {
 41    pub fn iter(&self) -> std::slice::Iter<R> {
 42        self.0.iter()
 43    }
 44}
 45
 46impl<R: Ring, const K: usize> Add<TR<R, K>> for TR<R, K> {
 47    type Output = Self;
 48    fn add(self, other: Self) -> Self {
 49        Self(
 50            zip_eq(self.0, other.0)
 51                .map(|(s, o)| s + o)
 52                .collect::<Vec<_>>(),
 53        )
 54    }
 55}
 56
 57impl<R: Ring, const K: usize> Sub<TR<R, K>> for TR<R, K> {
 58    type Output = Self;
 59    fn sub(self, other: Self) -> Self {
 60        Self(zip_eq(self.0, other.0).map(|(s, o)| s - o).collect())
 61    }
 62}
 63
 64impl<R: Ring, const K: usize> Neg for TR<R, K> {
 65    type Output = Self;
 66
 67    fn neg(self) -> Self::Output {
 68        Self(self.0.iter().map(|&e_i| -e_i).collect())
 69    }
 70}
 71
 72/// for (TR,TR), the Mul operation is defined as the dot product:
 73/// for A, B \in R^k, result = Σ A_i * B_i \in R
 74impl<R: Ring, const K: usize> Mul<TR<R, K>> for TR<R, K> {
 75    type Output = R;
 76    fn mul(self, other: Self) -> R {
 77        zip_eq(self.0, other.0).map(|(s, o)| s * o).sum()
 78    }
 79}
 80impl<R: Ring, const K: usize> Mul<&TR<R, K>> for &TR<R, K> {
 81    type Output = R;
 82    fn mul(self, other: &TR<R, K>) -> R {
 83        zip_eq(self.0.clone(), other.0.clone())
 84            .map(|(s, o)| s * o)
 85            .sum()
 86    }
 87}
 88
 89/// for (TR, R), the Mul operation is defined as each element of TR is
 90/// multiplied by R
 91impl<R: Ring, const K: usize> Mul<R> for TR<R, K> {
 92    type Output = TR<R, K>;
 93    fn mul(self, other: R) -> TR<R, K> {
 94        Self(self.0.iter().map(|s| s.clone() * other.clone()).collect())
 95    }
 96}
 97impl<R: Ring, const K: usize> Mul<&R> for &TR<R, K> {
 98    type Output = TR<R, K>;
 99    fn mul(self, other: &R) -> TR<R, K> {
100        TR::<R, K>(self.0.iter().map(|s| s.clone() * other.clone()).collect())
101    }
102}