p3-goldilocks 0.6.1

use p3_field::extension::{
    Binomial, BinomiallyExtendable, CubicTrinomial, CubicTrinomialExtendable, ExtensionAlgebra,
    HasTwoAdicBinomialExtension, HasTwoAdicCubicExtension, binomial_mul, cubic_square,
    trinomial_cubic_mul,
};
use p3_field::{PrimeCharacteristicRing, TwoAdicField, field_to_array};

use crate::Goldilocks;

impl ExtensionAlgebra<Self, 2, Binomial<Self>> for Goldilocks {
    #[inline]
    fn ext_mul(a: &[Self; 2], b: &[Self; 2], res: &mut [Self; 2]) {
        binomial_mul::<Self, Self, Self, 2>(a, b, res, <Self as BinomiallyExtendable<2>>::W);
    }
}

impl BinomiallyExtendable<2> for Goldilocks {
    // Verifiable in Sage with
    // `R.<x> = GF(p)[]; assert (x^2 - 7).is_irreducible()`.
    const W: Self = Self::new(7);

    // DTH_ROOT = W^((p - 1)/2).
    const DTH_ROOT: Self = Self::new(18446744069414584320);

    const EXT_GENERATOR: [Self; 2] = [
        Self::new(18081566051660590251),
        Self::new(16121475356294670766),
    ];
}

impl HasTwoAdicBinomialExtension<2> for Goldilocks {
    const EXT_TWO_ADICITY: usize = 33;

    fn ext_two_adic_generator(bits: usize) -> [Self; 2] {
        assert!(bits <= 33);

        if bits == 33 {
            [Self::ZERO, Self::new(15659105665374529263)]
        } else {
            [Self::two_adic_generator(bits), Self::ZERO]
        }
    }
}

impl ExtensionAlgebra<Self, 3, CubicTrinomial> for Goldilocks {
    #[inline]
    fn ext_mul(a: &[Self; 3], b: &[Self; 3], res: &mut [Self; 3]) {
        trinomial_cubic_mul::<Self>(a, b, res);
    }

    #[inline]
    fn ext_square(a: &[Self; 3], res: &mut [Self; 3]) {
        cubic_square::<Self>(a, res);
    }
}

impl CubicTrinomialExtendable for Goldilocks {
    // Verifiable via:
    // ```sage
    // p = 2**64 - 2**32 + 1
    // R.<x> = GF(p)[]
    // assert (x^3 - x - 1).is_irreducible()
    // ```
    const FROBENIUS_MATRIX: [[Self; 3]; 3] = [
        [
            Self::ONE,
            Self::new(10615703402128488253),
            Self::new(6700183068485440220),
        ],
        [
            Self::ZERO,
            Self::new(10050274602728160328),
            Self::new(14531223735771536287),
        ],
        [
            Self::ZERO,
            Self::new(11746561000929144102),
            Self::new(8396469466686423992),
        ],
    ];

    // Verifiable via:
    // ```sage
    // p = 2**64 - 2**32 + 1
    // F = GF(p)
    // R.<x> = F[]
    // K.<a> = F.extension(x^3 - x - 1)
    // g = 2 + a
    // order = p^3 - 1
    // assert g.multiplicative_order() == order
    // ```
    const EXT_GENERATOR: [Self; 3] = [Self::TWO, Self::ONE, Self::ZERO];
}

impl HasTwoAdicCubicExtension for Goldilocks {
    const EXT_TWO_ADICITY: usize = 32;

    fn ext_two_adic_generator(bits: usize) -> [Self; 3] {
        assert!(bits <= 32);

        field_to_array(Self::two_adic_generator(bits))
    }
}

impl ExtensionAlgebra<Self, 5, Binomial<Self>> for Goldilocks {
    #[inline]
    fn ext_mul(a: &[Self; 5], b: &[Self; 5], res: &mut [Self; 5]) {
        binomial_mul::<Self, Self, Self, 5>(a, b, res, <Self as BinomiallyExtendable<5>>::W);
    }
}

impl BinomiallyExtendable<5> for Goldilocks {
    // Verifiable via:
    //  ```sage
    //  # Define Fp
    //  p = 2**64 - 2**32 + 1
    //  F = GF(p)

    //  # Define Fp[z]
    //  R.<z> = PolynomialRing(F)

    //  # The polynomial x^5-3 is irreducible
    //  assert(R(z^5-3).is_irreducible())
    //  ```
    const W: Self = Self::new(3);

    // 5-th root = w^((p - 1)/5)
    const DTH_ROOT: Self = Self::new(1041288259238279555);

    // Generator of the extension field
    // Obtained by finding the smallest Hamming weight vector
    // with appropriate order, starting at [0,1,0,0,0]
    const EXT_GENERATOR: [Self; 5] = [Self::TWO, Self::ONE, Self::ZERO, Self::ZERO, Self::ZERO];
}

impl HasTwoAdicBinomialExtension<5> for Goldilocks {
    const EXT_TWO_ADICITY: usize = 32;

    fn ext_two_adic_generator(bits: usize) -> [Self; 5] {
        assert!(bits <= 32);

        field_to_array(Self::two_adic_generator(bits))
    }
}

#[cfg(test)]
mod test_quadratic_extension {

    use num_bigint::BigUint;
    use p3_field::extension::BinomialExtensionField;
    use p3_field::{ExtensionField, PrimeCharacteristicRing};
    use p3_field_testing::{
        test_extension_field, test_field, test_packed_extension_field,
        test_two_adic_extension_field,
    };

    use crate::Goldilocks;

    type F = Goldilocks;
    type EF = BinomialExtensionField<F, 2>;

    // There is a redundant representation of zero but we already tested it
    // when testing the base field.
    const ZEROS: [EF; 1] = [EF::ZERO];
    const ONES: [EF; 1] = [EF::ONE];

