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    pub fn mod_switch<const Q2: u64>(&self) -> T64 {
 37        todo!()
 38    }
 39}
 40
 41impl Add<T64> for T64 {
 42    type Output = Self;
 43
 44    fn add(self, rhs: Self) -> Self::Output {
 45        Self(self.0.wrapping_add(rhs.0))
 46    }
 47}
 48impl AddAssign for T64 {
 49    fn add_assign(&mut self, rhs: Self) {
 50        self.0 = self.0.wrapping_add(rhs.0)
 51    }
 52}
 53
 54impl Sub<T64> for T64 {
 55    type Output = Self;
 56
 57    fn sub(self, rhs: Self) -> Self::Output {
 58        Self(self.0.wrapping_sub(rhs.0))
 59    }
 60}
 61impl SubAssign for T64 {
 62    fn sub_assign(&mut self, rhs: Self) {
 63        self.0 = self.0.wrapping_sub(rhs.0)
 64    }
 65}
 66
 67impl Neg for T64 {
 68    type Output = Self;
 69
 70    fn neg(self) -> Self::Output {
 71        Self(self.0.wrapping_neg())
 72    }
 73}
 74
 75impl Sum for T64 {
 76    fn sum<I>(iter: I) -> Self
 77    where
 78        I: Iterator<Item = Self>,
 79    {
 80        iter.fold(Self(0), |acc, x| acc + x)
 81    }
 82}
 83
 84impl Mul<T64> for T64 {
 85    type Output = Self;
 86
 87    fn mul(self, rhs: Self) -> Self::Output {
 88        Self(self.0.wrapping_mul(rhs.0))
 89    }
 90}
 91impl MulAssign for T64 {
 92    fn mul_assign(&mut self, rhs: Self) {
 93        self.0 = self.0.wrapping_mul(rhs.0)
 94    }
 95}
 96
 97// mul by u64
 98impl Mul<u64> for T64 {
 99    type Output = Self;
100
101    fn mul(self, s: u64) -> Self {
102        Self(self.0.wrapping_mul(s))
103    }
104}
105impl Mul<&u64> for &T64 {
106    type Output = T64;
107
108    fn mul(self, s: &u64) -> Self::Output {
109        T64(self.0.wrapping_mul(*s))
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use rand::distributions::Standard;
117
118    fn recompose(d: Vec<T64>) -> T64 {
119        T64(d.iter().fold(0u64, |acc, &b| (acc << 1) | b.0))
120    }
121    #[test]
122    fn test_decompose() {
123        let beta: u32 = 2;
124        let l: u32 = 64;
125
126        let x = T64(12345);
127        let d = x.decompose(beta, l);
128        assert_eq!(recompose(d), T64(12345));
129
130        let x = T64(0);
131        let d = x.decompose(beta, l);
132        assert_eq!(recompose(d), T64(0));
133
134        let x = T64(u64::MAX - 1);
135        let d = x.decompose(beta, l);
136        assert_eq!(recompose(d), T64(u64::MAX - 1));
137
138        let mut rng = rand::thread_rng();
139        for _ in 0..1000 {
140            let x = T64::rand(&mut rng, Standard);
141            let d = x.decompose(beta, l);
142            assert_eq!(recompose(d), x);
143        }
144    }
145}