1extern crate num;
  2extern crate num_bigint;
  3extern crate num_traits;
  4
  5use num_bigint::{BigInt, ToBigInt};
  6use num_traits::{One, Zero};
  7
  8pub fn modulus(a: &BigInt, m: &BigInt) -> BigInt {
  9    ((a % m) + m) % m
 10}
 11
 12pub fn modinv(a: &BigInt, q: &BigInt) -> BigInt {
 13    let mut mn = (q.clone(), a.clone());
 14    let mut xy: (BigInt, BigInt) = (Zero::zero(), One::one());
 15
 16    let big_zero: BigInt = Zero::zero();
 17    while mn.1 != big_zero {
 18        xy = (xy.1.clone(), xy.0 - (mn.0.clone() / mn.1.clone()) * xy.1);
 19        mn = (mn.1.clone(), modulus(&mn.0, &mn.1));
 20    }
 21
 22    while xy.0 < Zero::zero() {
 23        xy.0 = modulus(&xy.0, q);
 24    }
 25    xy.0
 26}
 27
 28/*
 29pub fn modinv_v2(a0: &BigInt, m0: &BigInt) -> BigInt {
 30    if m0 == &One::one() {
 31        return One::one();
 32    }
 33
 34    let (mut a, mut m, mut x0, mut inv): (BigInt, BigInt, BigInt, BigInt) =
 35        (a0.clone(), m0.clone(), Zero::zero(), One::one());
 36
 37    while a > One::one() {
 38        inv = inv - (&a / m.clone()) * x0.clone();
 39        a = a % m.clone();
 40        std::mem::swap(&mut a, &mut m);
 41        std::mem::swap(&mut x0, &mut inv);
 42    }
 43
 44    if inv < Zero::zero() {
 45        inv += m0.clone()
 46    }
 47    inv
 48}
 49
 50pub fn modinv_v3(a: &BigInt, q: &BigInt) -> BigInt {
 51    let mut aa: BigInt = a.clone();
 52    let mut qq: BigInt = q.clone();
 53    if qq < Zero::zero() {
 54        qq = -qq;
 55    }
 56    if aa < Zero::zero() {
 57        aa = -aa;
 58    }
 59    let d = num::Integer::gcd(&aa, &qq);
 60    if d != One::one() {
 61        println!("ERR no mod_inv");
 62    }
 63    let res: BigInt;
 64    if d < Zero::zero() {
 65        res = d + qq;
 66    } else {
 67        res = d;
 68    }
 69    res
 70}
 71pub fn modinv_v4(x: &BigInt, q: &BigInt) -> BigInt {
 72    let (gcd, inverse, _) = extended_gcd(x.clone(), q.clone());
 73    let one: BigInt = One::one();
 74    if gcd == one {
 75        modulus(&inverse, q)
 76    } else {
 77        panic!("error: gcd!=one")
 78    }
 79}
 80pub fn extended_gcd(a: BigInt, b: BigInt) -> (BigInt, BigInt, BigInt) {
 81    let (mut s, mut old_s) = (BigInt::zero(), BigInt::one());
 82    let (mut t, mut old_t) = (BigInt::one(), BigInt::zero());
 83    let (mut r, mut old_r) = (b, a);
 84
 85    while r != BigInt::zero() {
 86        let quotient = &old_r / &r;
 87        old_r -= &quotient * &r;
 88        std::mem::swap(&mut old_r, &mut r);
 89        old_s -= &quotient * &s;
 90        std::mem::swap(&mut old_s, &mut s);
 91        old_t -= quotient * &t;
 92        std::mem::swap(&mut old_t, &mut t);
 93    }
 94
 95    let _quotients = (t, s); // == (a, b) / gcd
 96
 97    (old_r, old_s, old_t)
 98}
 99*/
