//! Arithmetic modulo the order `L` of the prime-order subgroup generated by the Ed25519 base
//! point, `L = 2^252 + 27742317777372353535851937790883648493`.

use subtle::{Choice, ConditionallySelectable};
use zeroize::Zeroize;

/// `L`, as four little-endian 64-bit limbs.
const L: [u64; 4] = [
    0x5812631a5cf5d3ed,
    0x14def9dea2f79cd6,
    0x0000000000000000,
    0x1000000000000000,
];

/// `floor(2^512 / L)`, the Barrett reduction constant for reducing a 512-bit value modulo `L`
/// (see [`barrett_reduce`]). `L` is just over `2^252`, so this is just over `2^260`, five
/// little-endian 64-bit limbs.
const MU: [u64; 5] = [
    0xed9ce5a30a2c131b,
    0x2106215d086329a7,
    0xffffffffffffffeb,
    0xffffffffffffffff,
    0x000000000000000f,
];

/// An integer modulo `L`, always canonically reduced (`< L`).
#[derive(Copy, Clone, Debug, Zeroize)]
pub struct Scalar(pub [u64; 4]);

/// Returns `true` if `a < b`, comparing as 256-bit unsigned integers (little-endian limbs).
fn limbs_lt(a: &[u64; 4], b: &[u64; 4]) -> bool {
    for i in (0..4).rev() {
        if a[i] != b[i] {
            return a[i] < b[i];
        }
    }
    false
}

/// Returns `a - b` modulo `2^(64*N)` and whether the subtraction underflowed.
fn limbs_sub_with_borrow<const N: usize>(a: &[u64; N], b: &[u64; N]) -> ([u64; N], bool) {
    let mut out = [0u64; N];
    let mut borrow = false;
    for i in 0..N {
        let (d1, b1) = a[i].overflowing_sub(b[i]);
        let (d2, b2) = d1.overflowing_sub(borrow as u64);
        out[i] = d2;
        borrow = b1 | b2;
    }
    (out, borrow)
}

/// Returns `a - b`, assuming `a >= b`.
fn limbs_sub(a: &[u64; 4], b: &[u64; 4]) -> [u64; 4] {
    limbs_sub_with_borrow(a, b).0
}

/// Returns `a + b`, assuming the true sum is `< 2^256` (no carry out of the top limb).
fn limbs_add(a: &[u64; 4], b: &[u64; 4]) -> [u64; 4] {
    let mut out = [0u64; 4];
    let mut carry = false;
    for i in 0..4 {
        let (sum, c1) = a[i].overflowing_add(b[i]);
        let (sum, c2) = sum.overflowing_add(carry as u64);
        out[i] = sum;
        carry = c1 | c2;
    }
    out
}

/// Selects `b` when `select_b` is set, in constant time via `subtle`'s optimization barrier.
fn limbs_select<const N: usize>(a: &[u64; N], b: &[u64; N], select_b: Choice) -> [u64; N] {
    core::array::from_fn(|i| u64::conditional_select(&a[i], &b[i], select_b))
}

/// Subtracts `b` from `a` when `a >= b`, leaving `a` unchanged otherwise.
fn limbs_conditional_sub<const N: usize>(a: &[u64; N], b: &[u64; N]) -> [u64; N] {
    let (difference, underflow) = limbs_sub_with_borrow(a, b);
    limbs_select(&difference, a, Choice::from(underflow as u8))
}

/// Returns `(a - b) mod 2^320`, as five little-endian 64-bit limbs: the final borrow (if any) is
/// discarded, which is exactly right both when `a >= b` (ordinary subtraction, no borrow) and,
/// in [`barrett_reduce`], when the true difference is negative (the discarded borrow is
/// equivalent to adding `2^320` back in).
fn limbs_sub5(a: &[u64; 5], b: &[u64; 5]) -> [u64; 5] {
    limbs_sub_with_borrow(a, b).0
}

/// Returns the full 512-bit product `a * b` as eight little-endian 64-bit limbs, via the standard
/// schoolbook multiply-accumulate-with-carry ("Comba") method.
fn limbs_mul_wide(a: &[u64; 4], b: &[u64; 4]) -> [u64; 8] {
    let mut t = [0u64; 8];
    for i in 0..4 {
        let mut carry = 0u128;
        for j in 0..4 {
            let sum = t[i + j] as u128 + (a[i] as u128) * (b[j] as u128) + carry;
            t[i + j] = sum as u64;
            carry = sum >> 64;
        }
        t[i + 4] = carry as u64;
    }
    t
}

