p3_field/
sqrt.rs

1//! Square roots in finite fields via the Tonelli–Shanks algorithm.
2
3use num_bigint::BigUint;
4
5use crate::{Field, TwoAdicField};
6
7/// Compute `base^exponent` for an arbitrarily large exponent.
8///
9/// Uses the standard square-and-multiply approach over the bits of `exponent`.
10/// This accepts exponents which do not fit in a `u64`, as required by fields
11/// whose order exceeds `2^64`.
12fn exp_biguint<F: Field>(base: F, exponent: &BigUint) -> F {
13    let mut product = F::ONE;
14    let mut current = base;
15    for j in 0..exponent.bits() {
16        if exponent.bit(j) {
17            product *= current;
18        }
19        current = current.square();
20    }
21    product
22}
23
24/// The core of the Tonelli–Shanks algorithm.
25///
26/// Given a nonzero `a`, the decomposition `|F| - 1 = q * 2^s` with `q` odd, and a
27/// primitive `2^s`-th root of unity `c` (the `2`-Sylow generator), return a square
28/// root of `a` if one exists.
29///
30/// `c` must be a quadratic non-residue raised to the `q`-th power; both
31/// [`tonelli_shanks`] and [`tonelli_shanks_two_adic`] provide such a value.
32fn tonelli_shanks_inner<F: Field>(a: F, s: usize, q: &BigUint, mut c: F) -> Option<F> {
33    // A single exponentiation yields both candidates: with `u = a^((q-1)/2)`,
34    // `r = u * a = a^((q+1)/2)` is the prospective root and `t = r * u = a^q`
35    // tracks the residue's `2`-power component.
36    let u = exp_biguint(a, &((q - 1u32) >> 1));
37    let mut r = u * a;
38    let mut t = r * u;
39    let mut m = s;
40
41    while !t.is_one() {
42        // Find the least `i`, with `0 < i < m`, such that `t^(2^i) == 1`.
43        let mut i = 0;
44        let mut t2i = t;
45        while !t2i.is_one() {
46            t2i = t2i.square();
47            i += 1;
48            if i == m {
49                // `t` has order `2^m`, which only happens when `a` is a
50                // quadratic non-residue, so no square root exists.
51                return None;
52            }
53        }
54
55        let b = c.exp_power_of_2(m - i - 1);
56        m = i;
57        c = b.square();
58        t *= c;
59        r *= b;
60    }
61
62    Some(r)
63}
64
65/// Return a square root of `a` if one exists, otherwise `None`.
66///
67/// This is the generic Tonelli–Shanks algorithm. Writing the multiplicative
68/// group order as `|F| - 1 = q * 2^s` with `q` odd, it uses [`Field::GENERATOR`]
69/// (a generator of `F^*`, hence a quadratic non-residue) to seed the `2`-Sylow
70/// subgroup via `c = GENERATOR^q`.
71///
72/// For a quadratic residue this returns one of its two square roots; which one
73/// is unspecified. `ZERO` maps to `ZERO`.
74pub fn tonelli_shanks<F: Field>(a: F) -> Option<F> {
75    // Zero is its own square root and would otherwise break the logic below.
76    if a.is_zero() {
77        return Some(F::ZERO);
78    }
79
80    // Write `|F| - 1 = q * 2^s` with `q` odd.
81    let order_minus_one = F::order() - BigUint::from(1u8);
82    // `|F| >= 2`, so `order_minus_one >= 1` and `trailing_zeros` is never `None`.
83    let s = order_minus_one
84        .trailing_zeros()
85        .expect("field order must be at least two") as usize;
86    let q = &order_minus_one >> s;
87
88    let c = exp_biguint(F::GENERATOR, &q);
89    tonelli_shanks_inner(a, s, &q, c)
90}
91
92/// Return a square root of `a` if one exists, otherwise `None`.
93///
94/// A variant of [`tonelli_shanks`] for two-adic fields. It seeds the `2`-Sylow
95/// subgroup directly from [`TwoAdicField::two_adic_generator`] (a primitive
96/// `2^TWO_ADICITY`-th root of unity), avoiding the `GENERATOR^q` exponentiation
97/// that the generic version performs on every call.
98///
99/// For a quadratic residue this returns one of its two square roots; which one
100/// is unspecified. `ZERO` maps to `ZERO`.
101pub fn tonelli_shanks_two_adic<F: TwoAdicField>(a: F) -> Option<F> {
102    // Zero is its own square root and would otherwise break the logic below.
103    if a.is_zero() {
104        return Some(F::ZERO);
105    }
106
107    // For a two-adic field, `s` is exactly `TWO_ADICITY` and the `2`-Sylow
108    // generator is available as a constant.
109    let s = F::TWO_ADICITY;
110    let q = (F::order() - BigUint::from(1u8)) >> s;
111    let c = F::two_adic_generator(s);
112    tonelli_shanks_inner(a, s, &q, c)
113}
p3_field/sqrt.rs

p3_field/
sqrt.rs
