arcium-primitives 0.8.1

use std::{
    hash::{Hash, Hasher},
    iter::Sum,
    mem::MaybeUninit,
    ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
};

use elliptic_curve::{
    bigint::U384,
    group::{Group, GroupEncoding},
    hash2curve::{ExpandMsgXmd, GroupDigest},
    ops::MulByGenerator,
    sec1::{FromEncodedPoint, ToEncodedPoint},
    FieldBytesEncoding,
};
use ff::Field;
use hybrid_array::Array;
use rand::RngCore;
use sha2::Sha384;
use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};
use typenum::{U3, U48};

use crate::{
    algebra::{
        elliptic_curve::{
            curve::{FromCoordinates, PointAtInfinityError, ToCoordinates},
            BaseFieldElement,
            Curve,
        },
        field::{FieldExtension, SubfieldElement},
    },
    errors::PrimitiveError,
    utils::codec::InPlaceCodec,
};

/// Implements the locally-defined field trait surface shared by both P-384 fields (scalar and
/// base): hashing, byte encodings, uniform sampling, naive wide ops, and `FieldExtension`.
///
/// `$field` is an `ff`-derived prime field over 7 limbs (the derive requires `2 * modulus` to
/// fit the backing representation, so a 384-bit modulus needs 448 bits); `$repr` is the derive's
/// generated 56-byte little-endian repr type.
macro_rules! impl_p384_field {
    ($field:ident, $repr:ident, $name:literal) => {
        impl std::hash::Hash for $field {
            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
                use crate::algebra::field::FieldExtension;
                self.to_le_bytes().into_iter().for_each(|x| {
                    x.hash(state);
                });
            }
        }

        impl From<u128> for $field {
            fn from(value: u128) -> Self {
                let mut bytes = [0u8; 56];
                bytes[..16].copy_from_slice(&value.to_le_bytes());
                <Self as ff::PrimeField>::from_repr($repr(bytes)).unwrap()
            }
        }

        impl crate::algebra::uniform_bytes::FromUniformBytes for $field {
            type UniformBytes = typenum::U64;

            fn from_uniform_bytes(bytes: &hybrid_array::Array<u8, Self::UniformBytes>) -> Self {
                use ff::PrimeField;
                // Split the 512-bit input as lo + 2^256 * hi. Both halves are < 2^256 and the
                // modulus is 384 bits, so all three `from_repr` inputs are canonical.
                let mut lo = [0u8; 56];
                lo[..32].copy_from_slice(&bytes[..32]);
                let mut hi = [0u8; 56];
                hi[..32].copy_from_slice(&bytes[32..]);
                let mut pow = [0u8; 56];
                pow[32] = 1;
                let lo = Self::from_repr($repr(lo)).unwrap();
                let hi = Self::from_repr($repr(hi)).unwrap();
                let two_pow_256 = Self::from_repr($repr(pow)).unwrap();
                hi * two_pow_256 + lo
            }
        }

        impl crate::algebra::field::FieldExtension for $field {
            type Subfield = Self;

            type Degree = typenum::U1;
            type FieldBitSize = typenum::U384;
            type FieldBytesSize = typenum::U48;

            fn to_subfield_elements(&self) -> hybrid_array::Array<Self::Subfield, Self::Degree> {
                hybrid_array::Array([*self])
            }

            fn from_subfield_elements(
                elems: hybrid_array::Array<Self::Subfield, Self::Degree>,
            ) -> Self {
                elems[0]
            }

            fn to_le_bytes(&self) -> hybrid_array::Array<u8, Self::FieldBytesSize> {
                // The derived repr is 56 bytes; the top 8 are always zero for a 384-bit modulus.
                <[u8; 48]>::try_from(&ff::PrimeField::to_repr(self).as_ref()[..48])
                    .unwrap()
                    .into()
            }

            fn from_le_bytes(bytes: &[u8]) -> Option<Self> {
                if bytes.len() != 48 {
                    return None;
                }
                let mut repr = [0u8; 56];
                repr[..48].copy_from_slice(bytes);
                ff::PrimeField::from_repr($repr(repr)).into()
            }

            fn mul_by_subfield(&self, other: &Self::Subfield) -> Self {
                *self * other
            }

            fn generator() -> Self {
                <Self as ff::PrimeField>::MULTIPLICATIVE_GENERATOR
            }
        }

        // SAFETY: `write_le_bytes`/`read_le_bytes` delegate to `FieldExtension::to_le_bytes`/
        // `from_le_bytes` above, a fixed 48-byte, architecture-independent little-endian encoding
        // backed by `ff::PrimeField`'s `to_repr`/`from_repr`. `write_le_bytes` initializes every
        // byte, and `from_le_bytes` range-validates via `from_repr`, so the round-trip is unbiased.
        unsafe impl crate::utils::codec::InPlaceCodec for $field {
            const ENCODED_SIZE: usize = 48;

            fn write_le_bytes(&self, out: &mut [std::mem::MaybeUninit<u8>]) {
                let bytes = crate::algebra::field::FieldExtension::to_le_bytes(self);
                for (slot, &b) in out.iter_mut().zip(bytes.iter()) {
                    slot.write(b);
                }
            }

            fn read_le_bytes(bytes: &[u8]) -> Result<Self, crate::errors::PrimitiveError> {
                <Self as crate::algebra::field::FieldExtension>::from_le_bytes(bytes).ok_or_else(
                    || {
                        crate::errors::PrimitiveError::DeserializationFailed(
                            concat!("non-canonical ", $name, " encoding").into(),
                        )
                    },
                )
            }
        }

        impl crate::random::Random for $field {
            fn random(rng: impl crate::random::CryptoRngCore) -> Self {
                <Self as ff::Field>::random(rng)
            }
        }

        impl crate::types::identifiers::Named for $field {
            fn get_name() -> String {
                $name.to_string()
            }
        }

        // ponytail: eager reduction (WideType = Self); add a lazy-reduction accumulator like
        // BF25519MulAccRepr if P-384 dot-product throughput ever matters.
        impl crate::algebra::ops::IntoWide for $field {
            #[inline]
            fn to_wide(&self) -> Self {
                *self
            }

            #[inline]
            fn zero_wide() -> Self {
                <Self as ff::Field>::ZERO
            }
        }

        impl crate::algebra::ops::ReduceWide for $field {
            #[inline]
            fn reduce_mod_order(a: Self) -> Self {
                a
            }
        }

        impl crate::algebra::ops::MulAccReduce for $field {
            type WideType = Self;

            #[inline]
            fn mul_acc(acc: &mut Self, a: Self, b: Self) {
                *acc += a * b;
            }
        }

        impl<'a> crate::algebra::ops::MulAccReduce<Self, &'a Self> for $field {
            type WideType = Self;

            #[inline]
            fn mul_acc(acc: &mut Self, a: Self, b: &'a Self) {
                *acc += a * b;
            }
        }

        impl<'a> crate::algebra::ops::MulAccReduce<&'a Self, Self> for $field {
            type WideType = Self;

            #[inline]
            fn mul_acc(acc: &mut Self, a: &'a Self, b: Self) {
                *acc += *a * b;
            }
        }

        impl<'a> crate::algebra::ops::MulAccReduce<&'a Self, &'a Self> for $field {
            type WideType = Self;

            #[inline]
            fn mul_acc(acc: &mut Self, a: &'a Self, b: &'a Self) {
                *acc += *a * b;
            }
        }

        impl crate::algebra::ops::AccReduce for $field {
            type WideType = Self;

            #[inline]
            fn acc(acc: &mut Self, a: Self) {
                *acc += a;
            }
        }

        impl<'a> crate::algebra::ops::AccReduce<&'a Self> for $field {
            type WideType = Self;

            #[inline]
            fn acc(acc: &mut Self, a: &'a Self) {
                *acc += a;
            }
        }

        impl crate::algebra::ops::DefaultDotProduct for $field {}
        impl crate::algebra::ops::DefaultDotProduct<Self, &Self> for $field {}
        impl<'a> crate::algebra::ops::DefaultDotProduct<&'a Self, &'a Self> for $field {}
        impl crate::algebra::ops::DefaultDotProduct<&Self, Self> for $field {}
    };
}

