haar_lib/math/mod_ops/
sqrt.rs

1//! x² = a (mod p)を満たすxを一つ求める。
2
3use crate::math::mod_ops::pow::*;
4use crate::rand::rand;
5
6/// x² = a (mod p)を満たすxを一つ求める。
7pub fn mod_sqrt(a: u64, p: u64) -> Option<u64> {
8    if p == 2 {
9        return Some(a % 2);
10    }
11    if a == 0 {
12        return Some(0);
13    }
14
15    let b = mod_pow(a, (p - 1) / 2, p);
16
17    if b == p - 1 {
18        return None;
19    }
20    if p % 4 == 3 {
21        return Some(mod_pow(a, (p + 1) / 4, p));
22    }
23
24    let mut q = p - 1;
25    let mut s = 0;
26    while q % 2 == 0 {
27        q /= 2;
28        s += 1;
29    }
30
31    let z = {
32        let ret;
33        loop {
34            let z = rand() % p;
35            if mod_pow(z, (p - 1) / 2, p) == p - 1 {
36                ret = z;
37                break;
38            }
39        }
40        ret
41    };
42
43    let mut m = s;
44    let mut c = mod_pow(z, q, p);
45    let mut t = mod_pow(a, q, p);
46    let mut r = mod_pow(a, q.div_ceil(2), p);
47
48    loop {
49        if t == 0 {
50            return Some(0);
51        }
52        if t == 1 {
53            return Some(r);
54        }
55
56        let mut i = 1;
57        let mut k = t;
58        while i < m {
59            k *= k;
60            k %= p;
61            if k == 1 {
62                break;
63            }
64
65            i += 1;
66        }
67
68        let b = mod_pow(c, 1 << (m - i - 1), p);
69
70        m = i;
71        c = b * b % p;
72        t *= b * b % p;
73        t %= p;
74        r *= b;
75        r %= p;
76    }
77}