//! The prime fields Prio3 uses (`draft-irtf-cfrg-vdaf-22` §6.1): Field64 and Field128, their
//! little-endian encoding (§6.1.1), and their NTT-friendly interface (§6.1.2).
//!
//! Field128's modulus is 2^128 − 28·2^64 + 1, so 2^128 is congruent to 28·2^64 − 1 and a
//! 256-bit product folds back below 2^128 in three steps.

use std::fmt;
use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};

/// A prime field of the draft's Table 4.
pub trait Field:
    Copy
    + Clone
    + PartialEq
    + Eq
    + fmt::Debug
    + Add<Output = Self>
    + Sub<Output = Self>
    + Mul<Output = Self>
    + Neg<Output = Self>
    + AddAssign
    + SubAssign
    + MulAssign
    + 'static
{
    /// The prime modulus.
    const MODULUS: u128;
    /// Bytes of one encoded element.
    const ENCODED_SIZE: usize;
    /// log2 of the order of the multiplicative subgroup `generator()` generates.
    const GEN_ORDER_LOG2: u32;
    /// The element of value `v`, which is below `MODULUS`.
    fn from_u128(v: u128) -> Self;
    /// The element's value, below `MODULUS`.
    fn int(self) -> u128;
    /// The generator of the subgroup of order 2^`GEN_ORDER_LOG2`: 7^((p − 1) / 2^k).
    fn generator() -> Self;

    fn zero() -> Self {
        Self::from_u128(0)
    }
    fn one() -> Self {
        Self::from_u128(1)
    }
    /// An integer in (−MODULUS, MODULUS); a negative one is a negation.
    fn from_i128(v: i128) -> Self {
        if v < 0 { -Self::from_u128(v.unsigned_abs()) } else { Self::from_u128(v as u128) }
    }
    fn pow(self, mut e: u128) -> Self {
        let mut base = self;
        let mut r = Self::one();
        while e > 0 {
            if e & 1 == 1 {
                r *= base;
            }
            base *= base;
            e >>= 1;
        }
        r
    }
    /// The multiplicative inverse (Fermat); the inverse of zero is zero.
    fn inv(self) -> Self {
        self.pow(Self::MODULUS - 2)
    }
    /// The principal `n`-th root of unity, gen^(GEN_ORDER / n), `n` a power of two.
    fn nth_root(n: usize) -> Self {
        assert!(n.is_power_of_two() && (n.trailing_zeros()) <= Self::GEN_ORDER_LOG2);
        Self::generator().pow(1u128 << (Self::GEN_ORDER_LOG2 - n.trailing_zeros()))
    }
}

/// Encode a vector of field elements, little-endian, `ENCODED_SIZE` bytes each (§6.1.1).
pub fn encode_vec<F: Field>(v: &[F]) -> Vec<u8> {
    let mut out = Vec::with_capacity(v.len() * F::ENCODED_SIZE);
    for x in v {
        out.extend_from_slice(&x.int().to_le_bytes()[..F::ENCODED_SIZE]);
    }
    out
}

/// Decode a vector of field elements; a length not a multiple of the element size, or an
/// element not below the modulus, is refused (§6.1.1).
pub fn decode_vec<F: Field>(b: &[u8]) -> Option<Vec<F>> {
    if !b.len().is_multiple_of(F::ENCODED_SIZE) {
        return None;
    }
    b.chunks(F::ENCODED_SIZE)
        .map(|c| {
            let mut w = [0u8; 16];
            w[..c.len()].copy_from_slice(c);
            let x = u128::from_le_bytes(w);
            (x < F::MODULUS).then(|| F::from_u128(x))
        })
        .collect()
}

pub fn vec_add<F: Field>(a: &[F], b: &[F]) -> Vec<F> {
    assert_eq!(a.len(), b.len(), "mismatched vector sizes");
    a.iter().zip(b).map(|(x, y)| *x + *y).collect()
}

pub fn vec_sub<F: Field>(a: &[F], b: &[F]) -> Vec<F> {
    assert_eq!(a.len(), b.len(), "mismatched vector sizes");
    a.iter().zip(b).map(|(x, y)| *x - *y).collect()
}

// ---------------------------------------------------------------------------------------
// Field64: p = 2^32 · 4294967295 + 1 = 2^64 − 2^32 + 1.