pub mod base_field;
pub mod scalar_field;

pub use base_field::BaseFieldP384;
pub use scalar_field::ScalarP384;

/// Marker type for NIST P-384 (secp384r1).
///
/// A local marker (instead of [`::p384::NistP384`]) because the [`Curve`] trait requires `Hash`
/// and the associated types require impls that the orphan rule forbids on foreign types.
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct P384;

impl elliptic_curve::Curve for P384 {
    type FieldBytesSize = U48;
    type Uint = U384;

    const ORDER: U384 = U384::from_be_hex(
        "ffffffffffffffffffffffffffffffffffffffffffffffffc7634d81f4372ddf581a0db248b0a77aecec196accc52973",
    );
}

impl FieldBytesEncoding<P384> for U384 {}

impl Curve for P384 {
    const NAME: &'static str = "NIST P-384";
    const SCALAR_BIG_ENDIAN: bool = false;
    const POINT_BIG_ENDIAN: bool = true;
    const BASE_FIELD_BIG_ENDIAN: bool = false;

    type Point = PointP384;
    type Scalar = ScalarP384;
    type BaseField = BaseFieldP384;

    fn hash_to_curve(bytes: &[u8]) -> Self::Point {
        // RFC 9380 P384_XMD:SHA-384_SSWU_RO_ suite.
        PointP384(
            ::p384::NistP384::hash_from_bytes::<ExpandMsgXmd<Sha384>>(
                &[bytes],
                &[b"P384_XMD:SHA-384_SSWU_RO_"],
            )
            .expect("hash-to-curve with a fixed non-empty DST cannot fail"),
        )
    }
}

