haar_lib/math/mod_ops/
sqrt.rs1use crate::math::mod_ops::pow::*;
4use crate::rand::rand;
5
6pub 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}