const P64: u64 = 0xFFFF_FFFF_0000_0001;

#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct F64(u64);

impl fmt::Debug for F64 {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "F64({})", self.0)
    }
}

impl Add for F64 {
    type Output = F64;
    fn add(self, o: F64) -> F64 {
        F64(((self.0 as u128 + o.0 as u128) % P64 as u128) as u64)
    }
}
impl Sub for F64 {
    type Output = F64;
    fn sub(self, o: F64) -> F64 {
        F64(((self.0 as u128 + P64 as u128 - o.0 as u128) % P64 as u128) as u64)
    }
}
impl Mul for F64 {
    type Output = F64;
    fn mul(self, o: F64) -> F64 {
        F64(((self.0 as u128 * o.0 as u128) % P64 as u128) as u64)
    }
}
impl Neg for F64 {
    type Output = F64;
    fn neg(self) -> F64 {
        F64::zero() - self
    }
}
impl AddAssign for F64 {
    fn add_assign(&mut self, o: F64) {
        *self = *self + o;
    }
}
impl SubAssign for F64 {
    fn sub_assign(&mut self, o: F64) {
        *self = *self - o;
    }
}
impl MulAssign for F64 {
    fn mul_assign(&mut self, o: F64) {
        *self = *self * o;
    }
}

impl Field for F64 {
    const MODULUS: u128 = P64 as u128;
    const ENCODED_SIZE: usize = 8;
    const GEN_ORDER_LOG2: u32 = 32;
    fn from_u128(v: u128) -> F64 {
        assert!(v < Self::MODULUS, "field element out of range");
        F64(v as u64)
    }
    fn int(self) -> u128 {
        self.0 as u128
    }
    fn generator() -> F64 {
        F64(7).pow(4_294_967_295)
    }
}

// ---------------------------------------------------------------------------------------
// Field128: p = 2^66 · 4611686018427387897 + 1 = 2^128 − 28·2^64 + 1.

const P128: u128 = 0xFFFF_FFFF_FFFF_FFE4_0000_0000_0000_0001;
/// 2^128 − p: 2^128 is congruent to it modulo p.
const C128: u128 = 0x1B_FFFF_FFFF_FFFF_FFFF;

/// The 256-bit product of two 128-bit integers, as (high, low).
fn wide_mul(a: u128, b: u128) -> (u128, u128) {
    const M: u128 = u64::MAX as u128;
    let (a0, a1) = (a & M, a >> 64);
    let (b0, b1) = (b & M, b >> 64);
    let p00 = a0 * b0;
    let p01 = a0 * b1;
    let p10 = a1 * b0;
    let p11 = a1 * b1;
    let mid = (p00 >> 64) + (p01 & M) + (p10 & M);
    let lo = (p00 & M) | (mid << 64);
    let hi = p11 + (p01 >> 64) + (p10 >> 64) + (mid >> 64);
    (hi, lo)
}

/// hi · 2^128 + lo modulo p.
fn reduce128(hi: u128, lo: u128) -> u128 {
    // hi · 2^128 + lo ≡ hi · C + lo, with hi · C below 2^197.
    let (h1, l1) = wide_mul(hi, C128);
    let (l, c) = l1.overflowing_add(lo);
    let h1 = h1 + c as u128;
    // ≡ h1 · C + l, h1 below 2^70, h1 · C below 2^139.
    let (h2, l2) = wide_mul(h1, C128);
    let (l, c) = l2.overflowing_add(l);
    let h2 = h2 + c as u128;
    // ≡ h2 · C + l, h2 below 2^12, h2 · C below 2^82.
    let (l, c) = l.overflowing_add(h2 * C128);
    // A carry out leaves l below 2^82, so l + C does not overflow.
    let mut r = if c { l + C128 } else { l };
    while r >= P128 {
        r -= P128;
    }
    r
}

#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct F128(u128);

impl fmt::Debug for F128 {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "F128({})", self.0)
    }
}

