1use rand::{distributions::Distribution, Rng};
  2use std::{
  3    iter::Sum,
  4    ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
  5};
  6
  7use crate::ring::Ring;
  8
  9/// Let 𝕋 = ℝ/ℤ, where 𝕋 is a ℤ-module, with homogeneous external product.
 10/// Let 𝕋q
 11/// T64 is 𝕋q with q=2^Ω, with Ω=64. We identify 𝕋q=(1/q)ℤ/ℤ ≈ ℤq.
 12#[derive(Debug, Clone, Copy, PartialEq)]
 13pub struct T64(pub u64);
 14
 15// implement the `Ring` trait for T64, so that it can be used where we would use
 16// `Tn<1>`.
 17impl Ring for T64 {
 18    type C = T64;
 19    type Params = ();
 20
 21    // const Q: u64 = u64::MAX; // WIP
 22    // const N: usize = 1;
 23
 24    fn coeffs(&self) -> Vec<T64> {
 25        vec![self.clone()]
 26    }
 27    fn zero(_: ()) -> Self {
 28        Self(0u64)
 29    }
 30    fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, _: ()) -> Self {
 31        let r: f64 = dist.sample(&mut rng);
 32        Self(r.round() as u64)
 33    }
 34    fn from_vec(_n: (), coeffs: Vec<Self::C>) -> Self {
 35        assert_eq!(coeffs.len(), 1);
 36        coeffs[0]
 37    }
 38
 39    // TODO rm beta & l from inputs, make it always beta=2,l=64.
 40    /// Note: only beta=2 and l=64 is supported.
 41    fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
 42        assert_eq!(beta, 2u32); // only beta=2 supported
 43                                // assert_eq!(l, 64u32); // only l=64 supported
 44
 45        // (0..64)
 46        (0..l)
 47            .rev()
 48            .map(|i| T64(((self.0 >> i) & 1) as u64))
 49            .collect()
 50    }
 51    fn remodule(&self, p: u64) -> T64 {
 52        todo!()
 53    }
 54
 55    // modulus switch from Q to Q2: self * Q2/Q
 56    fn mod_switch(&self, q2: u64) -> T64 {
 57        // for the moment we assume Q|Q2, since Q=2^64, check that Q2 is a power
 58        // of two:
 59        assert!(q2.is_power_of_two());
 60        // since Q=2^64, dividing Q2/Q is equivalent to dividing 2^log2(Q2)/2^64
 61        // which would be like right-shifting 64-log2(Q2).
 62        let log2_q2 = 63 - q2.leading_zeros();
 63        T64(self.0 >> (64 - log2_q2))
 64    }
 65
 66    fn mul_div_round(&self, num: u64, den: u64) -> Self {
 67        T64(((num as f64 * self.0 as f64) / den as f64).round() as u64)
 68    }
 69}
 70
 71impl T64 {
 72    pub fn rand_u64(mut rng: impl Rng, dist: impl Distribution<u64>) -> Self {
 73        let r: u64 = dist.sample(&mut rng);
 74        Self(r)
 75    }
 76}
 77
 78impl Add<T64> for T64 {
 79    type Output = Self;
 80
 81    fn add(self, rhs: Self) -> Self::Output {
 82        Self(self.0.wrapping_add(rhs.0))
 83    }
 84}
 85impl AddAssign for T64 {
 86    fn add_assign(&mut self, rhs: Self) {
 87        self.0 = self.0.wrapping_add(rhs.0)
 88    }
 89}
 90
 91impl Sub<T64> for T64 {
 92    type Output = Self;
 93
 94    fn sub(self, rhs: Self) -> Self::Output {
 95        Self(self.0.wrapping_sub(rhs.0))
 96    }
 97}
 98impl SubAssign for T64 {
 99    fn sub_assign(&mut self, rhs: Self) {
100        self.0 = self.0.wrapping_sub(rhs.0)
101    }
102}
103
104impl Neg for T64 {
105    type Output = Self;
106
107    fn neg(self) -> Self::Output {
108        Self(self.0.wrapping_neg())
109    }
110}
111
112impl Sum for T64 {
113    fn sum<I>(iter: I) -> Self
114    where
115        I: Iterator<Item = Self>,
116    {
117        iter.fold(Self(0), |acc, x| acc + x)
118    }
119}
120
121impl Mul<T64> for T64 {
122    type Output = Self;
123
124    fn mul(self, rhs: Self) -> Self::Output {
125        Self(self.0.wrapping_mul(rhs.0))
126    }
127}
128impl MulAssign for T64 {
129    fn mul_assign(&mut self, rhs: Self) {
130        self.0 = self.0.wrapping_mul(rhs.0)
131    }
132}
133
134// mul by u64
135impl Mul<u64> for T64 {
136    type Output = Self;
137
138    fn mul(self, s: u64) -> Self {
139        if self.0 == 0 {
140            return Self(0);
141        }
142        Self(self.0.wrapping_mul(s))
143    }
144}
145impl Mul<&u64> for &T64 {
146    type Output = T64;
147
148    fn mul(self, s: &u64) -> Self::Output {
149        T64(self.0.wrapping_mul(*s))
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use rand::distributions::Standard;
157
158    fn recompose(d: Vec<T64>) -> T64 {
159        T64(d.iter().fold(0u64, |acc, &b| (acc << 1) | b.0))
160    }
161    #[test]
162    fn test_decompose() {
163        let beta: u32 = 2;
164        let l: u32 = 64;
165
166        let x = T64(12345);
167        let d = x.decompose(beta, l);
168        assert_eq!(recompose(d), T64(12345));
169
170        let x = T64(0);
171        let d = x.decompose(beta, l);
172        assert_eq!(recompose(d), T64(0));
173
174        let x = T64(u64::MAX - 1);
175        let d = x.decompose(beta, l);
176        assert_eq!(recompose(d), T64(u64::MAX - 1));
177
178        let mut rng = rand::thread_rng();
179        for _ in 0..1000 {
180            let x = T64::rand(&mut rng, Standard, ());
181            let d = x.decompose(beta, l);
182            assert_eq!(recompose(d), x);
183        }
184    }
185}