cryptix-bn254 0.1.0

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

use cryptix_bigint::property::IsBigInt;
use cryptix_ecc::CurvePoint;
use cryptix_field::field::{montgomery::{MontgomeryOps, Montgomery}, Field, MulIdentity};
use cryptix_field::group::{AbelianGroup, CommunicativeAdd, Group, AddIdentity, AssociativeAdd};

use crate::{galoisfield::{FpElement, BN254, U256}, fp};

/// Point using Jacobian projective coordinate on E1
#[derive(PartialEq, Eq, Clone, Copy)]
pub struct BN254Fp {
    pub x: FpElement,
    pub y: FpElement,
    pub z: FpElement
}

impl core::fmt::Debug for BN254Fp {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "E1Point {{\n    ")?;
        write!(f, "x: fp!(\"")?;
        self.x.fmt(f)?;
        write!(f, "\"), \n    ")?;
        write!(f, "y: fp!(\"")?;
        self.y.fmt(f)?;
        write!(f, "\"), \n    ")?;
        write!(f, "z: fp!(\"")?;
        self.z.fmt(f)?;
        write!(f, "\"), \n}}")
    }
}

impl BN254Fp {
    /// create a instance of e1 point at infinity
    pub const fn inf() -> Self {
        Self { x: FpElement::ZERO, y: FpElement::ONE, z: FpElement::ZERO }
    }
}

impl Add for BN254Fp {
    type Output = BN254Fp;

    fn add(self, rhs: Self) -> Self::Output {
        if self.at_inf() {
            return rhs
        }

        if rhs.at_inf() {
            return self
        }

        let t1 = self.z.mont_sqr();   // 1
        let t2 = self.z.mont_mul(t1);
        let t1 = t1.mont_mul(rhs.x);
        let t2 = t2.mont_mul(rhs.y);
        let t3 = rhs.z;               // 5
        let y = self.y.mont_mul(t3);
        let z = self.z.mont_mul(t3);
        let t3 = t3.mont_sqr();
        let y = y.mont_mul(t3);
        let x = self.x.mont_mul(t3);  // 10
        let t1 = t1 - x;
        let z = z.mont_mul(t1);
        let t2 = t2 - y;
        let t3 = t1.mont_sqr();       
        let t1 = t1.mont_mul(t3);     // 15
        let t3 = t3.mont_mul(x);
        let x = t2.mont_sqr();
        let x = x - t3 - t3 - t1;
        let t3 = (t3 - x).mont_mul(t2);
        let t1 = t1.mont_mul(y);
        let y = t3 - t1;

        Self { x, y, z }
    }
}

impl AssociativeAdd for BN254Fp { }

impl CommunicativeAdd for BN254Fp { }

impl AddIdentity for BN254Fp {
    const ADD_IDENTITY: Self = Self::inf();
}

impl Neg for BN254Fp {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self { x: self.x, y: -self.y, z: self.z }
    }
}

impl Group for BN254Fp { }

impl AbelianGroup for BN254Fp { }

impl CurvePoint for BN254Fp {
    type MulScalar = U256;
    const GENERATOR: Self = BN254Fp {
        x: fp!("03F7BF8FC000000C176E1E800000003AA7E70000000000899100000000000085"),
        y: fp!("212BA4F27FFFFFF5A2C62EFFFFFFFFCDB939FFFFFFFFFF8A15FFFFFFFFFFFF8E"),
        z: fp!("212BA4F27FFFFFF5A2C62EFFFFFFFFCDB939FFFFFFFFFF8A15FFFFFFFFFFFF8E")
    };

    fn double(self) -> Self {
        if self.at_inf() {
            return self
        }

        let t2 = self.x.mont_sqr();
        let t1 = t2 + t2 + t2;
        let y = self.y + self.y;
        let z = self.z.mont_mul(y);
        let y = y.mont_sqr();
        let t2 = y.mont_mul(self.x);
        let y = y.mont_sqr().hlv();
        let x = t1.mont_sqr() - t2 - t2;
        let t2 = t2 - x;
        let t1 = t1.mont_mul(t2);
        let y = t1 - y;

        Self { x, y, z }
    }