/// A NIST P-384 point in projective coordinates, wrapping [`::p384::ProjectivePoint`] to attach
/// the impls the [`Curve`] trait needs but the orphan rule forbids on the foreign type (`Hash`,
/// the `InPlaceCodec` encoding, coordinate conversions, and `Group<Scalar = ScalarP384>`).
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
#[repr(transparent)]
pub struct PointP384(pub ::p384::ProjectivePoint);

impl Hash for PointP384 {
    fn hash<H: Hasher>(&self, state: &mut H) {
        // Canonical compressed encoding: consistent with `Eq` on projective points.
        self.to_bytes().as_slice().hash(state);
    }
}

impl ConstantTimeEq for PointP384 {
    #[inline]
    fn ct_eq(&self, other: &Self) -> Choice {
        self.0.ct_eq(&other.0)
    }
}

impl ConditionallySelectable for PointP384 {
    #[inline]
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        Self(::p384::ProjectivePoint::conditional_select(
            &a.0, &b.0, choice,
        ))
    }
}

#[macros::op_variants(owned)]
impl Add<&PointP384> for PointP384 {
    type Output = PointP384;

    #[inline]
    fn add(mut self, rhs: &PointP384) -> PointP384 {
        self.0 += rhs.0;
        self
    }
}

#[macros::op_variants(owned)]
impl AddAssign<&PointP384> for PointP384 {
    #[inline]
    fn add_assign(&mut self, rhs: &PointP384) {
        self.0 += rhs.0;
    }
}

#[macros::op_variants(owned)]
impl Sub<&PointP384> for PointP384 {
    type Output = PointP384;

    #[inline]
    fn sub(mut self, rhs: &PointP384) -> PointP384 {
        self.0 -= rhs.0;
        self
    }
}

#[macros::op_variants(owned)]
impl SubAssign<&PointP384> for PointP384 {
    #[inline]
    fn sub_assign(&mut self, rhs: &PointP384) {
        self.0 -= rhs.0;
    }
}

impl Neg for PointP384 {
    type Output = PointP384;

    #[inline]
    fn neg(self) -> PointP384 {
        Self(-self.0)
    }
}

#[macros::op_variants(owned)]
impl Mul<&ScalarP384> for PointP384 {
    type Output = PointP384;

    #[inline]
    fn mul(self, rhs: &ScalarP384) -> PointP384 {
        Self(self.0 * rhs.to_p384())
    }
}

#[macros::op_variants(owned)]
impl MulAssign<&ScalarP384> for PointP384 {
    #[inline]
    fn mul_assign(&mut self, rhs: &ScalarP384) {
        self.0 *= rhs.to_p384();
    }
}

impl Sum for PointP384 {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Self::identity(), |acc, x| acc + x)
    }
}

impl<'a> Sum<&'a PointP384> for PointP384 {
    fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
        iter.fold(Self::identity(), |acc, x| acc + x)
    }
}

impl Group for PointP384 {
    type Scalar = ScalarP384;

    fn random(rng: impl RngCore) -> Self {
        Self(::p384::ProjectivePoint::random(rng))
    }

