1//! This file implements the struct for an Tuple of Ring Rq elements and its
  2//! operations, which are performed element-wise.
  3
  4use itertools::zip_eq;
  5use rand::{distr::Distribution, Rng};
  6use std::ops::{Add, Mul, Neg, Sub};
  7
  8use crate::{Ring, RingParam};
  9
 10/// Tuple of K Ring (Rq) elements. We use Vec<R> to allocate it in the heap,
 11/// since if using a fixed-size array it would overflow the stack.
 12#[derive(Clone, Debug)]
 13pub struct TR<R: Ring> {
 14    pub k: usize,
 15    pub r: Vec<R>,
 16}
 17// TODO rm pub from Vec<R>, so that TR can not be created from a Vec with
 18// invalid length, since it has to be created using the `new` method.
 19
 20impl<R: Ring> TR<R> {
 21    pub fn new(k: usize, r: Vec<R>) -> Self {
 22        assert_eq!(r.len(), k);
 23        Self { k, r }
 24    }
 25    pub fn zero(k: usize, r_param: &RingParam) -> Self {
 26        Self {
 27            k,
 28            r: (0..k).into_iter().map(|_| R::zero(r_param)).collect(),
 29        }
 30    }
 31    pub fn rand(
 32        mut rng: impl Rng,
 33        dist: impl Distribution<f64>,
 34        k: usize,
 35        r_param: &RingParam,
 36    ) -> Self {
 37        Self {
 38            k,
 39            r: (0..k)
 40                .into_iter()
 41                .map(|_| R::rand(&mut rng, &dist, r_param))
 42                .collect(),
 43        }
 44    }
 45    // returns the decomposition of each polynomial element
 46    pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
 47        unimplemented!()
 48        // self.0.iter().map(|r| r.decompose(beta, l)).collect() // this is Vec<Vec<Vec<R::C>>>
 49    }
 50}
 51
 52impl TR<crate::torus::T64> {
 53    pub fn mod_switch(&self, q2: u64) -> TR<crate::torus::T64> {
 54        TR::<crate::torus::T64> {
 55            k: self.k,
 56            r: self.r.iter().map(|c_i| c_i.mod_switch(q2)).collect(),
 57        }
 58    }
 59    // pub fn mod_switch(&self, Q2: u64) -> TR<crate::torus::T64, K> {
 60    //     TR(self.0.iter().map(|c_i| c_i.mod_switch(Q2)).collect())
 61    // }
 62}
 63impl TR<crate::ring_torus::Tn> {
 64    pub fn left_rotate(&self, h: usize) -> Self {
 65        TR {
 66            k: self.k,
 67            r: self.r.iter().map(|c_i| c_i.left_rotate(h)).collect(),
 68        }
 69    }
 70}
 71
 72impl<R: Ring> TR<R> {
 73    pub fn iter(&self) -> std::slice::Iter<R> {
 74        self.r.iter()
 75    }
 76}
 77
 78impl<R: Ring> Add<TR<R>> for TR<R> {
 79    type Output = Self;
 80    fn add(self, other: Self) -> Self {
 81        debug_assert_eq!(self.k, other.k);
 82
 83        Self {
 84            k: self.k,
 85            r: zip_eq(self.r, other.r)
 86                .map(|(s, o)| s + o)
 87                .collect::<Vec<_>>(),
 88        }
 89    }
 90}
 91
 92impl<R: Ring> Sub<TR<R>> for TR<R> {
 93    type Output = Self;
 94    fn sub(self, other: Self) -> Self {
 95        debug_assert_eq!(self.k, other.k);
 96
 97        Self {
 98            k: self.k,
 99            r: zip_eq(self.r, other.r).map(|(s, o)| s - o).collect(),
100        }
101    }
102}
103
104impl<R: Ring> Neg for TR<R> {
105    type Output = Self;
106
107    fn neg(self) -> Self::Output {
108        Self {
109            k: self.k,
110            r: self.r.iter().map(|e_i| -e_i.clone()).collect(),
111        }
112    }
113}
114
115/// for (TR,TR), the Mul operation is defined as the dot product:
116/// for A, B \in R^k, result = Σ A_i * B_i \in R
117impl<R: Ring> Mul<TR<R>> for TR<R> {
118    type Output = R;
119    fn mul(self, other: Self) -> R {
120        debug_assert_eq!(self.k, other.k);
121
122        zip_eq(self.r, other.r).map(|(s, o)| s * o).sum()
123    }
124}
125impl<R: Ring> Mul<&TR<R>> for &TR<R> {
126    type Output = R;
127    fn mul(self, other: &TR<R>) -> R {
128        debug_assert_eq!(self.k, other.k);
129
130        zip_eq(self.r.clone(), other.r.clone())
131            .map(|(s, o)| s * o)
132            .sum()
133    }
134}
135
136/// for (TR, R), the Mul operation is defined as each element of TR is
137/// multiplied by R
138impl<R: Ring> Mul<R> for TR<R> {
139    type Output = TR<R>;
140    fn mul(self, other: R) -> TR<R> {
141        Self {
142            k: self.k,
143            r: self.r.iter().map(|s| s.clone() * other.clone()).collect(),
144        }
145    }
146}
147impl<R: Ring> Mul<&R> for &TR<R> {
148    type Output = TR<R>;
149    fn mul(self, other: &R) -> TR<R> {
150        TR::<R> {
151            k: self.k,
152            r: self.r.iter().map(|s| s.clone() * other.clone()).collect(),
153        }
154    }
155}