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