1use rand::{distributions::Distribution, Rng};
  2use std::{
  3    iter::Sum,
  4    ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
  5};
  6
  7/// Let 𝕋 = ℝ/ℤ, where 𝕋 is a ℤ-module, with homogeneous external product.
  8/// Let 𝕋q
  9/// T64 is 𝕋q with q=2^Ω, with Ω=64. We identify 𝕋q=(1/q)ℤ/ℤ ≈ ℤq.
 10#[derive(Debug, Clone, Copy, PartialEq)]
 11pub struct T64(pub u64);
 12
 13impl T64 {
 14    pub fn zero() -> Self {
 15        Self(0u64)
 16    }
 17    pub fn rand(mut rng: impl Rng, dist: impl Distribution<u64>) -> Self {
 18        let r: u64 = dist.sample(&mut rng);
 19        Self(r)
 20    }
 21    pub fn rand_f64(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
 22        let r: f64 = dist.sample(&mut rng);
 23        Self(r.round() as u64)
 24    }
 25
 26    /// Note: only beta=2 and l=64 is supported.
 27    pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
 28        assert_eq!(beta, 2u32); // only beta=2 supported
 29        assert_eq!(l, 64u32); // only l=64 supported
 30
 31        (0..64)
 32            .rev()
 33            .map(|i| T64(((self.0 >> i) & 1) as u64))
 34            .collect()
 35    }
 36}
 37
 38impl Add<T64> for T64 {
 39    type Output = Self;
 40
 41    fn add(self, rhs: Self) -> Self::Output {
 42        Self(self.0.wrapping_add(rhs.0))
 43    }
 44}
 45impl AddAssign for T64 {
 46    fn add_assign(&mut self, rhs: Self) {
 47        self.0 += rhs.0;
 48    }
 49}
 50
 51impl Sub<T64> for T64 {
 52    type Output = Self;
 53
 54    fn sub(self, rhs: Self) -> Self::Output {
 55        Self(self.0.wrapping_sub(rhs.0))
 56    }
 57}
 58impl SubAssign for T64 {
 59    fn sub_assign(&mut self, rhs: Self) {
 60        self.0 -= rhs.0;
 61    }
 62}
 63
 64impl Neg for T64 {
 65    type Output = Self;
 66
 67    fn neg(self) -> Self::Output {
 68        Self(self.0.wrapping_neg())
 69    }
 70}
 71
 72impl Sum for T64 {
 73    fn sum<I>(iter: I) -> Self
 74    where
 75        I: Iterator<Item = Self>,
 76    {
 77        iter.fold(Self(0), |acc, x| acc + x)
 78    }
 79}
 80
 81impl Mul<T64> for T64 {
 82    type Output = Self;
 83
 84    fn mul(self, rhs: Self) -> Self::Output {
 85        Self(self.0.wrapping_mul(rhs.0))
 86    }
 87}
 88impl MulAssign for T64 {
 89    fn mul_assign(&mut self, rhs: Self) {
 90        self.0 *= rhs.0;
 91    }
 92}
 93
 94// mul by u64
 95impl Mul<u64> for T64 {
 96    type Output = Self;
 97
 98    fn mul(self, s: u64) -> Self {
 99        Self(self.0 * s)
100    }
101}
102impl Mul<&u64> for &T64 {
103    type Output = T64;
104
105    fn mul(self, s: &u64) -> Self::Output {
106        T64(self.0 * s)
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use rand::distributions::Standard;
114
115    fn recompose(d: Vec<T64>) -> T64 {
116        T64(d.iter().fold(0u64, |acc, &b| (acc << 1) | b.0))
117    }
118    #[test]
119    fn test_decompose() {
120        let beta: u32 = 2;
121        let l: u32 = 64;
122
123        let x = T64(12345);
124        let d = x.decompose(beta, l);
125        assert_eq!(recompose(d), T64(12345));
126
127        let x = T64(0);
128        let d = x.decompose(beta, l);
129        assert_eq!(recompose(d), T64(0));
130
131        let x = T64(u64::MAX - 1);
132        let d = x.decompose(beta, l);
133        assert_eq!(recompose(d), T64(u64::MAX - 1));
134
135        let mut rng = rand::thread_rng();
136        for _ in 0..1000 {
137            let x = T64::rand(&mut rng, Standard);
138            let d = x.decompose(beta, l);
139            assert_eq!(recompose(d), x);
140        }
141    }
142}