cryptix-bn254 0.1.0

use core::ops::{Add, Sub, Mul, Neg};

use cryptix_bigint::property::IsBigInt;
use cryptix_field::group::*;
use cryptix_field::ring::*;
use cryptix_field::field::*;
use cryptix_field::field::montgomery::MontgomeryOps;

use super::{U256, BN254, FpElement};

use super::fp4::Fp4Element;

/// Element in field F_{p^12}
/// 
/// ```text,ignore
/// Fp12[t] = Fp4[t] / (t^3 - s) where s is element 0 + s in Fp4
/// ```
#[derive(PartialEq, Eq, Clone, Copy)]
pub struct Fp12Element(pub Fp4Element, pub Fp4Element, pub Fp4Element);

impl core::fmt::Debug for Fp12Element {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "fp12!(\n    \"")?;
        self.0.0.0.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.0.0.1.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.0.1.0.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.0.1.1.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.1.0.0.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.1.0.1.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.1.1.0.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.1.1.1.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.2.0.0.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.2.0.1.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.2.1.0.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.2.1.1.fmt(f)?;
        write!(f, "\"\n)\n    ")
    }
}

impl AbelianGroup for Fp12Element { }

impl Group for Fp12Element { }

impl Add for Fp12Element {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(
            self.0 + rhs.0, 
            self.1 + rhs.1,
            self.2 + rhs.2
        )
    }
}

impl Sub for Fp12Element{
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self(
            self.0 - rhs.0, 
            self.1 - rhs.1,
            self.2 - rhs.2
        )
    }
}

impl AddIdentity for Fp12Element{
    const ADD_IDENTITY: Self = Self(Fp4Element::ZERO, Fp4Element::ZERO, Fp4Element::ZERO);
}

impl Neg for Fp12Element { 
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self(-self.0, -self.1, -self.2)
    }
}

/// # Safety
/// 
/// Element is backed by biguint, which is associative under addition
impl AssociativeAdd for Fp12Element { }

/// # Safety
/// 
/// Element is backed by biguint, which is communicative under addition
impl CommunicativeAdd for Fp12Element { }


impl Ring for Fp12Element { }

impl Mul for Fp12Element {
    type Output = Self;

    /// calculate a * b mod m
    /// this can be achieved more efficiently with montgomery multiplication
    fn mul(self, rhs: Self) -> Self::Output {
        self.mont_mul(rhs).mont_form()
    }
}

/// # Safety
/// 
/// our element type is backed by biguint, so mod mul is associative
impl AssociativeMul for Fp12Element { }

/// # Safety
/// 
/// our element type is backed by biguint, so mod mul is distributive over add
impl DistributiveMul for Fp12Element { }


/// # Safety
/// 
/// 1 is the multiplicative ideneity for biguint
impl MulIdentity for Fp12Element {
    const MUL_IDENTITY: Self = Self(Fp4Element::ONE, Fp4Element::ZERO, Fp4Element::ZERO);
}

/// # Safety
/// 
/// BigUInt mod mul is communicative
impl CommunicativeMul for Fp12Element { }

/// The montgomery trait bound restricts the modular to odd prime
impl MulInverse for Fp12Element {
    fn mul_inv(self) -> Self {
        self.mont_form().mont_inv().mont_rdc()
    }
}


impl Field for Fp12Element {
    fn hlv(self) -> Self {
        Self(self.0.hlv(), self.1.hlv(), self.2.hlv())
    }

    fn is_zero(&self) -> bool {
        self.0.is_zero() && self.1.is_zero() && self.2.is_zero()
    }
}

impl From<FpElement> for Fp12Element {
    fn from(value: FpElement) -> Self {
        Self(Fp4Element::from(value), Fp4Element::ZERO, Fp4Element::ZERO)
    }
}

impl MontgomeryOps<U256, BN254> for Fp12Element {
    /// input: fp2_add
    ///   (a0 + b0t + c0t^2)(a1 + b1t + c1t^2) mod (t^3 - s)
    /// = a0a1 + (a0b1 + b0a1)t + (a0c1 + b0b1 + c0a1)t^2 + (b0c1 + c0b1)t^3 + c0c1t^4 mod (t^3 - s)
    ///                     - (t^3 - s)(b0c1 + c0b1)
    /// = a0a1 + (b0c1 + c0b1)s + (a0b1 + b0a1)t + (a0c1 + b0b1 + c0a1)t^2 + c0c1t^4 mod (t^3 - s)
    ///                     - (t^3 - s)(c0c1t)
    /// = a0a1 + (b0c1 + c0b1)s + (a0b1 + b0a1 + c0c1s)t + (a0c1 + b0b1 + c0a1)t^2
    /// ```
    fn mont_mul(self, rhs: Self) -> Self {
        let d0 = self.0.mont_mul(rhs.0);
        let d1 = self.1.mont_mul(rhs.1);
        let d2 = self.2.mont_mul(rhs.2);

        let d3 = (rhs.0 - rhs.1).mont_mul(self.0 - self.1);
        let d4 = (rhs.0 - rhs.2).mont_mul(self.0 - self.2);
        let d5 = (rhs.1 - rhs.2).mont_mul(self.1 - self.2);

        let c0 = d0;
        let t = d1 + d2;
        let c2 = d0 + t - d4;
        let t = (t - d5).mul_s();
        let c0 = c0 + t;
        let c1 = d0 + d1 - d3 + d2.mul_s();

        Self(c0, c1, c2)
    }