/// Returns the full product `a * b` as ten little-endian 64-bit limbs, via the same
/// schoolbook method as [`limbs_mul_wide`].
fn mul5x5(a: &[u64; 5], b: &[u64; 5]) -> [u64; 10] {
    let mut t = [0u64; 10];
    for i in 0..5 {
        let mut carry = 0u128;
        for j in 0..5 {
            let sum = t[i + j] as u128 + (a[i] as u128) * (b[j] as u128) + carry;
            t[i + j] = sum as u64;
            carry = sum >> 64;
        }
        t[i + 5] = carry as u64;
    }
    t
}

/// Returns the full product `a * b` as nine little-endian 64-bit limbs, via the same schoolbook
/// method as [`limbs_mul_wide`].
fn mul5x4(a: &[u64; 5], b: &[u64; 4]) -> [u64; 9] {
    let mut t = [0u64; 9];
    for i in 0..5 {
        let mut carry = 0u128;
        for j in 0..4 {
            let sum = t[i + j] as u128 + (a[i] as u128) * (b[j] as u128) + carry;
            t[i + j] = sum as u64;
            carry = sum >> 64;
        }
        t[i + 4] = carry as u64;
    }
    t
}

/// Reduces a 512-bit value (eight little-endian 64-bit limbs) modulo `L` with a fixed number of
/// limb multiplications using Barrett reduction (Handbook of Applied Cryptography, Algorithm
/// 14.42).
///
/// With word size `b = 2^64` and `L` fitting in `k = 4` words, this reduces any `x < b^(2k) =
/// 2^512`: `q = floor(floor(x / b^(k-1)) * MU / b^(k+1))` approximates `floor(x / L)` (using the
/// precomputed `MU = floor(b^(2k) / L)`), accurate enough that `x - q*L` is guaranteed `< 3L`, so
/// at most two trial subtractions of `L` remain to reach the canonical residue.
fn barrett_reduce(x: [u64; 8]) -> Scalar {
    let q1: [u64; 5] = x[3..8].try_into().expect("slice has 5 elements");
    let q2 = mul5x5(&q1, &MU);
    let q3: [u64; 5] = q2[5..10].try_into().expect("slice has 5 elements");

    let r1: [u64; 5] = x[0..5].try_into().expect("slice has 5 elements");
    let r2: [u64; 5] = mul5x4(&q3, &L)[0..5]
        .try_into()
        .expect("slice has 5 elements");
    let mut r = limbs_sub5(&r1, &r2);

    let l5 = [L[0], L[1], L[2], L[3], 0];
    // The Barrett quotient estimate leaves a residue below `3L`, so two fixed conditional
    // subtractions are sufficient and avoid data-dependent control flow for secret scalars.
    r = limbs_conditional_sub(&r, &l5);
    r = limbs_conditional_sub(&r, &l5);
    debug_assert_eq!(r[4], 0, "residue must fit in L's 4 limbs after reduction");
    Scalar([r[0], r[1], r[2], r[3]])
}

impl Scalar {
    pub const ZERO: Self = Self([0, 0, 0, 0]);

    /// Returns the additive inverse modulo `L`.
    pub fn neg_mod_l(&self) -> Self {
        if self.0 == Self::ZERO.0 {
            *self
        } else {
            Self(limbs_sub(&L, &self.0))
        }
    }

    /// Interprets `bytes` as a little-endian integer and rejects it unless it is already the
    /// canonical representative (`< L`), as required for the `s` component of a signature.
    pub fn from_canonical_bytes(bytes: &[u8; 32]) -> Option<Self> {
        let mut limbs = [0u64; 4];
        for (i, limb) in limbs.iter_mut().enumerate() {
            let mut chunk = [0u8; 8];
            chunk.copy_from_slice(&bytes[i * 8..i * 8 + 8]);
            *limb = u64::from_le_bytes(chunk);
        }
        limbs_lt(&limbs, &L).then_some(Self(limbs))
    }

    /// Reduces a 64-byte little-endian integer (e.g. a SHA-512 digest) modulo `L`, via
    /// [`barrett_reduce`].
    pub fn from_bytes_mod_order_wide(bytes: &[u8; 64]) -> Self {
        let limbs = core::array::from_fn(|i| {
            let mut chunk = [0u8; 8];
            chunk.copy_from_slice(&bytes[i * 8..i * 8 + 8]);
            u64::from_le_bytes(chunk)
        });
        barrett_reduce(limbs)
    }

