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