p3_baby_bear/

baby_bear.rs

use p3_field::exponentiation::exp_1725656503;
use p3_field::{Algebra, PrimeCharacteristicRing, UniformSamplingField};
use p3_monty_31::{
    BarrettParameters, BinomialExtensionData, FieldParameters, MontyField31, MontyParameters,
    PackedMontyParameters, RelativelyPrimePower, TwoAdicData,
};

/// The prime field `2^31 - 2^27 + 1`, a.k.a. the Baby Bear field.
pub type BabyBear = MontyField31<BabyBearParameters>;

#[derive(Copy, Clone, Default, Debug, Eq, Hash, PartialEq)]
pub struct BabyBearParameters;

impl MontyParameters for BabyBearParameters {
    /// The Baby Bear prime: 2^31 - 2^27 + 1.
    /// This is the unique 31-bit prime with the highest possible 2 adicity (27).
    const PRIME: u32 = 0x78000001;

    const MONTY_BITS: u32 = 32;
    const MONTY_MU: u32 = 0x88000001;
}

impl PackedMontyParameters for BabyBearParameters {}

impl BarrettParameters for BabyBearParameters {}

impl FieldParameters for BabyBearParameters {
    const MONTY_GEN: BabyBear = BabyBear::new(31);

    // Measured -3.5% on quotient evaluation (twoadic PCS, Poseidon2 AIR, aarch64+neon).
    const BENEFITS_FROM_LOCKSTEP_EVALUATION: bool = true;
}

impl RelativelyPrimePower<7> for BabyBearParameters {
    /// In the field `BabyBear`, `a^{1/7}` is equal to a^{1725656503}.
    ///
    /// This follows from the calculation `7 * 1725656503 = 6*(2^31 - 2^27) + 1 = 1 mod (p - 1)`.
    fn exp_root_d<R: PrimeCharacteristicRing>(val: R) -> R {
        exp_1725656503(val)
    }
}

impl TwoAdicData for BabyBearParameters {
    const TWO_ADICITY: usize = 27;

    type ArrayLike = &'static [BabyBear];

    const TWO_ADIC_GENERATORS: Self::ArrayLike = &BabyBear::new_array([
        0x1, 0x78000000, 0x67055c21, 0x5ee99486, 0xbb4c4e4, 0x2d4cc4da, 0x669d6090, 0x17b56c64,
        0x67456167, 0x688442f9, 0x145e952d, 0x4fe61226, 0x4c734715, 0x11c33e2a, 0x62c3d2b1,
        0x77cad399, 0x54c131f4, 0x4cabd6a6, 0x5cf5713f, 0x3e9430e8, 0xba067a3, 0x18adc27d,
        0x21fd55bc, 0x4b859b3d, 0x3bd57996, 0x4483d85a, 0x3a26eef8, 0x1a427a41,
    ]);

    const ROOTS_8: Self::ArrayLike = &BabyBear::new_array([0x1, 0x5ee99486, 0x67055c21, 0xc9ea3ba]);
    const INV_ROOTS_8: Self::ArrayLike =
        &BabyBear::new_array([0x1, 0x6b615c47, 0x10faa3e0, 0x19166b7b]);

    const ROOTS_16: Self::ArrayLike = &BabyBear::new_array([
        0x1, 0xbb4c4e4, 0x5ee99486, 0x4b49e08, 0x67055c21, 0x5376917a, 0xc9ea3ba, 0x563112a7,
    ]);
    const INV_ROOTS_16: Self::ArrayLike = &BabyBear::new_array([
        0x1, 0x21ceed5a, 0x6b615c47, 0x24896e87, 0x10faa3e0, 0x734b61f9, 0x19166b7b, 0x6c4b3b1d,
    ]);
}

impl BinomialExtensionData<4> for BabyBearParameters {
    const W: BabyBear = BabyBear::new(11);
    const DTH_ROOT: BabyBear = BabyBear::new(1728404513);
    const EXT_GENERATOR: [BabyBear; 4] = BabyBear::new_array([8, 1, 0, 0]);
    const EXT_TWO_ADICITY: usize = 29;