    /// Returns the bits of this scalar's canonical representative, most significant first.
    pub fn bits_be(&self) -> impl Iterator<Item = bool> + '_ {
        (0..256)
            .rev()
            .map(move |i| (self.0[i / 64] >> (i % 64)) & 1 == 1)
    }

    /// Returns the canonical little-endian encoding of this scalar.
    pub fn to_bytes(self) -> [u8; 32] {
        let mut bytes = [0u8; 32];
        for (chunk, limb) in bytes.as_chunks_mut::<8>().0.iter_mut().zip(self.0) {
            *chunk = limb.to_le_bytes();
        }
        bytes
    }

    /// Constructs a scalar from a little-endian 128-bit value. Always canonical, since every
    /// 128-bit value is `< L` (`L > 2^252`).
    pub const fn from_u128(value: u128) -> Self {
        Self([value as u64, (value >> 64) as u64, 0, 0])
    }

    /// Returns the base-`2^width` digit at position `index`, i.e. bits `[index*width,
    /// index*width+width)` of this scalar's canonical representative, as an unsigned integer.
    pub const fn window(&self, index: usize, width: u32) -> usize {
        let bit_start = index * width as usize;
        if bit_start >= 256 {
            return 0;
        }
        let limb_index = bit_start / 64;
        let bit_offset = (bit_start % 64) as u32;
        let mut digit = self.0[limb_index] >> bit_offset;
        if bit_offset + width > 64 && limb_index + 1 < 4 {
            digit |= self.0[limb_index + 1] << (64 - bit_offset);
        }
        (digit as usize) & ((1usize << width) - 1)
    }

    /// Recodes this scalar into `N` signed, base-`2^width` digits (each in
    /// `[-2^(width-1), 2^(width-1) - 1]`), least-significant first. Used by the MSM's bucket
    /// method (see [`super::msm`]): a signed digit only ever needs a bucket for its *magnitude*
    /// (`1..=2^(width-1)`, half as many as the `1..2^width` an unsigned digit needs), at the cost
    /// of negating the affine point's X and T coordinates when its digit is negative.
    ///
    /// This is exactly the standard signed-digit recoding: process [`Scalar::window`]'s unsigned
    /// digits from least to most significant, and whenever one is `>= 2^(width-1)` (the upper half
    /// of its range), replace it with `digit - 2^width` (negative, same residue) and carry a `+1`
    /// into the next digit -- since `digit * 2^(width*i) = (digit - 2^width) * 2^(width*i) +
    /// 2^width * 2^(width*i)`, and that `2^width` term is exactly one unit of the next digit's
    /// weight. `N` must be large enough that the final carry (at most `1`) has a digit to land in;
    /// `256usize.div_ceil(width) + 1` unsigned windows' worth is always enough.
    pub const fn signed_digits<const N: usize>(&self, width: u32) -> [i32; N] {
        let half = 1i64 << (width - 1);
        let full = 1i64 << width;
        let mut digits = [0i32; N];
        let mut carry = 0i64;
        let mut i = 0;
        while i < N {
            let raw = self.window(i, width) as i64 + carry;
            if raw >= half {
                digits[i] = (raw - full) as i32;
                carry = 1;
            } else {
                digits[i] = raw as i32;
                carry = 0;
            }
            i += 1;
        }
        digits
    }

    /// Returns `(self + rhs) mod L`, assuming both operands are already `< L`.
    pub fn add_mod_l(&self, rhs: &Self) -> Self {
        let sum = limbs_add(&self.0, &rhs.0);
        Self(limbs_conditional_sub(&sum, &L))
    }

    /// Returns `(self * rhs) mod L`.
    ///
    /// Computes the full 512-bit product via schoolbook multiplication, then reduces it via
    /// [`barrett_reduce`].
    pub fn mul_mod_l(&self, rhs: &Self) -> Self {
        barrett_reduce(limbs_mul_wide(&self.0, &rhs.0))
    }
}

