1use std::fmt;
  2use std::ops;
  3
  4// Z_q, integers modulus q, not necessarily prime
  5#[derive(Clone, Copy, PartialEq)]
  6pub struct Zq<const Q: u64>(pub u64);
  7
  8// WIP
  9// impl<const Q: u64> From<Vec<u64>> for Vec<Zq<Q>> {
 10//     fn from(v: Vec<u64>) -> Self {
 11//         v.into_iter().map(Zq::new).collect()
 12//     }
 13// }
 14
 15pub(crate) fn modulus_u64<const Q: u64>(e: u64) -> u64 {
 16    (e % Q + Q) % Q
 17}
 18impl<const Q: u64> Zq<Q> {
 19    pub fn from_u64(e: u64) -> Self {
 20        if e >= Q {
 21            // (e % Q + Q) % Q
 22            return Zq(modulus_u64::<Q>(e));
 23            // return Zq(e % Q);
 24        }
 25        Zq(e)
 26    }
 27    pub fn from_f64(e: f64) -> Self {
 28        // WIP method
 29        let e: i64 = e.round() as i64;
 30        let q = Q as i64;
 31        if e < 0 || e >= q {
 32            return Zq(((e % q + q) % q) as u64);
 33        }
 34        Zq(e as u64)
 35
 36        // if e < 0 {
 37        //     // dbg!(&e);
 38        //     // dbg!(Zq::<Q>(((Q as i64 + e) % Q as i64) as u64));
 39        //     // return Zq(((Q as i64 + e) % Q as i64) as u64);
 40        //     return Zq(e as u64 % Q);
 41        // } else if e >= Q as i64 {
 42        //     return Zq((e % Q as i64) as u64);
 43        // }
 44        // Zq(e as u64)
 45    }
 46    pub fn from_bool(b: bool) -> Self {
 47        if b { Zq(1) } else { Zq(0) }
 48    }
 49    pub fn zero() -> Self {
 50        Zq(0u64)
 51    }
 52    pub fn square(self) -> Self {
 53        self * self
 54    }
 55    // modular exponentiation
 56    pub fn exp(self, e: Self) -> Self {
 57        // mul-square approach
 58        let mut res = Self(1);
 59        let mut rem = e.clone();
 60        let mut exp = self;
 61        // for rem != Self(0) {
 62        while rem != Self(0) {
 63            // if odd
 64            // TODO use a more readible expression
 65            if 1 - ((rem.0 & 1) << 1) as i64 == -1 {
 66                res = res * exp;
 67            }
 68            exp = exp.square();
 69            rem = Self(rem.0 >> 1);
 70        }
 71        res
 72    }
 73    // multiplicative inverse
 74    // WARNING: if this is needed, it means that 'Zq' is a Finite Field. For the moment we assume
 75    // we work in a Finite Field
 76    pub fn inv_OLD(self) -> Self {
 77        // TODO
 78        // let a = self.0;
 79        // let q = Q;
 80        let mut t = 0;
 81        let mut r = Q;
 82        let mut new_t = 0;
 83        let mut new_r = self.0.clone();
 84        while new_r != 0 {
 85            let q = r / new_r;
 86
 87            t = new_t.clone();
 88            new_t = t - q;
 89
 90            r = new_r.clone();
 91            new_r = r - (q * new_r);
 92        }
 93        // if t < 0 {
 94        //     t = t + q;
 95        // }
 96        return Zq::from_u64(t);
 97    }
 98    pub fn inv(self) -> Zq<Q> {
 99        let (g, x, _) = Self::egcd(self.0 as i128, Q as i128);
100        if g != 1 {
101            // None
102            panic!("E");
103        } else {
104            let q = Q as i128;
105            Zq(((x % q + q) % q) as u64) // TODO maybe just Zq::new(x)
106        }
107    }
108    fn egcd(a: i128, b: i128) -> (i128, i128, i128) {
109        if a == 0 {
110            (b, 0, 1)
111        } else {
112            let (g, x, y) = Self::egcd(b % a, a);
113            (g, y - (b / a) * x, x)
114        }
115    }
116}
117
118impl<const Q: u64> Zq<Q> {
119    fn r#mod(self) -> Self {
120        if self.0 >= Q {
121            return Zq(self.0 % Q);
122        }
123        self
124    }
125}
126
127impl<const Q: u64> ops::Add<Zq<Q>> for Zq<Q> {
128    type Output = Self;
129
130    fn add(self, rhs: Self) -> Self::Output {
131        let mut r = self.0 + rhs.0;
132        if r >= Q {
133            r -= Q;
134        }
135        Zq(r)
136    }
137}
138impl<const Q: u64> ops::Add<&Zq<Q>> for &Zq<Q> {
139    type Output = Zq<Q>;
140
141    fn add(self, rhs: &Zq<Q>) -> Self::Output {
142        let mut r = self.0 + rhs.0;
143        if r >= Q {
144            r -= Q;
145        }
146        Zq(r)
147    }
148}
149impl<const Q: u64> ops::AddAssign<Zq<Q>> for Zq<Q> {
150    fn add_assign(&mut self, rhs: Self) {
151        *self = *self + rhs
152    }
153}
154impl<const Q: u64> std::iter::Sum for Zq<Q> {
155    fn sum<I>(iter: I) -> Self
156    where
157        I: Iterator<Item = Self>,
158    {
159        iter.fold(Zq(0), |acc, x| acc + x)
160    }
161}
162impl<const Q: u64> ops::Sub<Zq<Q>> for Zq<Q> {
163    type Output = Self;
164
165    fn sub(self, rhs: Self) -> Zq<Q> {
166        if self.0 >= rhs.0 {
167            Zq(self.0 - rhs.0)
168        } else {
169            Zq((Q + self.0) - rhs.0)
170        }
171    }
172}
173impl<const Q: u64> ops::Sub<&Zq<Q>> for &Zq<Q> {
174    type Output = Zq<Q>;
175
176    fn sub(self, rhs: &Zq<Q>) -> Self::Output {
177        if self.0 >= rhs.0 {
178            Zq(self.0 - rhs.0)
179        } else {
180            Zq((Q + self.0) - rhs.0)
181        }
182    }
183}
184impl<const Q: u64> ops::SubAssign<Zq<Q>> for Zq<Q> {
185    fn sub_assign(&mut self, rhs: Self) {
186        *self = *self - rhs
187    }
188}
189impl<const Q: u64> ops::Neg for Zq<Q> {
190    type Output = Self;
191
192    fn neg(self) -> Self::Output {
193        Zq(Q - self.0)
194    }
195}
196impl<const Q: u64> ops::Mul<Zq<Q>> for Zq<Q> {
197    type Output = Self;
198
199    fn mul(self, rhs: Self) -> Zq<Q> {
200        // TODO non-naive way
201        Zq(((self.0 as u128 * rhs.0 as u128) % Q as u128) as u64)
202        // Zq((self.0 * rhs.0) % Q)
203    }
204}
205impl<const Q: u64> ops::Div<Zq<Q>> for Zq<Q> {
206    type Output = Self;
207
208    fn div(self, rhs: Self) -> Zq<Q> {
209        // TODO non-naive way
210        // Zq((self.0 / rhs.0) % Q)
211        self * rhs.inv()
212    }
213}
214
215impl<const Q: u64> fmt::Display for Zq<Q> {
216    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
217        write!(f, "{}", self.0)
218    }
219}
220impl<const Q: u64> fmt::Debug for Zq<Q> {
221    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
222        write!(f, "{}", self.0)
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    #[test]
231    fn exp() {
232        const Q: u64 = 1021;
233        let a = Zq::<Q>(3);
234        let b = Zq::<Q>(3);
235        assert_eq!(a.exp(b), Zq(27));
236
237        let a = Zq::<Q>(1000);
238        let b = Zq::<Q>(3);
239        assert_eq!(a.exp(b), Zq(949));
240    }
241    #[test]
242    fn neg() {
243        const Q: u64 = 1021;
244        let a = Zq::<Q>::from_f64(101.0);
245        let b = Zq::<Q>::from_f64(-1.0);
246        assert_eq!(-a, a * b);
247    }
248}