    fn mont_mul_fp(self, rhs: FpElement) -> Self {
        Self(
            self.0.mont_mul_fp(rhs), 
            self.1.mont_mul_fp(rhs), 
            self.2.mont_mul_fp(rhs)
        )
    }

    fn mont_inv(self) -> Self {
        let Self(a, b, c) = self;
        let (c0, c1, c2) = (
            a.mont_sqr() - b.mont_mul(c).mul_s(),       // a^2 - bcs
            c.mont_sqr().mul_s() - a.mont_mul(b),       // c^2s - ab
            b.mont_sqr() - a.mont_mul(c)                // b^2 - ac
        );

        let t = (
            // (bc2s + ac0 + cc1s)^-1
            b.mont_mul(c2).mul_s() + a.mont_mul(c0) + c.mont_mul(c1).mul_s()
        ).mont_inv();
        
        Self(t.mont_mul(c0), t.mont_mul(c1), t.mont_mul(c2))
    }

    fn mont_rdc(self) -> Self {
        Self(self.0.mont_rdc(), self.1.mont_rdc(), self.2.mont_rdc())
    }
}

impl Fp12Element  {
    pub fn sparse_mul(self, rhs: Self) -> Self {
        let t0 = rhs.1.mont_mul(self.0);
        let t1 = rhs.2.sparse_mul(self.1.0);
        let u0 = (rhs.0 + rhs.2).sparse_mul(self.1.0);
        let u1 = rhs.1 + rhs.2;
        let u2 = Fp4Element(self.0.0 + self.1.0, self.0.1);
        let u1 = u1.mont_mul(u2);
        let u2 = (rhs.0 + rhs.1).mont_mul(self.0);

        Self(t1.mul_s() + u2 - t0, u0 + t0 - t1, u1 - t0 - t1)
    }
    
    /// TODO
    pub fn map_frob(self) -> Self {
        Self(
            self.0.map_frob(),
            Fp4Element(
                self.1.0.map_frob().mont_mul(BN254::XI_MONT[0]),
                self.1.1.map_frob().mont_mul(BN254::XI_MONT[3])
            ),
            Fp4Element(
                self.2.0.map_frob().mont_mul(BN254::XI_MONT[1]), 
                self.2.1.map_frob().mont_mul(BN254::XI_MONT[4])
            )
        )
    } 

    /// TODO
    pub fn map2_frob(self) -> Self {
        Self(
            Fp4Element(self.0.0, self.0.1.mont_mul_fp(BN254::P2_FROB_MAP_VAL[2])), 
            Fp4Element(
                self.1.0.mont_mul_fp(BN254::P2_FROB_MAP_VAL[0]), 
                self.1.1.mont_mul_fp(BN254::P2_FROB_MAP_VAL[3])
            ), 
            Fp4Element(
                self.2.0.mont_mul_fp(BN254::P2_FROB_MAP_VAL[1]), 
                self.2.1.mont_mul_fp(BN254::P2_FROB_MAP_VAL[4])
            ),
        )
    }
}

#[cfg(feature = "rand")]
impl Fp12Element {
    pub fn rand(rng: &mut impl rand_core::CryptoRngCore) -> Self {
        Self(Fp4Element::rand(rng), Fp4Element::rand(rng), Fp4Element::rand(rng))
    }
}

impl From<Fp12Element> for [u8; U256::BYTE_LEN * 12] {
    fn from(val: Fp12Element) -> Self {
        let mut buf = [0_u8; U256::BYTE_LEN * 12];
        let tmp: [u8; U256::BYTE_LEN * 4] = val.0.into();
        buf[..U256::BYTE_LEN * 4].copy_from_slice(&tmp);
        let tmp: [u8; U256::BYTE_LEN * 4] = val.1.into();
        buf[U256::BYTE_LEN * 4..U256::BYTE_LEN * 8].copy_from_slice(&tmp);
        let tmp: [u8; U256::BYTE_LEN * 4] = val.1.into();
        buf[U256::BYTE_LEN * 8..U256::BYTE_LEN * 12].copy_from_slice(&tmp);
        buf
    }
}