    type ArrayLike = [[BabyBear; 4]; 2];
    const TWO_ADIC_EXTENSION_GENERATORS: Self::ArrayLike =
        BabyBear::new_2d_array([[0, 0, 1996171314, 0], [0, 0, 0, 124907976]]);
}

impl BinomialExtensionData<5> for BabyBearParameters {
    const W: BabyBear = BabyBear::new(2);

    #[inline(always)]
    fn mul_w<A: Algebra<MontyField31<Self>>>(a: A) -> A {
        a.double()
    }

    const DTH_ROOT: BabyBear = BabyBear::new(815036133);
    const EXT_GENERATOR: [BabyBear; 5] = BabyBear::new_array([8, 1, 0, 0, 0]);
    const EXT_TWO_ADICITY: usize = 27;

    type ArrayLike = [[BabyBear; 5]; 0];
    const TWO_ADIC_EXTENSION_GENERATORS: Self::ArrayLike = [];
}

impl BinomialExtensionData<8> for BabyBearParameters {
    const W: BabyBear = BabyBear::new(11);
    const DTH_ROOT: BabyBear = BabyBear::new(420899707);
    const EXT_GENERATOR: [BabyBear; 8] = BabyBear::new_array([5, 1, 0, 0, 0, 0, 0, 0]);
    const EXT_TWO_ADICITY: usize = 30;

    type ArrayLike = [[BabyBear; 8]; 3];
    const TWO_ADIC_EXTENSION_GENERATORS: Self::ArrayLike = BabyBear::new_2d_array([
        [0, 0, 0, 0, 1996171314, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 124907976, 0],
        [0, 0, 0, 518392818, 0, 0, 0, 0],
    ]);
}

impl UniformSamplingField for BabyBearParameters {
    const MAX_SINGLE_SAMPLE_BITS: usize = 27;
    // NOTE: We only include `0` to not have to deal with one-off indexing. `k` must be > 0.
    // Also, we don't care about k > 30 for BabyBear.
    const SAMPLING_BITS_M: [u64; 64] = {
        let prime: u64 = Self::PRIME as u64;
        let mut a = [0u64; 64];
        let mut k = 0;
        while k < 64 {
            if k == 0 {
                a[k] = prime; // This value is irrelevant in practice. `bits = 0` returns 0 always
            } else {
                // Create a mask to zero out the last k bits
                let mask = !((1u64 << k) - 1);
                a[k] = prime & mask;
            }
            k += 1;
        }
        a
    };
}

#[cfg(test)]
mod tests {
    use core::array;

    use num_bigint::BigUint;
    use p3_field::extension::BinomialExtensionField;
    use p3_field::{InjectiveMonomial, PermutationMonomial, PrimeField64, TwoAdicField};
    use p3_field_testing::{
        test_field, test_field_dft, test_field_dft_consistency, test_field_dft_large,
        test_field_json_serialization, test_prime_field, test_prime_field_32, test_prime_field_64,
        test_two_adic_field,
    };
    use proptest::prelude::*;

    use super::*;

    type F = BabyBear;
    type EF = BinomialExtensionField<F, 4>;

    #[test]
    fn test_baby_bear_two_adicity_generators() {
        let base = BabyBear::from_u32(0x1a427a41);
        for bits in 0..=BabyBear::TWO_ADICITY {
            assert_eq!(
                BabyBear::two_adic_generator(bits),
                base.exp_power_of_2(BabyBear::TWO_ADICITY - bits)
            );
        }
    }

    #[test]
    fn test_to_babybear_array() {
        let range_array: [u32; 32] = array::from_fn(|i| i as u32);
        assert_eq!(
            BabyBear::new_array(range_array),
            range_array.map(F::from_u32)
        );
    }

