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// TODO rm pub from Vec<R>, so that TR can not be created from a Vec with
 21// invalid length, since it has to be created using the `new` method.
 22
 23impl<R: Ring, const K: usize> TR<R, K> {
 24    pub fn new(v: Vec<R>) -> Self {
 25        assert_eq!(v.len(), K);
 26        Self(v)
 27    }
 28    pub fn zero() -> Self {
 29        Self((0..K).into_iter().map(|_| R::zero()).collect())
 30    }
 31    pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
 32        Self(
 33            (0..K)
 34                .into_iter()
 35                .map(|_| R::rand(&mut rng, &dist))
 36                .collect(),
 37        )
 38    }
 39    // returns the decomposition of each polynomial element
 40    pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
 41        unimplemented!()
 42        // self.0.iter().map(|r| r.decompose(beta, l)).collect() // this is Vec<Vec<Vec<R::C>>>
 43    }
 44}
 45
 46impl<const K: usize> TR<crate::torus::T64, K> {
 47    pub fn mod_switch<const Q2: u64>(&self) -> TR<crate::torus::T64, K> {
 48        TR(self.0.iter().map(|c_i| c_i.mod_switch::<Q2>()).collect())
 49    }
 50    // pub fn mod_switch(&self, Q2: u64) -> TR<crate::torus::T64, K> {
 51    //     TR(self.0.iter().map(|c_i| c_i.mod_switch(Q2)).collect())
 52    // }
 53}
 54impl<const N: usize, const K: usize> TR<crate::ring_torus::Tn<N>, K> {
 55    pub fn left_rotate(&self, h: usize) -> Self {
 56        TR(self.0.iter().map(|c_i| c_i.left_rotate(h)).collect())
 57    }
 58}
 59
 60impl<R: Ring, const K: usize> TR<R, K> {
 61    pub fn iter(&self) -> std::slice::Iter<R> {
 62        self.0.iter()
 63    }
 64}
 65
 66impl<R: Ring, const K: usize> Add<TR<R, K>> for TR<R, K> {
 67    type Output = Self;
 68    fn add(self, other: Self) -> Self {
 69        Self(
 70            zip_eq(self.0, other.0)
 71                .map(|(s, o)| s + o)
 72                .collect::<Vec<_>>(),
 73        )
 74    }
 75}
 76
 77impl<R: Ring, const K: usize> Sub<TR<R, K>> for TR<R, K> {
 78    type Output = Self;
 79    fn sub(self, other: Self) -> Self {
 80        Self(zip_eq(self.0, other.0).map(|(s, o)| s - o).collect())
 81    }
 82}
 83
 84impl<R: Ring, const K: usize> Neg for TR<R, K> {
 85    type Output = Self;
 86
 87    fn neg(self) -> Self::Output {
 88        Self(self.0.iter().map(|&e_i| -e_i).collect())
 89    }
 90}
 91
 92/// for (TR,TR), the Mul operation is defined as the dot product:
 93/// for A, B \in R^k, result = Σ A_i * B_i \in R
 94impl<R: Ring, const K: usize> Mul<TR<R, K>> for TR<R, K> {
 95    type Output = R;
 96    fn mul(self, other: Self) -> R {
 97        zip_eq(self.0, other.0).map(|(s, o)| s * o).sum()
 98    }
 99}
100impl<R: Ring, const K: usize> Mul<&TR<R, K>> for &TR<R, K> {
101    type Output = R;
102    fn mul(self, other: &TR<R, K>) -> R {
103        zip_eq(self.0.clone(), other.0.clone())
104            .map(|(s, o)| s * o)
105            .sum()
106    }
107}
108
109/// for (TR, R), the Mul operation is defined as each element of TR is
110/// multiplied by R
111impl<R: Ring, const K: usize> Mul<R> for TR<R, K> {
112    type Output = TR<R, K>;
113    fn mul(self, other: R) -> TR<R, K> {
114        Self(self.0.iter().map(|s| s.clone() * other.clone()).collect())
115    }
116}
117impl<R: Ring, const K: usize> Mul<&R> for &TR<R, K> {
118    type Output = TR<R, K>;
119    fn mul(self, other: &R) -> TR<R, K> {
120        TR::<R, K>(self.0.iter().map(|s| s.clone() * other.clone()).collect())
121    }
122}