100
101pub fn concatenate_arrays<T: Clone>(x: &[T], y: &[T]) -> Vec<T> {
102    x.iter().chain(y).cloned().collect()
103}
104
105pub fn modsqrt(a: &BigInt, q: &BigInt) -> BigInt {
106    // Tonelli-Shanks Algorithm (https://en.wikipedia.org/wiki/Tonelli%E2%80%93Shanks_algorithm)
107    //
108    // This implementation is following the Go lang core implementation https://golang.org/src/math/big/int.go?s=23173:23210#L859
109    // Also described in https://www.maa.org/sites/default/files/pdf/upload_library/22/Polya/07468342.di020786.02p0470a.pdf
110    // -> section 6
111
112    let zero: BigInt = Zero::zero();
113    let one: BigInt = One::one();
114    if legendre_symbol(&a, q) != 1 {
115        // not a mod p square
116        return zero;
117    } else if a == &zero {
118        return zero;
119    } else if q == &2.to_bigint().unwrap() {
120        return zero;
121    } else if q % 4.to_bigint().unwrap() == 3.to_bigint().unwrap() {
122        let r = a.modpow(&((q + one) / 4), &q);
123        return r;
124    }
125
126    let mut s = q - &one;
127    let mut e: BigInt = Zero::zero();
128    while &s % 2 == zero {
129        s = s >> 1;
130        e = e + &one;
131    }
132
133    let mut n: BigInt = 2.to_bigint().unwrap();
134    while legendre_symbol(&n, q) != -1 {
135        n = &n + &one;
136    }
137
138    let mut y = a.modpow(&((&s + &one) >> 1), q);
139    let mut b = a.modpow(&s, q);
140    let mut g = n.modpow(&s, q);
141    let mut r = e;
142
143    loop {
144        let mut t = b.clone();
145        let mut m: BigInt = Zero::zero();
146        while &t != &one {
147            t = modulus(&(&t * &t), q);
148            m = m + &one;
149        }
150
151        if m == zero {
152            return y.clone();
153        }
154
155        t = g.modpow(&(2.to_bigint().unwrap().modpow(&(&r - &m - 1), q)), q);
156        g = g.modpow(&(2.to_bigint().unwrap().modpow(&(r - &m), q)), q);
157        y = modulus(&(y * t), q);
158        b = modulus(&(b * &g), q);
159        r = m.clone();
160    }
161}
162
163#[allow(dead_code)]
164pub fn modsqrt_v2(a: &BigInt, q: &BigInt) -> BigInt {
165    // Tonelli-Shanks Algorithm (https://en.wikipedia.org/wiki/Tonelli%E2%80%93Shanks_algorithm)
166    //
167    // This implementation is following this Python implementation by Dusk https://github.com/dusk-network/dusk-zerocaf/blob/master/tools/tonelli.py
168
169    let zero: BigInt = Zero::zero();
170    let one: BigInt = One::one();
171    if legendre_symbol(&a, q) != 1 {
172        // not a mod p square
173        return zero;
174    } else if a == &zero {
175        return zero;
176    } else if q == &2.to_bigint().unwrap() {
177        return zero;
178    } else if q % 4.to_bigint().unwrap() == 3.to_bigint().unwrap() {
179        let r = a.modpow(&((q + one) / 4), &q);
180        return r;
181    }
182
183    let mut p = q - &one;
184    let mut s: BigInt = Zero::zero();
185    while &p % 2.to_bigint().unwrap() == zero {
186        s = s + &one;
187        p = p >> 1;
188    }
189
190    let mut z: BigInt = One::one();
191    while legendre_symbol(&z, q) != -1 {
192        z = &z + &one;
193    }
194    let mut c = z.modpow(&p, q);
195
196    let mut x = a.modpow(&((&p + &one) >> 1), q);
197    let mut t = a.modpow(&p, q);
198    let mut m = s;
199
200    while &t != &one {
201        let mut i: BigInt = One::one();
202        let mut e: BigInt = 2.to_bigint().unwrap();
203        while i < m {
204            if t.modpow(&e, q) == one {
205                break;
206            }
207            e = e * 2.to_bigint().unwrap();
208            i = i + &one;
209        }
210
211        let b = c.modpow(&(2.to_bigint().unwrap().modpow(&(&m - &i - 1), q)), q);
212        x = modulus(&(x * &b), q);
213        t = modulus(&(t * &b * &b), q);
214        c = modulus(&(&b * &b), q);
215        m = i.clone();
216    }
217    return x;
218}
219
220pub fn legendre_symbol(a: &BigInt, q: &BigInt) -> i32 {
221    // returns 1 if has a square root modulo q
222    let one: BigInt = One::one();
223    let ls: BigInt = a.modpow(&((q - &one) >> 1), &q);
224    if &(ls) == &(q - one) {
225        return -1;
226    }
227    1
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233
234    #[test]
235    fn test_mod_inverse() {
236        let a = BigInt::parse_bytes(b"123456789123456789123456789123456789123456789", 10).unwrap();
237        let b = BigInt::parse_bytes(b"12345678", 10).unwrap();
238        assert_eq!(modinv(&a, &b), BigInt::parse_bytes(b"641883", 10).unwrap());
239    }
240
241    #[test]
242    fn test_sqrtmod() {
243        let a = BigInt::parse_bytes(
244            b"6536923810004159332831702809452452174451353762940761092345538667656658715568",
245            10,
246        )
247        .unwrap();
248        let q = BigInt::parse_bytes(
249            b"7237005577332262213973186563042994240857116359379907606001950938285454250989",
250            10,
251        )
252        .unwrap();
253
254        assert_eq!(
255            (modsqrt(&a, &q)).to_string(),
256            "5464794816676661649783249706827271879994893912039750480019443499440603127256"
257        );
258        assert_eq!(
259            (modsqrt_v2(&a, &q)).to_string(),
260            "5464794816676661649783249706827271879994893912039750480019443499440603127256"
261        );
262    }
263}