    // Get the prime factorization of the order of the multiplicative group.
    // i.e. the prime factorization of P^2 - 1.
    fn multiplicative_group_prime_factorization() -> [(BigUint, u32); 9] {
        [
            (BigUint::from(2u8), 33),
            (BigUint::from(3u8), 1),
            (BigUint::from(5u8), 1),
            (BigUint::from(7u8), 1),
            (BigUint::from(17u8), 1),
            (BigUint::from(179u8), 1),
            (BigUint::from(257u16), 1),
            (BigUint::from(65537u32), 1),
            (BigUint::from(7361031152998637u64), 1),
        ]
    }

    test_field!(
        super::EF,
        &super::ZEROS,
        &super::ONES,
        &super::multiplicative_group_prime_factorization()
    );

    test_extension_field!(super::F, super::EF);
    test_two_adic_extension_field!(super::F, super::EF);

    type Pef = <EF as ExtensionField<F>>::ExtensionPacking;
    const PACKED_ZEROS: [Pef; 1] = [Pef::ZERO];
    const PACKED_ONES: [Pef; 1] = [Pef::ONE];
    test_packed_extension_field!(
        super::F,
        super::EF,
        super::Pef,
        &super::PACKED_ZEROS,
        &super::PACKED_ONES
    );
    p3_field_testing::test_packed_binomial_extension_division!(F, 2);
}

#[cfg(test)]
mod test_cubic_trinomial_extension {

    use num_bigint::BigUint;
    use p3_field::extension::CubicTrinomialExtensionField;
    use p3_field::{ExtensionField, PrimeCharacteristicRing};
    use p3_field_testing::{
        test_extension_field, test_field, test_frobenius, test_packed_extension_field,
        test_two_adic_extension_field,
    };

    use crate::Goldilocks;

    type F = Goldilocks;
    type EF = CubicTrinomialExtensionField<F>;

    const ZEROS: [EF; 1] = [EF::ZERO];
    const ONES: [EF; 1] = [EF::ONE];

    fn multiplicative_group_prime_factorization() -> [(BigUint, u32); 9] {
        [
            (BigUint::from(2u8), 32),
            (BigUint::from(3u8), 2),
            (BigUint::from(5u8), 1),
            (BigUint::from(17u8), 1),
            (BigUint::from(257u16), 1),
            (BigUint::from(937u16), 1),
            (BigUint::from(65537u32), 1),
            (BigUint::from(724723u32), 1),
            (BigUint::from(167034643597991036904547663171u128), 1),
        ]
    }

    // TODO: Consider generalizing and putting into test_extension_field!
    #[test]
    fn test_defining_relation() {
        let x = EF::new([F::ZERO, F::ONE, F::ZERO]);
        let x_cubed = x * x * x;
        let x_plus_one = x + EF::ONE;
        assert_eq!(x_cubed, x_plus_one, "X^3 should equal X + 1");
    }

    test_field!(
        super::EF,
        &super::ZEROS,
        &super::ONES,
        &super::multiplicative_group_prime_factorization()
    );

    test_extension_field!(super::F, super::EF);
    test_two_adic_extension_field!(super::F, super::EF);
    test_frobenius!(super::F, super::EF);

    type Pef = <EF as ExtensionField<F>>::ExtensionPacking;
    const PACKED_ZEROS: [Pef; 1] = [Pef::ZERO];
    const PACKED_ONES: [Pef; 1] = [Pef::ONE];
    test_packed_extension_field!(
        super::F,
        super::EF,
        super::Pef,
        &super::PACKED_ZEROS,
        &super::PACKED_ONES
    );
}

#[cfg(test)]
mod test_quintic_extension {

    use num_bigint::BigUint;
    use p3_field::extension::BinomialExtensionField;
    use p3_field::{ExtensionField, PrimeCharacteristicRing};
    use p3_field_testing::{
        test_extension_field, test_field, test_packed_extension_field,
        test_two_adic_extension_field,
    };

    use crate::Goldilocks;

    type F = Goldilocks;
    type EF = BinomialExtensionField<F, 5>;

    // There is a redundant representation of zero but we already tested it
    // when testing the base field.
    const ZEROS: [EF; 1] = [EF::ZERO];
    const ONES: [EF; 1] = [EF::ONE];

    // Get the prime factorization of the order of the multiplicative group.
    // i.e. the prime factorization of P^5 - 1.
    fn multiplicative_group_prime_factorization() -> [(num_bigint::BigUint, u32); 10] {
        [
            (BigUint::from(2u8), 32),
            (BigUint::from(3u8), 1),
            (BigUint::from(5u8), 2),
            (BigUint::from(17u8), 1),
            (BigUint::from(257u16), 1),
            (BigUint::from(45971u16), 1),
            (BigUint::from(65537u32), 1),
            (BigUint::from(255006435240067831u64), 1),
            (BigUint::from(280083648770327405561u128), 1),
            (BigUint::from(7053197395277272939628824863222181u128), 1),
        ]
    }

    test_field!(
        super::EF,
        &super::ZEROS,
        &super::ONES,
        &super::multiplicative_group_prime_factorization()
    );

    test_extension_field!(super::F, super::EF);
    test_two_adic_extension_field!(super::F, super::EF);

    type Pef = <EF as ExtensionField<F>>::ExtensionPacking;
    const PACKED_ZEROS: [Pef; 1] = [Pef::ZERO];
    const PACKED_ONES: [Pef; 1] = [Pef::ONE];
    test_packed_extension_field!(
        super::F,
        super::EF,
        super::Pef,
        &super::PACKED_ZEROS,
        &super::PACKED_ONES
    );
    p3_field_testing::test_packed_binomial_extension_division!(F, 5);
}