impl Add for F128 {
    type Output = F128;
    fn add(self, o: F128) -> F128 {
        let (s, c) = self.0.overflowing_add(o.0);
        F128(if c || s >= P128 { s.wrapping_sub(P128) } else { s })
    }
}
impl Sub for F128 {
    type Output = F128;
    fn sub(self, o: F128) -> F128 {
        F128(if self.0 >= o.0 { self.0 - o.0 } else { (P128 - o.0) + self.0 })
    }
}
impl Mul for F128 {
    type Output = F128;
    fn mul(self, o: F128) -> F128 {
        let (h, l) = wide_mul(self.0, o.0);
        F128(reduce128(h, l))
    }
}
impl Neg for F128 {
    type Output = F128;
    fn neg(self) -> F128 {
        F128::zero() - self
    }
}
impl AddAssign for F128 {
    fn add_assign(&mut self, o: F128) {
        *self = *self + o;
    }
}
impl SubAssign for F128 {
    fn sub_assign(&mut self, o: F128) {
        *self = *self - o;
    }
}
impl MulAssign for F128 {
    fn mul_assign(&mut self, o: F128) {
        *self = *self * o;
    }
}

impl Field for F128 {
    const MODULUS: u128 = P128;
    const ENCODED_SIZE: usize = 16;
    const GEN_ORDER_LOG2: u32 = 66;
    fn from_u128(v: u128) -> F128 {
        assert!(v < Self::MODULUS, "field element out of range");
        F128(v)
    }
    fn int(self) -> u128 {
        self.0
    }
    fn generator() -> F128 {
        F128(7).pow(4_611_686_018_427_387_897)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// a · b mod p by doubling and adding, which needs no wide arithmetic.
    fn slow_mul(a: u128, b: u128, p: u128) -> u128 {
        let add = |x: u128, y: u128| {
            let (s, c) = x.overflowing_add(y);
            if c || s >= p { s.wrapping_sub(p) } else { s }
        };
        let mut r = 0u128;
        for i in (0..128).rev() {
            r = add(r, r);
            if (b >> i) & 1 == 1 {
                r = add(r, a);
            }
        }
        r
    }

    #[test]
    fn the_moduli_are_the_drafts() {
        assert_eq!(F64::MODULUS, (1u128 << 32) * 4_294_967_295 + 1);
        assert_eq!(F128::MODULUS, (1u128 << 66) * 4_611_686_018_427_387_897 + 1);
        assert_eq!(C128, 0u128.wrapping_sub(P128));
    }

    #[test]
    fn field128_multiplication_agrees_with_doubling() {
        let mut x = 0x0123_4567_89ab_cdef_fedc_ba98_7654_3210u128;
        let edge = [0, 1, 2, P128 - 1, P128 - 2, C128, 1u128 << 127, (1u128 << 64) - 1, 1u128 << 64];
        let mut vals: Vec<u128> = edge.to_vec();
        for _ in 0..200 {
            x = x.wrapping_mul(0x9e37_79b9_7f4a_7c15_f39c_c060_5ced_c834).wrapping_add(0x2545_f491_4f6c_dd1d);
            vals.push(x % P128);
        }
        for &a in &vals {
            for &b in vals.iter().step_by(7) {
                assert_eq!((F128(a) * F128(b)).0, slow_mul(a, b, P128), "{a} * {b}");
                assert_eq!((F64((a % P64 as u128) as u64) * F64((b % P64 as u128) as u64)).0 as u128, slow_mul(a % P64 as u128, b % P64 as u128, P64 as u128));
            }
        }
    }

    #[test]
    fn generators_have_the_stated_order() {
        let g = F64::generator();
        assert_eq!(g.pow(1u128 << 32), F64::one());
        assert_ne!(g.pow(1u128 << 31), F64::one());
        let g = F128::generator();
        assert_eq!(g.pow(1u128 << 66), F128::one());
        assert_ne!(g.pow(1u128 << 65), F128::one());
        let w = F128::nth_root(8);
        assert_eq!(w.pow(8), F128::one());
        assert_eq!(w * w.inv(), F128::one());
    }

    #[test]
    fn decoding_refuses_an_element_past_the_modulus() {
        let mut b = (P128).to_le_bytes().to_vec();
        assert!(decode_vec::<F128>(&b).is_none());
        b[0] -= 1;
        assert_eq!(decode_vec::<F128>(&b).unwrap(), vec![F128(P128 - 1)]);
        assert!(decode_vec::<F64>(&[0u8; 7]).is_none());
    }
}
