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