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