bls12_381 0.9.1

//! Table-based square roots (Sarkar 2020), ported from pasta_curves 0.6.1
//! (`src/arithmetic/fields.rs`, MIT OR Apache-2.0, Copyright (c) 2020-2025 The Electric Coin
//! Company). Everything below the imports is pasta_curves' code unchanged; the only
//! bls12_381-specific inputs are `Scalar`'s `SqrtTableHelpers` impl and its hash multiplier
//! (`scalar.rs`).
//!
//! Timing is pasta_curves': every input runs the same sequence of field operations
//! (exponentiation by the constant (t - 1) / 2, then four fixed 8-bit table steps, no early
//! exits), `is_square` is a `subtle::Choice` and `sqrt` returns a `CtOption`; table entries are
//! selected by plain indexing on values derived from the input, so the memory access pattern
//! depends on the input, as in pasta_curves. The root returned may be the negation of the
//! Tonelli-Shanks root (the choice of root is unspecified).

use alloc::{boxed::Box, vec::Vec};
use core::marker::PhantomData;

use subtle::Choice;

// pasta_curves: `const_assert!(size_of::<usize>() >= 4);` (static_assertions).
const _: () = assert!(core::mem::size_of::<usize>() >= 4);

#[cfg(feature = "sqrt-table")]
const SQRT_TABLE_DIGIT_BITS: u32 = 8;
#[cfg(feature = "sqrt-table")]
const SQRT_TABLE_RADIX: usize = 1 << SQRT_TABLE_DIGIT_BITS;
#[cfg(feature = "sqrt-table")]
const SQRT_TABLE_TOP_LEN: usize = SQRT_TABLE_RADIX / 2 + 1;
#[cfg(feature = "sqrt-table")]
const SQRT_HASH_BITS: u32 = 10;
#[cfg(feature = "sqrt-table")]
const SQRT_HASH_SHIFT: u32 = u32::BITS - SQRT_HASH_BITS;
#[cfg(feature = "sqrt-table")]
const SQRT_HASH_TABLE_LEN: usize = 1 << SQRT_HASH_BITS;

/// An internal trait that exposes additional operations related to calculating square roots of
/// prime-order finite fields.
#[cfg(feature = "sqrt-table")]
pub(crate) trait SqrtTableHelpers: ff::PrimeField {
    /// Raise this field element to the power $(t-1)/2$.
    ///
    /// Field implementations may override this to use an efficient addition chain.
    fn pow_by_t_minus1_over2(&self) -> Self;

    /// Returns the representation-derived key for the square-root hash table.
    ///
    /// The key need not have meaning outside the table lookup. It must be
    /// deterministic for normalized field elements. Table construction
    /// verifies that the resulting hash is injective on the complete
    /// 256-element lookup domain before constructing the inverse table.
    fn sqrt_hash_key(&self) -> u32;
}

/// Parameters for a perfect hash function used in square root computation.
#[cfg(feature = "sqrt-table")]
#[cfg_attr(docsrs, doc(cfg(feature = "sqrt-table")))]
#[derive(Debug)]
struct SqrtHasher<F: SqrtTableHelpers> {
    hash_mul: u32,
    marker: PhantomData<F>,
}

#[cfg(feature = "sqrt-table")]
impl<F: SqrtTableHelpers> SqrtHasher<F> {
    /// Returns a perfect (except for zero) hash of `x` for use with [`SqrtTables::inv`].
    ///
    /// The hash must be injective on the order-256 subgroup generated by
    /// `ROOT_OF_UNITY^(1 << (3 * SQRT_TABLE_DIGIT_BITS))`. Every nonzero
    /// input hashed by [`SqrtTables::sqrt_common`] is in that subgroup.
    ///
    /// Zero can also reach the lookup when `sqrt_alt` receives zero, or when
    /// either argument to `sqrt_ratio` is zero. Its hash need not be distinct
    /// from the subgroup's hashes: any inverse-table entry produces a bounded
    /// table index, and the result is zero because [`SqrtTables::sqrt_common`]
    /// multiplies it by `uv == 0`.
    fn hash(&self, x: &F) -> usize {
        // The field-specific multiplier was generated to make the high bits of
        // the wrapping product unique within the order-256 subgroup.
        (x.sqrt_hash_key().wrapping_mul(self.hash_mul) >> SQRT_HASH_SHIFT) as usize
    }
}

/// Tables used for square root computation.
#[cfg(feature = "sqrt-table")]
#[cfg_attr(docsrs, doc(cfg(feature = "sqrt-table")))]
#[derive(Debug)]
pub(crate) struct SqrtTables<F: SqrtTableHelpers> {
    hasher: SqrtHasher<F>,
    inv: Vec<u8>,
    g0: Box<[F; SQRT_TABLE_RADIX]>,
    g1: Box<[F; SQRT_TABLE_RADIX]>,
    g2: Box<[F; SQRT_TABLE_RADIX]>,
    g3: Box<[F; SQRT_TABLE_TOP_LEN]>,
}