    #[test]
    fn test_baby_bear() {
        let f = F::from_u32(100);
        assert_eq!(f.as_canonical_u64(), 100);

        let f_1 = F::ONE;
        let f_2 = F::TWO;
        let f_p_minus_1 = F::NEG_ONE;
        let f_p_minus_2 = F::NEG_ONE + F::NEG_ONE;
        let m1 = F::from_u32(0x34167c58);
        let m2 = F::from_u32(0x61f3207b);
        let expected_prod = F::from_u32(0x1b5c8046);
        assert_eq!(m1 * m2, expected_prod);

        assert_eq!(m1.injective_exp_n().injective_exp_root_n(), m1);
        assert_eq!(m2.injective_exp_n().injective_exp_root_n(), m2);
        assert_eq!(F::TWO.injective_exp_n().injective_exp_root_n(), F::TWO);

        test_field_json_serialization(&[f, f_1, f_2, f_p_minus_1, f_p_minus_2, m1, m2]);
    }

    proptest! {
        /// Exercises the inherent (two-adic) `try_sqrt` which shadows
        /// `Field::try_sqrt` for the concrete `BabyBear` type.
        #[test]
        fn test_baby_bear_inherent_sqrt(x in prop::num::u32::ANY.prop_map(F::from_u32)) {
            // The square of any element is a quadratic residue, so its root
            // always exists and squares back to the original square.
            let square = x.square();
            prop_assert_eq!(square.try_sqrt().map(|r| r.square()), Some(square));

            // Whenever a root is reported, it must square back to `x`.
            if let Some(r) = x.try_sqrt() {
                prop_assert_eq!(r.square(), x);
            }
        }
    }

    // MontyField31's have no redundant representations.
    const ZEROS: [BabyBear; 1] = [BabyBear::ZERO];
    const ONES: [BabyBear; 1] = [BabyBear::ONE];

    // Get the prime factorization of the order of the multiplicative group.
    // i.e. the prime factorization of P - 1.
    fn multiplicative_group_prime_factorization() -> [(BigUint, u32); 3] {
        [
            (BigUint::from(2u8), 27),
            (BigUint::from(3u8), 1),
            (BigUint::from(5u8), 1),
        ]
    }

    test_field!(
        crate::BabyBear,
        &super::ZEROS,
        &super::ONES,
        &super::multiplicative_group_prime_factorization()
    );
    test_two_adic_field!(crate::BabyBear);

    test_field_dft!(radix2dit, crate::BabyBear, super::EF, p3_dft::Radix2Dit<_>);
    test_field_dft!(
        radix2smallbatch,
        crate::BabyBear,
        super::EF,
        p3_dft::Radix2DFTSmallBatch<_>
    );
    test_field_dft!(bowers, crate::BabyBear, super::EF, p3_dft::Radix2Bowers);
    test_field_dft!(
        parallel,
        crate::BabyBear,
        super::EF,
        p3_dft::Radix2DitParallel::<_>
    );
    test_field_dft!(
        recur_dft,
        crate::BabyBear,
        super::EF,
        p3_monty_31::dft::RecursiveDft<_>
    );
    test_field_dft_consistency!(
        radix2_smallbatch_and_ditparallel,
        crate::BabyBear,
        super::EF,
        p3_dft::Radix2DFTSmallBatch<_>,
        p3_dft::Radix2DitParallel<_>
    );
    test_field_dft_large!(
        radix2smallbatch_large,
        crate::BabyBear,
        super::EF,
        p3_dft::Radix2DFTSmallBatch<_>
    );

    test_field_dft_large!(
        parallel_large,
        crate::BabyBear,
        super::EF,
        p3_dft::Radix2DitParallel<_>
    );

    test_field_dft_large!(
        recur_dft_large,
        crate::BabyBear,
        super::EF,
        p3_monty_31::dft::RecursiveDft<_>
    );
    test_prime_field!(crate::BabyBear);
    test_prime_field_64!(crate::BabyBear, &super::ZEROS, &super::ONES);
    test_prime_field_32!(crate::BabyBear, &super::ZEROS, &super::ONES);
}