S
```/// Cipolla algorithm
///
///     x^2 = a (mod p) , p is an odd prime
/// with O(M*log(n)) time complexity, M depends on the complexity of complex numbers multiplication.
///
/// Wikipedia reference: https://en.wikipedia.org/wiki/Cipolla%27s_algorithm
/// When a is the primitive root modulo n, the answer is unique.
/// Otherwise it will return the smallest positive solution
use std::rc::Rc;
use std::time::{SystemTime, UNIX_EPOCH};

use rand::Rng;

use super::{fast_power, PCG32};

#[derive(Debug)]
struct CustomFiniteField {
modulus: u64,
i_square: u64,
}

impl CustomFiniteField {
pub fn new(modulus: u64, i_square: u64) -> Self {
Self { modulus, i_square }
}
}

#[derive(Clone, Debug)]
struct CustomComplexNumber {
real: u64,
imag: u64,
f: Rc<CustomFiniteField>,
}

impl CustomComplexNumber {
pub fn new(real: u64, imag: u64, f: Rc<CustomFiniteField>) -> Self {
Self { real, imag, f }
}

pub fn mult_other(&mut self, rhs: &Self) {
let tmp = (self.imag * rhs.real + self.real * rhs.imag) % self.f.modulus;
self.real = (self.real * rhs.real
+ ((self.imag * rhs.imag) % self.f.modulus) * self.f.i_square)
% self.f.modulus;
self.imag = tmp;
}

pub fn mult_self(&mut self) {
let tmp = (self.imag * self.real + self.real * self.imag) % self.f.modulus;
self.real = (self.real * self.real
+ ((self.imag * self.imag) % self.f.modulus) * self.f.i_square)
% self.f.modulus;
self.imag = tmp;
}

pub fn fast_power(mut base: Self, mut power: u64) -> Self {
let mut result = CustomComplexNumber::new(1, 0, base.f.clone());
while power != 0 {
if (power & 1) != 0 {
result.mult_other(&base); // result *= base;
}
base.mult_self(); // base *= base;
power >>= 1;
}
result
}
}

fn is_residue(x: u64, modulus: u64) -> bool {
let power = (modulus - 1) >> 1;
x != 0 && fast_power(x as usize, power as usize, modulus as usize) == 1
}

/// The Legendre symbol `(a | p)`
///
/// Returns 0 if a = 0 mod p, 1 if a is a square mod p, -1 if it not a square mod p.
///
/// <https://en.wikipedia.org/wiki/Legendre_symbol>
pub fn legendre_symbol(a: u64, odd_prime: u64) -> i64 {
debug_assert!(odd_prime % 2 != 0, "prime must be odd");
if a == 0 {
0
} else if is_residue(a, odd_prime) {
1
} else {
-1
}
}

// return two solutions (x1, x2) for Quadratic Residue problem x^2 = a (mod p), where p is an odd prime
// if a is Quadratic Nonresidues, return None
pub fn cipolla(a: u32, p: u32, seed: Option<u64>) -> Option<(u32, u32)> {
// The params should be kept in u32 range for multiplication overflow issue
// But inside we use u64 for convenience
let a = a as u64;
let p = p as u64;
if a == 0 {
return Some((0, 0));
}
if !is_residue(a, p) {
return None;
}
let seed = match seed {
Some(seed) => seed,
None => SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
};
let mut rng = PCG32::new_default(seed);
let r = loop {
let r = rng.get_u64() % p;
if r == 0 || !is_residue((p + r * r - a) % p, p) {
break r;
}
};
let filed = Rc::new(CustomFiniteField::new(p, (p + r * r - a) % p));
let comp = CustomComplexNumber::new(r, 1, filed);
let power = (p + 1) >> 1;
let x0 = CustomComplexNumber::fast_power(comp, power).real as u32;
let x1 = p as u32 - x0;
if x0 < x1 {
Some((x0, x1))
} else {
Some((x1, x0))
}
}

/// Returns one of the two possible solutions of _x² = a mod p_, if any.
///
/// The other solution is _-x mod p_. If there is no solution, returns `None`.
///
/// Reference: H. Cohen, _A course in computational algebraic number theory_, Algorithm 1.4.3
///
/// ## Implementation details
///
/// To avoid multiplication overflows, internally the algorithm uses the `128`-bit arithmetic.
///
/// Also see [`cipolla`].
pub fn tonelli_shanks(a: i64, odd_prime: u64) -> Option<u64> {
let p: u128 = odd_prime as u128;
let e = (p - 1).trailing_zeros();
let q = (p - 1) >> e; // p = 2^e * q, with q odd

let a = if a < 0 {
a.rem_euclid(p as i64) as u128
} else {
a as u128
};

let power_mod_p = |b, e| fast_power(b as usize, e as usize, p as usize) as u128;

// find generator: choose a random non-residue n mod p
let n = loop {
let n = rng.gen_range(0..p);
if legendre_symbol(n as u64, p as u64) == -1 {
break n;
}
};
let z = power_mod_p(n, q);

// init
let mut y = z;
let mut r = e;
let mut x = power_mod_p(a, (q - 1) / 2) % p;
let mut b = (a * x * x) % p;
x = (a * x) % p;

while b % p != 1 {
// find exponent
let m = (1..r)
.scan(b, |prev, m| {
*prev = (*prev * *prev) % p;
Some((m, *prev == 1))
})
.find_map(|(m, cond)| cond.then_some(m));
let Some(m) = m else {
return None; // non-residue
};

// reduce exponent
let t = power_mod_p(y as u128, 2_u128.pow(r - m - 1));
y = (t * t) % p;
r = m;
x = (x * t) % p;
b = (b * y) % p;
}

Some(x as u64)
}

#[cfg(test)]
mod tests {
use super::*;

fn tonelli_shanks_residues(x: u64, odd_prime: u64) -> Option<(u64, u64)> {
let x = tonelli_shanks(x as i64, odd_prime)?;
let x2 = (-(x as i64)).rem_euclid(odd_prime as i64) as u64;
Some(if x < x2 { (x, x2) } else { (x2, x) })
}

#[test]
fn cipolla_small_numbers() {
assert_eq!(cipolla(1, 43, None), Some((1, 42)));
assert_eq!(cipolla(2, 23, None), Some((5, 18)));
assert_eq!(cipolla(17, 83, Some(42)), Some((10, 73)));
}

#[test]
fn tonelli_shanks_small_numbers() {
assert_eq!(tonelli_shanks_residues(1, 43).unwrap(), (1, 42));
assert_eq!(tonelli_shanks_residues(2, 23).unwrap(), (5, 18));
assert_eq!(tonelli_shanks_residues(17, 83).unwrap(), (10, 73));
}

#[test]
fn cipolla_random_numbers() {
assert_eq!(cipolla(392203, 852167, None), Some((413252, 438915)));
assert_eq!(
cipolla(379606557, 425172197, None),
Some((143417827, 281754370))
);
assert_eq!(
cipolla(585251669, 892950901, None),
Some((192354555, 700596346))
);
assert_eq!(
cipolla(404690348, 430183399, Some(19260817)),
Some((57227138, 372956261))
);
assert_eq!(
cipolla(210205747, 625380647, Some(998244353)),
Some((76810367, 548570280))
);
}

#[test]
fn tonelli_shanks_random_numbers() {
assert_eq!(
tonelli_shanks_residues(392203, 852167),
Some((413252, 438915))
);
assert_eq!(
tonelli_shanks_residues(379606557, 425172197),
Some((143417827, 281754370))
);
assert_eq!(
tonelli_shanks_residues(585251669, 892950901),
Some((192354555, 700596346))
);
assert_eq!(
tonelli_shanks_residues(404690348, 430183399),
Some((57227138, 372956261))
);
assert_eq!(
tonelli_shanks_residues(210205747, 625380647),
Some((76810367, 548570280))
);
}

#[test]