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