#[cfg(test)]
impl<'a> arbitrary::Arbitrary<'a> for Scalar {
    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
        Ok(match u.int_in_range(0u8..=5)? {
            0 => Self::ZERO,
            1 => Self::from_u128(1),
            2 => {
                let mut l_minus_one = L;
                l_minus_one[0] -= 1;
                Self(l_minus_one)
            }
            3 => Self::from_u128(u.arbitrary()?),
            _ => Self::from_bytes_mod_order_wide(&u.arbitrary()?),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::{L, Scalar, limbs_lt, limbs_sub};
    use commonware_invariants::minifuzz::Builder;
    use zeroize::Zeroize;

    /// Returns `(2*a + bit) mod 2^256`, the left-shift-by-one step of the bit-serial reduction
    /// [`reduce_wide_naive`] uses.
    fn limbs_shl1(a: &[u64; 4], bit: u64) -> [u64; 4] {
        let mut out = [0u64; 4];
        let mut carry = bit;
        for i in 0..4 {
            out[i] = (a[i] << 1) | carry;
            carry = a[i] >> 63;
        }
        out
    }

    /// A bit-serial double-and-add-mod reduction used as a differential-test oracle for
    /// [`Scalar::from_bytes_mod_order_wide`]. It performs one conditional subtraction per input
    /// bit, making it slow but straightforward to audit.
    fn reduce_wide_naive(bytes: &[u8; 64]) -> Scalar {
        let mut r = [0u64; 4];
        for byte_index in (0..64).rev() {
            let byte = bytes[byte_index];
            for bit_index in (0..8).rev() {
                let bit = (byte >> bit_index) & 1;
                r = limbs_shl1(&r, bit as u64);
                if !limbs_lt(&r, &L) {
                    r = limbs_sub(&r, &L);
                }
            }
        }
        Scalar(r)
    }

    #[test]
    fn from_bytes_mod_order_wide_matches_naive_reference() {
        Builder::default()
            .with_seed(0)
            .with_search_limit(1024)
            .test(|u| {
                let bytes: [u8; 64] = u.arbitrary()?;
                let expected = reduce_wide_naive(&bytes);
                let actual = Scalar::from_bytes_mod_order_wide(&bytes);
                assert_eq!(actual.0, expected.0);
                Ok(())
            });
    }

    #[test]
    fn from_bytes_mod_order_wide_matches_naive_reference_on_edge_cases() {
        // `x = 0`, `x = L - 1`, `x = L`, `x = L + 1`, and `x = 2^512 - 1` (the Barrett quotient
        // approximation's error is largest for inputs near the top of the input range).
        let l_minus_1 = {
            let mut bytes = [0u8; 64];
            for (i, limb) in L.iter().enumerate() {
                bytes[i * 8..i * 8 + 8].copy_from_slice(&(limb - u64::from(i == 0)).to_le_bytes());
            }
            bytes
        };
        let mut cases = vec![[0u8; 64], l_minus_1];
        let mut l_bytes = [0u8; 64];
        for (i, limb) in L.iter().enumerate() {
            l_bytes[i * 8..i * 8 + 8].copy_from_slice(&limb.to_le_bytes());
        }
        cases.push(l_bytes);
        let mut l_plus_1 = l_bytes;
        l_plus_1[0] = l_plus_1[0].wrapping_add(1);
        cases.push(l_plus_1);
        cases.push([0xffu8; 64]);

        for bytes in cases {
            let expected = reduce_wide_naive(&bytes);
            let actual = Scalar::from_bytes_mod_order_wide(&bytes);
            assert_eq!(actual.0, expected.0);
        }
    }

    #[test]
    fn mul_mod_l_matches_naive_reference() {
        Builder::default()
            .with_seed(0)
            .with_search_limit(1024)
            .test(|u| {
                let a: Scalar = u.arbitrary()?;
                let b: Scalar = u.arbitrary()?;

                let wide = super::limbs_mul_wide(&a.0, &b.0);
                let mut bytes = [0u8; 64];
                for (i, limb) in wide.iter().enumerate() {
                    bytes[i * 8..i * 8 + 8].copy_from_slice(&limb.to_le_bytes());
                }
                let expected = reduce_wide_naive(&bytes);

                assert_eq!(a.mul_mod_l(&b).0, expected.0);
                Ok(())
            });
    }

    #[test]
    fn signed_digits_reconstruct_value() {
        const WIDTH: u32 = 6;
        const N: usize = 256usize.div_ceil(WIDTH as usize) + 1;

        Builder::default()
            .with_seed(0)
            .with_search_limit(64)
            .test(|u| {
                let s: Scalar = u.arbitrary()?;
                let digits = s.signed_digits::<N>(WIDTH);

                let base = Scalar::from_u128(1u128 << WIDTH);
                let mut reconstructed = Scalar::ZERO;
                let mut power = Scalar::from_u128(1);
                for &digit in &digits {
                    let magnitude = Scalar::from_u128(digit.unsigned_abs() as u128);
                    let term = power.mul_mod_l(&magnitude);
                    let term = if digit < 0 { term.neg_mod_l() } else { term };
                    reconstructed = reconstructed.add_mod_l(&term);
                    power = power.mul_mod_l(&base);
                }

                assert_eq!(reconstructed.0, s.0);
                Ok(())
            });
    }

    #[test]
    fn signed_digits_are_in_range() {
        const WIDTH: u32 = 6;
        const N: usize = 256usize.div_ceil(WIDTH as usize) + 1;
        let half = 1i32 << (WIDTH - 1);

        Builder::default()
            .with_seed(0)
            .with_search_limit(64)
            .test(|u| {
                let s: Scalar = u.arbitrary()?;
                for digit in s.signed_digits::<N>(WIDTH) {
                    assert!((-half..half).contains(&digit));
                }
                Ok(())
            });
    }

    #[test]
    fn zeroize_clears_scalar() {
        let mut scalar = Scalar([u64::MAX; 4]);
        scalar.zeroize();
        assert_eq!(scalar.0, [0; 4]);
    }
}