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    fn mod_switch<const Q2: u64>(&self) -> T64 {
 54        todo!()
 55    }
 56
 57    fn mul_div_round(&self, num: u64, den: u64) -> Self {
 58        T64(((num as f64 * self.0 as f64) / den as f64).round() as u64)
 59    }
 60}
 61
 62impl T64 {
 63    pub fn rand_u64(mut rng: impl Rng, dist: impl Distribution<u64>) -> Self {
 64        let r: u64 = dist.sample(&mut rng);
 65        Self(r)
 66    }
 67}
 68
 69impl Add<T64> for T64 {
 70    type Output = Self;
 71
 72    fn add(self, rhs: Self) -> Self::Output {
 73        Self(self.0.wrapping_add(rhs.0))
 74    }
 75}
 76impl AddAssign for T64 {
 77    fn add_assign(&mut self, rhs: Self) {
 78        self.0 = self.0.wrapping_add(rhs.0)
 79    }
 80}
 81
 82impl Sub<T64> for T64 {
 83    type Output = Self;
 84
 85    fn sub(self, rhs: Self) -> Self::Output {
 86        Self(self.0.wrapping_sub(rhs.0))
 87    }
 88}
 89impl SubAssign for T64 {
 90    fn sub_assign(&mut self, rhs: Self) {
 91        self.0 = self.0.wrapping_sub(rhs.0)
 92    }
 93}
 94
 95impl Neg for T64 {
 96    type Output = Self;
 97
 98    fn neg(self) -> Self::Output {
 99        Self(self.0.wrapping_neg())
100    }
101}
102
103impl Sum for T64 {
104    fn sum<I>(iter: I) -> Self
105    where
106        I: Iterator<Item = Self>,
107    {
108        iter.fold(Self(0), |acc, x| acc + x)
109    }
110}
111
112impl Mul<T64> for T64 {
113    type Output = Self;
114
115    fn mul(self, rhs: Self) -> Self::Output {
116        Self(self.0.wrapping_mul(rhs.0))
117    }
118}
119impl MulAssign for T64 {
120    fn mul_assign(&mut self, rhs: Self) {
121        self.0 = self.0.wrapping_mul(rhs.0)
122    }
123}
124
125// mul by u64
126impl Mul<u64> for T64 {
127    type Output = Self;
128
129    fn mul(self, s: u64) -> Self {
130        if self.0 == 0 {
131            return Self(0);
132        }
133        Self(self.0.wrapping_mul(s))
134    }
135}
136impl Mul<&u64> for &T64 {
137    type Output = T64;
138
139    fn mul(self, s: &u64) -> Self::Output {
140        T64(self.0.wrapping_mul(*s))
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use rand::distributions::Standard;
148
149    fn recompose(d: Vec<T64>) -> T64 {
150        T64(d.iter().fold(0u64, |acc, &b| (acc << 1) | b.0))
151    }
152    #[test]
153    fn test_decompose() {
154        let beta: u32 = 2;
155        let l: u32 = 64;
156
157        let x = T64(12345);
158        let d = x.decompose(beta, l);
159        assert_eq!(recompose(d), T64(12345));
160
161        let x = T64(0);
162        let d = x.decompose(beta, l);
163        assert_eq!(recompose(d), T64(0));
164
165        let x = T64(u64::MAX - 1);
166        let d = x.decompose(beta, l);
167        assert_eq!(recompose(d), T64(u64::MAX - 1));
168
169        let mut rng = rand::thread_rng();
170        for _ in 0..1000 {
171            let x = T64::rand(&mut rng, Standard);
172            let d = x.decompose(beta, l);
173            assert_eq!(recompose(d), x);
174        }
175    }
176}