cryptix-bn254 0.1.0

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

use cryptix_bigint::property::IsBigInt;
use cryptix_field::field::montgomery::Montgomery;
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::fp2::Fp2Element;

/// Element in field F_{p^4}
/// 
/// ```text,ignore
/// Fp4[s] = Fp2[s] / (s^2 - ξ) where ξ = 1 + i
/// ```
#[derive(PartialEq, Eq, Clone, Copy)]
pub struct Fp4Element(pub Fp2Element, pub Fp2Element);

impl core::fmt::Debug for Fp4Element {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "fp4!(\n    \"")?;
        self.0.0.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.0.1.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.1.0.fmt(f)?;
        write!(f, "\", \n    \"")?;
        self.1.1.fmt(f)?;
        write!(f, "\"\n)\n    ")
    }
}

impl AbelianGroup for Fp4Element { }

impl Group for Fp4Element { }

impl Add for Fp4Element {
    type Output = Self;

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

impl Sub for Fp4Element{
    type Output = Self;

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

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

impl Neg for Fp4Element { 
    type Output = Self;

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

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

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


impl Ring for Fp4Element { }

impl Mul for Fp4Element {
    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 Fp4Element { }

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


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

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

/// The montgomery trait bound restricts the modular to odd prime
impl MulInverse for Fp4Element {
    fn mul_inv(self) -> Self {
        self.mont_inv().mont_mul_fp(BN254::R_INV_P)
    }
}


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

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

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

impl MontgomeryOps<U256, BN254> for Fp4Element {
    /// input: 
    ///   - lhs: `a0 + b0x`
    ///   - rhs: `a1 + b1x`
    /// 
    /// output: 
    /// 
    /// ```text,ignore
    ///   (a0 + b0x) * (a1 + b1x) mod (s^2 - \xi)
    /// = a0a1 + (a0b1 + b0a1)s + b0b1s^2 mod (s^2 - \xi)
    /// = a0a1 + b0b1\xi + (a0b1 + b0a1)s
    /// ```
    /// 
    /// This operations needs **3 mul + 5 add/sub + 1 mul_xi on Fp2**
    fn mont_mul(self, rhs: Self) -> Self {
        let (a0, b0) = (self.0, self.1);
        let (a1, b1) = (rhs.0, rhs.1);

        let a0a1 = a0.mont_mul(a1);
        let b0b1 = b0.mont_mul(b1);

        Self(
            a0a1 + b0b1.mul_xi(), 
             // a0b1 + a1b0 = (a0 + b0) * (a1 + b1) - a0a1 - b0b1
            (a0 + b0).mont_mul(a1 + b1) - a0a1 - b0b1
        )
    }

    /// input: Fp4 element `a + bs`
    /// 
    /// output: Fp4 element `(a + bs)^2`
    /// 
    /// Here is a faster algorithm to do this
    /// 
    /// ```text,ignore
    ///   (a + bs)^2
    /// = a^2 + 2abs + b^2s^2
    /// = (a^2 + b^2\xi) + 2abs
    /// = (a + b\xi)(a + b) - ab - ab\xi + 2abs
    /// ```
    /// 
    /// This operation needs **2 mul + 5 add/sub + 2 mul_xi on Fp2**
    fn mont_sqr(self) -> Self {
        let (a, b) = (self.0, self.1);

        let ab = a.mont_mul(b);
        Self(
            (a + b.mul_xi()).mont_mul(a + b) - ab - ab.mul_xi(), 
            ab + ab
        )
    }

    /// input: Fp4 element `a + bs`
    /// 
    /// output: Fp4 element `(c + ds) * RR` so that `(a + bs)(c + ds) = 1`
    /// 
    /// we need to make sure
    /// 
    /// ```text,ignore
    /// ac + bd\xi = 1
    /// bc + ad = 0
    /// ```
    /// 
    /// solve these equations we will get
    /// 
    /// ```text,ignore
    /// c = a / (a^2 - b^2 * \xi)
    /// d = -b / (a^2 - b^2 * \xi)
    /// ```
    fn mont_inv(self) -> Self {
        // t = (a^2 * R^-1 + b^2 * \xi * R^-1)^-1 * RR
        let t = (self.0.mont_sqr() - self.1.mont_sqr().mul_xi()).mont_inv();

        Self(self.0.mont_mul(t), -self.1.mont_mul(t))
    }

    /// input: Fp4 element `a + bs`, Fp element `c`
    /// 
    /// output: `(ac + bcs) * R^-1`
    fn mont_mul_fp(self, rhs: FpElement) -> Self {
        Self(self.0.mont_mul_fp(rhs), self.1.mont_mul_fp(rhs))
    }

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

impl Fp4Element  {
    /// input: Fp4 element `a + bs`
    /// 
    /// output: Fp4 element `(a + bs) * s`
    /// 
    /// ```text,ignore
    ///   (a + bs) * s mod (s^2 - \xi)
    /// = as + bs^2 mod (s^2 - \xi)
    /// = b\xi + as
    /// ```
    pub fn mul_s(self) -> Self {
        Self(self.1.mul_xi(), self.0)
    }

    /// input: Fp4 element `a + bs`, Fp2 element `c`
    /// 
    /// output: `(ac + bcs) * R^-1`
    pub fn sparse_mul(self, rhs: Fp2Element) -> Self {
        Self(self.0.mont_mul(rhs), self.1.mont_mul(rhs))
    }

    /// TODO
    pub fn conjugate(self) -> Self {
        Self(self.0, -self.1)
    }

    /// TODO
    pub fn map_frob(self) -> Self {
        Self(self.0.map_frob(), self.1.map_frob().mont_mul(BN254::XI_MONT[2]))
    }
}

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

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