    fn identity() -> Self {
        Self(::p384::ProjectivePoint::IDENTITY)
    }

    fn generator() -> Self {
        Self(::p384::ProjectivePoint::GENERATOR)
    }

    fn is_identity(&self) -> Choice {
        self.0.is_identity()
    }

    fn double(&self) -> Self {
        Self(self.0.double())
    }
}

impl MulByGenerator for PointP384 {}

impl GroupEncoding for PointP384 {
    type Repr = <::p384::ProjectivePoint as GroupEncoding>::Repr;

    fn from_bytes(bytes: &Self::Repr) -> CtOption<Self> {
        ::p384::ProjectivePoint::from_bytes(bytes).map(Self)
    }

    fn from_bytes_unchecked(bytes: &Self::Repr) -> CtOption<Self> {
        ::p384::ProjectivePoint::from_bytes_unchecked(bytes).map(Self)
    }

    fn to_bytes(&self) -> Self::Repr {
        self.0.to_bytes()
    }
}

// SAFETY: `write_le_bytes`/`read_le_bytes` use the curve's canonical `GroupEncoding`
// (`to_bytes`/`from_bytes`), a fixed `ENCODED_SIZE`-byte, architecture-independent encoding.
// `write_le_bytes` initializes every byte, and `read_le_bytes` validates the encoding via
// `from_bytes`, so the round-trip is unbiased.
unsafe impl InPlaceCodec for PointP384 {
    const ENCODED_SIZE: usize = size_of::<<Self as GroupEncoding>::Repr>();

    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]) {
        let bytes = self.to_bytes();
        for (slot, &b) in out.iter_mut().zip(bytes.as_slice()) {
            slot.write(b);
        }
    }

    fn read_le_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError> {
        let mut repr = <Self as GroupEncoding>::Repr::default();
        if bytes.len() != repr.len() {
            return Err(PrimitiveError::InvalidSize(repr.len(), bytes.len()));
        }
        repr.as_mut_slice().copy_from_slice(bytes);
        Option::from(Self::from_bytes(&repr)).ok_or_else(|| {
            PrimitiveError::DeserializationFailed("invalid curve point encoding".into())
        })
    }
}

fn base_field_to_be_bytes(element: BaseFieldP384) -> ::p384::FieldBytes {
    let mut bytes: [u8; 48] = element.to_le_bytes().into();
    bytes.reverse();
    bytes.into()
}

fn base_field_from_be_bytes(bytes: &[u8]) -> Option<BaseFieldP384> {
    let mut bytes: [u8; 48] = bytes.try_into().ok()?;
    bytes.reverse();
    BaseFieldP384::from_le_bytes(&bytes)
}

impl ToCoordinates for PointP384 {
    type BaseFieldElement = BaseFieldElement<P384>;
    type NumCoordinates = U3;

    /// Returns projective Weierstrass coordinates [X, Y, Z], dehomogenized to [x, y, 1] so that
    /// all nodes obtain identical coordinates for MAC checks. The identity maps to (0 : 1 : 0).
    fn to_coordinates(self) -> Result<Array<Self::BaseFieldElement, U3>, PointAtInfinityError> {
        if self.0.is_identity().into() {
            return Ok(Array(
                [BaseFieldP384::ZERO, BaseFieldP384::ONE, BaseFieldP384::ZERO]
                    .map(SubfieldElement::new),
            ));
        }
        let encoded = self.0.to_affine().to_encoded_point(false);
        let x = base_field_from_be_bytes(encoded.x().unwrap()).unwrap();
        let y = base_field_from_be_bytes(encoded.y().unwrap()).unwrap();
        Ok(Array([x, y, BaseFieldP384::ONE].map(SubfieldElement::new)))
    }
}

impl FromCoordinates for PointP384 {
    type BaseFieldElement = BaseFieldElement<P384>;
    type NumCoordinates = U3;