#[cfg(feature = "sqrt-table")]
impl<F: SqrtTableHelpers> SqrtTables<F> {
    /// Builds tables given a multiplier for the perfect hash.
    ///
    /// The multiplier must make the hash injective on every power of
    /// `ROOT_OF_UNITY^(1 << (3 * SQRT_TABLE_DIGIT_BITS))`. Construction checks
    /// the complete domain and panics on a collision, before a table with
    /// incorrect inverse entries can be used.
    pub fn new(hash_mul: u32) -> Self {
        use alloc::vec;

        let hasher = SqrtHasher {
            hash_mul,
            marker: PhantomData,
        };

        let mut gtab = (0..4).scan(F::ROOT_OF_UNITY, |gi, _| {
            // gi == ROOT_OF_UNITY^(SQRT_TABLE_RADIX^i)
            let gtab_i: Vec<F> = (0..SQRT_TABLE_RADIX)
                .scan(F::ONE, |acc, _| {
                    let res = *acc;
                    *acc *= *gi;
                    Some(res)
                })
                .collect();
            *gi = gtab_i[SQRT_TABLE_RADIX - 1] * *gi;
            Some(gtab_i)
        });
        let gtab_0 = gtab.next().unwrap();
        let gtab_1 = gtab.next().unwrap();
        let gtab_2 = gtab.next().unwrap();
        let mut gtab_3 = gtab.next().unwrap();
        assert_eq!(gtab.next(), None);

        // Now invert gtab[3].
        let mut inv: Vec<u8> = vec![1; SQRT_HASH_TABLE_LEN];
        let mut occupied = vec![false; SQRT_HASH_TABLE_LEN];
        for (j, gtab_3_j) in gtab_3.iter().enumerate() {
            let hash = hasher.hash(gtab_3_j);
            assert!(!occupied[hash], "square-root hash collision at {j}");
            occupied[hash] = true;
            inv[hash] = ((SQRT_TABLE_RADIX - j) & (SQRT_TABLE_RADIX - 1)) as u8;
        }

        gtab_3.truncate(SQRT_TABLE_TOP_LEN);

        SqrtTables::<F> {
            hasher,
            inv,
            g0: gtab_0.into_boxed_slice().try_into().unwrap(),
            g1: gtab_1.into_boxed_slice().try_into().unwrap(),
            g2: gtab_2.into_boxed_slice().try_into().unwrap(),
            g3: gtab_3.into_boxed_slice().try_into().unwrap(),
        }
    }

    /// Exhaustively checks the safety invariant for representation-based
    /// hashing and verifies every inverse-table entry.
    ///
    /// The field-specific `fp_sqrt_table_hash_is_perfect` and
    /// `fq_sqrt_table_hash_is_perfect` tests call this for the production
    /// tables.
    #[cfg(test)]
    pub(crate) fn assert_hash_table_is_consistent(&self) {
        // bls12_381: `vec!` is not in scope here (no_std crate without `macro_use`).
        use alloc::vec;

        let g3 = (0..3 * SQRT_TABLE_DIGIT_BITS).fold(F::ROOT_OF_UNITY, |g, _| g.square());
        let mut value = F::ONE;
        let mut occupied = vec![false; self.inv.len()];

        for j in 0..SQRT_TABLE_RADIX {
            let hash = self.hasher.hash(&value);
            assert!(!occupied[hash], "square-root hash collision at {j}");
            occupied[hash] = true;
            assert_eq!(
                self.inv[hash],
                ((SQRT_TABLE_RADIX - j) & (SQRT_TABLE_RADIX - 1)) as u8,
                "incorrect square-root inverse entry at {j}",
            );
            value *= g3;
        }

        assert_eq!(value, F::ONE);
    }

