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