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