    /// Computes:
    ///
    /// * (true,  sqrt(num/div)),                 if num and div are nonzero and num/div is a square in the field;
    /// * (true,  0),                             if num is zero;
    /// * (false, 0),                             if num is nonzero and div is zero;
    /// * (false, sqrt(ROOT_OF_UNITY * num/div)), if num and div are nonzero and num/div is a nonsquare in the field;
    ///
    /// where ROOT_OF_UNITY is a generator of the order 2^n subgroup (and therefore a nonsquare).
    ///
    /// The choice of root from sqrt is unspecified.
    pub fn sqrt_ratio(&self, num: &F, div: &F) -> (Choice, F) {
        // Based on:
        // * [Sarkar2020](https://eprint.iacr.org/2020/1407)
        // * [BDLSY2012](https://cr.yp.to/papers.html#ed25519)
        //
        // We need to calculate uv and v, where v = u^((T-1)/2), u = num/div, and p-1 = T * 2^S.
        // We can rewrite as follows:
        //
        //      v = (num/div)^((T-1)/2)
        //        = num^((T-1)/2) * div^(p-1 - (T-1)/2)    [Fermat's Little Theorem]
        //        =       "       * div^(T * 2^S - (T-1)/2)
        //        =       "       * div^((2^(S+1) - 1)*(T-1)/2 + 2^S)
        //        = (num * div^(2^(S+1) - 1))^((T-1)/2) * div^(2^S)
        //
        // Let  w = (num * div^(2^(S+1) - 1))^((T-1)/2) * div^(2^S - 1).
        // Then v = w * div, and uv = num * v / div = num * w.
        //
        // We calculate:
        //
        //      s = div^(2^S - 1) using an addition chain
        //      t = div^(2^(S+1) - 1) = s^2 * div
        //      w = (num * t)^((T-1)/2) * s using another addition chain
        //
        // then u and uv as above. The addition chains are given in
        // https://github.com/zcash/pasta/blob/master/addchain_sqrt.py .
        // The overall cost of this part is similar to a single full-width exponentiation,
        // regardless of S.

        let sqr = |x: F, i: u32| (0..i).fold(x, |x, _| x.square());

        // s = div^(2^S - 1)
        let s = (0..5).fold(*div, |d: F, i| sqr(d, 1 << i) * d);

        // t == div^(2^(S+1) - 1)
        let t = s.square() * div;

        // w = (num * t)^((T-1)/2) * s
        let w = (t * num).pow_by_t_minus1_over2() * s;

        // v == u^((T-1)/2)
        let v = w * div;

        // uv = u * v
        let uv = w * num;

        let res = self.sqrt_common(&uv, &v);

        let sqdiv = res.square() * div;
        let is_square = (sqdiv - num).is_zero();
        let is_nonsquare = (sqdiv - F::ROOT_OF_UNITY * num).is_zero();
        assert!(bool::from(
            num.is_zero() | div.is_zero() | (is_square ^ is_nonsquare)
        ));

        (is_square, res)
    }

    /// Same as sqrt_ratio(u, one()) but more efficient.
    pub fn sqrt_alt(&self, u: &F) -> (Choice, F) {
        let v = u.pow_by_t_minus1_over2();
        let uv = *u * v;

        let res = self.sqrt_common(&uv, &v);

        let sq = res.square();
        let is_square = (sq - u).is_zero();
        let is_nonsquare = (sq - F::ROOT_OF_UNITY * u).is_zero();
        assert!(bool::from(u.is_zero() | (is_square ^ is_nonsquare)));

        (is_square, res)
    }

    /// Common part of sqrt_ratio and sqrt_alt: return their result given v = u^((T-1)/2) and uv = u * v.
    fn sqrt_common(&self, uv: &F, v: &F) -> F {
        let sqr = |x: F, i: u32| (0..i).fold(x, |x, _| x.square());
        let inv = |x: F| self.inv[self.hasher.hash(&x)] as usize;

        let x3 = *uv * v;
        let x2 = sqr(x3, 8);
        let x1 = sqr(x2, 8);
        let x0 = sqr(x1, 8);

        // i = 0, 1
        let mut t_ = inv(x0); // = t >> 16
        // 1 == x0 * ROOT_OF_UNITY^(t_ << 24)
        assert!(t_ < 0x100);
        let alpha = x1 * self.g2[t_];

        // i = 2
        t_ += inv(alpha) << 8; // = t >> 8
        // 1 == x1 * ROOT_OF_UNITY^(t_ << 16)
        assert!(t_ < 0x10000);
        let alpha = x2 * self.g1[t_ & 0xFF] * self.g2[t_ >> 8];

        // i = 3
        t_ += inv(alpha) << 16; // = t
        // 1 == x2 * ROOT_OF_UNITY^(t_ << 8)
        assert!(t_ < 0x1000000);
        let alpha = x3 * self.g0[t_ & 0xFF] * self.g1[(t_ >> 8) & 0xFF] * self.g2[t_ >> 16];

        t_ += inv(alpha) << 24; // = t << 1
        // 1 == x3 * ROOT_OF_UNITY^t_
        t_ = (((t_ as u64) + 1) >> 1) as usize;
        assert!(t_ <= 0x80000000);

        *uv * self.g0[t_ & 0xFF]
            * self.g1[(t_ >> 8) & 0xFF]
            * self.g2[(t_ >> 16) & 0xFF]
            * self.g3[t_ >> 24]
    }
}