    fn scalar_mul(self, k: U256) -> Self {
        const W: usize = 4;
        const LEN: usize = 16;
        // mask is 0xf
        const MASK: <U256 as IsBigInt>::Dig = 0xf;
        let mut kp = [BN254Fp::inf(); LEN];

        kp[1] = self;
        kp[2] = self.double();

        for i in 3..LEN {
            // i is even
            if i & 1 == 0 {
                kp[i] = kp[i >> 1].double();
            } else {
                kp[i] = kp[i - 1].mix_add(self);
            }

            debug_assert!(kp[i].on_curve())
        }

        let ki = ((k[U256::DIG_LEN - 1] >> (W * (U256::DIG_BIT_LEN / W - 1))) & MASK) as usize;
        let mut q = kp[ki];

        for j in (0..(U256::DIG_BIT_LEN / W - 1)).rev() {
            let ki = ((k[U256::DIG_LEN - 1] >> (W * j)) & MASK) as usize;
            for _ in 0..W {
                q = q.double();
            }
            q = q + kp[ki]
        }

        for i in (0..(U256::DIG_LEN - 1)).rev() {
            for j in (0..(U256::DIG_BIT_LEN / W)).rev() {
                let ki = ((k[i] >> (W * j)) & MASK) as usize;
                for _ in 0..W {
                    q = q.double();
                }
                q = q + kp[ki]
            }
        }

        q
    }

    fn at_inf(&self) -> bool {
        self.z.is_zero()
    }

    /// check if y^2 == x^3 + 2z^6
    fn on_curve(&self) -> bool {
        let y2 = self.y.mont_sqr();
        
        let x3 = self.x.mont_sqr();
        let x3 = x3.mont_mul(self.x);

        let z6 = self.z.mont_sqr();
        let z6 = self.z.mont_mul(z6);
        let z6 = z6.mont_sqr();

        let r = (z6 + z6 + x3).mont_rdc();
        let l = y2.mont_rdc();

        l == r
    }

    fn normalize(self) -> Self {
        let t1 = self.z.mont_inv();
        let t2 = t1.mont_sqr();
        let x = self.x.mont_mul(t2);
        let t2 = t2.mont_mul(t1);
        let y = self.y.mont_mul(t2);

        Self { x, y, z: BN254::R_P }
    }

    fn mont_form(self) -> Self {
        Self { 
            x: self.x.mont_form(), 
            y: self.y.mont_form(), 
            z: self.z.mont_form() 
        }
    }

    fn mont_rdc(self) -> Self {
        Self { 
            x: self.x.mont_rdc(), 
            y: self.y.mont_rdc(), 
            z: self.z.mont_rdc() 
        }
    }
}

impl BN254Fp {
    pub fn mix_add(self, mut rhs: Self) -> Self {
        // TODO: we show apply this restriction with type
        if rhs.z != BN254::R_P {
            rhs = rhs.normalize();
        }

        if self.at_inf() { return rhs }
        if rhs.at_inf() { return self }

        let t1 = self.z.mont_sqr();   // 1
        let t2 = self.z.mont_mul(t1);
        let t1 = t1.mont_mul(rhs.x);
        let t2 = t2.mont_mul(rhs.y);
        let t1 = t1 - self.x;         // 5
        let z = self.z.mont_mul(t1);
        let t2 = t2 - self.y;
        let t3 = t1.mont_sqr();
        let t1 = t1.mont_mul(t3);
        let t3 = t3.mont_mul(self.x);
        let x = t2.mont_sqr();
        let x = x - t3 - t3 - t1;
        let t3 = t3 - x;
        let t3 = t3.mont_mul(t2);
        let t1 = t1.mont_mul(self.y);
        let y = t3 - t1;

        Self { x, y, z }
    }
}