    /// Inverse of [`ToCoordinates`]: accepts any projective representative, dehomogenizes by Z,
    /// and validates the Weierstrass curve equation (via the SEC1 uncompressed decoder).
    /// For Z = 0 the only projective point on the curve is the identity (0 : Y : 0) with Y != 0.
    /// Returns `None` for points not on the curve.
    #[allow(non_snake_case)]
    fn from_coordinates(coordinates: Array<Self::BaseFieldElement, U3>) -> Option<Self> {
        let [X, Y, Z] = coordinates.0.map(|coordinate| coordinate.inner());
        let Some(z_inv) = Option::<BaseFieldP384>::from(Z.invert()) else {
            let is_identity = bool::from(X.is_zero()) && !bool::from(Y.is_zero());
            return is_identity.then_some(Self(::p384::ProjectivePoint::IDENTITY));
        };
        let x = X * z_inv;
        let y = Y * z_inv;
        let encoded = ::p384::EncodedPoint::from_affine_coordinates(
            &base_field_to_be_bytes(x),
            &base_field_to_be_bytes(y),
            false,
        );
        let affine =
            Option::<::p384::AffinePoint>::from(::p384::AffinePoint::from_encoded_point(&encoded))?;
        Some(Self(affine.into()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::random::{test_rng, Random};

    fn random_scalar(rng: impl crate::random::CryptoRngCore) -> ScalarP384 {
        Random::random(rng)
    }

    #[test]
    fn test_group_law_and_encoding_roundtrip() {
        let mut rng = test_rng();
        let point = PointP384::random(&mut rng);
        assert_eq!(point + point, point.double());

        let bytes = point.to_bytes();
        let decoded = Option::<PointP384>::from(PointP384::from_bytes(&bytes)).unwrap();
        assert_eq!(point, decoded);

        // Invalid encodings must be rejected.
        let mut bad = bytes;
        bad.as_mut_slice()[1..].fill(0xFF);
        assert!(Option::<PointP384>::from(PointP384::from_bytes(&bad)).is_none());
    }

    #[test]
    fn test_scalar_mul_matches_p384_crate() {
        let mut rng = test_rng();
        for _ in 0..10 {
            let a = random_scalar(&mut rng);
            let b = random_scalar(&mut rng);
            assert_eq!((a * b).to_p384(), a.to_p384() * b.to_p384());
            assert_eq!((a + b).to_p384(), a.to_p384() + b.to_p384());
        }

        let scalar = random_scalar(&mut rng);
        let expected = ::p384::ProjectivePoint::GENERATOR * scalar.to_p384();
        assert_eq!(PointP384::generator() * scalar, PointP384(expected));
    }

    #[test]
    fn test_coordinates_roundtrip() {
        let mut rng = test_rng();
        let point = PointP384::random(&mut rng);
        let coordinates = point.to_coordinates().unwrap();
        assert_eq!(coordinates[2], SubfieldElement::new(BaseFieldP384::ONE));
        assert_eq!(PointP384::from_coordinates(coordinates).unwrap(), point);

        // Any projective representative of the same point is accepted.
        let lambda = SubfieldElement::new(BaseFieldP384::from(7u128));
        let scaled = Array([
            coordinates[0] * lambda,
            coordinates[1] * lambda,
            coordinates[2] * lambda,
        ]);
        assert_eq!(PointP384::from_coordinates(scaled).unwrap(), point);

        // The identity is (0 : 1 : 0) and round-trips.
        let identity_coordinates = PointP384::identity().to_coordinates().unwrap();
        assert_eq!(
            PointP384::from_coordinates(identity_coordinates).unwrap(),
            PointP384::identity()
        );

        // Z = 0 with X != 0 is not on the curve.
        let mut bad_infinity = coordinates;
        bad_infinity[2] = SubfieldElement::new(BaseFieldP384::ZERO);
        assert!(PointP384::from_coordinates(bad_infinity).is_none());

        // On-curve check: (x, y) not on the curve is rejected.
        let mut off_curve = coordinates;
        off_curve[1] += SubfieldElement::new(BaseFieldP384::ONE);
        assert!(PointP384::from_coordinates(off_curve).is_none());
    }

    #[test]
    fn test_hash_to_curve() {
        let point1 = P384::hash_to_curve(b"async-mpc test input");
        let point2 = P384::hash_to_curve(b"async-mpc test input");
        let point3 = P384::hash_to_curve(b"different input");
        assert_eq!(point1, point2);
        assert_ne!(point1, point3);
        assert!(!bool::from(point1.is_identity()));
    }

    #[test]
    fn test_inplace_roundtrip() {
        let mut rng = test_rng();
        let point = PointP384::random(&mut rng);
        let serialized = point.to_inplace_bytes();
        assert_eq!(PointP384::from_inplace_bytes(&serialized).unwrap(), point);
    }
}