arcium-primitives 0.8.1

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

use elliptic_curve::group::{Group, GroupEncoding};
use rand::{
    distributions::{Distribution, Standard},
    RngCore,
};
use serde::{Deserialize, Serialize};
use subtle::{Choice, ConditionallySelectable, ConstantTimeEq};

use crate::{
    algebra::elliptic_curve::{
        curve::{FromCoordinates, PointAtInfinityError, PointCoordinates, ToCoordinates},
        Curve,
        Scalar,
        ScalarAsExtension,
    },
    errors::PrimitiveError,
    random::{CryptoRngCore, Random},
    sharing::unauthenticated::AdditiveShares,
    utils::codec::InPlaceCodec,
};

/// A point on a given curve.
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
#[repr(transparent)]
pub struct Point<C: Curve>(pub(crate) C::Point);

// SAFETY: Point<C> is #[repr(transparent)] over C::Point.
unsafe impl<C: Curve> bytemuck::TransparentWrapper<C::Point> for Point<C> {}

// SAFETY: `write_le_bytes`/`read_le_bytes` use the curve's canonical `GroupEncoding`
// (`to_bytes`/`from_bytes`), which is 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<C: Curve> InPlaceCodec for Point<C> {
    const ENCODED_SIZE: usize = size_of::<<C::Point as GroupEncoding>::Repr>();

    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]) {
        let mut bytes = self.0.to_bytes();
        if C::POINT_BIG_ENDIAN {
            bytes.as_mut().reverse();
        }
        for (slot, &b) in out.iter_mut().zip(bytes.as_ref()) {
            slot.write(b);
        }
    }

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

// Encoded as a fixed-arity tuple (`serialize_tuple`/`deserialize_tuple`), not a `[u8]`/`Vec<u8>`:
// arity is fixed and known to both ends, so formats like `bincode` write no length prefix. Also
// keeps `Point<C>: Serialize` unconditional on `C: Curve` (no bound on
// `<C::Point as GroupEncoding>::Repr`), which downstream code (e.g. `CompressedCircuit<C>`) relies
// on.
impl<C: Curve> Serialize for Point<C> {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        use serde::ser::SerializeTuple;
        let bytes = self.to_inplace_bytes();
        let mut tup = serializer.serialize_tuple(bytes.len())?;
        for b in &bytes {
            tup.serialize_element(b)?;
        }
        tup.end()
    }
}

impl<'de, C: Curve> Deserialize<'de> for Point<C> {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        struct PointVisitor<C>(std::marker::PhantomData<C>);

        impl<'de, C: Curve> serde::de::Visitor<'de> for PointVisitor<C> {
            type Value = Point<C>;

            fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "{} bytes", Point::<C>::ENCODED_SIZE)
            }

            fn visit_seq<A: serde::de::SeqAccess<'de>>(
                self,
                mut seq: A,
            ) -> Result<Self::Value, A::Error> {
                let mut bytes = Vec::with_capacity(Point::<C>::ENCODED_SIZE);
                while let Some(b) = seq.next_element::<u8>()? {
                    bytes.push(b);
                }
                Point::<C>::from_inplace_bytes(&bytes).map_err(serde::de::Error::custom)
            }
        }

        deserializer.deserialize_tuple(
            Point::<C>::ENCODED_SIZE,
            PointVisitor(std::marker::PhantomData),
        )
    }
}

// ------------------------
// | Misc Implementations |
// ------------------------

impl<C: Curve> Point<C> {
    /// The additive identity in the curve group
    pub fn identity() -> Point<C> {
        Point(C::Point::identity())
    }

    pub fn new(point: C::Point) -> Point<C> {
        Point(point)
    }

    /// Check whether the given point is the identity point in the group
    pub fn is_identity(&self) -> Choice {
        self.ct_eq(&Point::identity())
    }

    /// Return the wrapped type
    pub fn inner(&self) -> C::Point {
        self.0
    }

    /// The group generator
    pub fn generator() -> Point<C> {
        Point(<C::Point as Group>::generator())
    }

    pub fn from_coordinates(coordinates: PointCoordinates<C>) -> Option<Point<C>> {
        C::Point::from_coordinates(coordinates).map(Point)
    }

    pub fn to_coordinates(self) -> Result<PointCoordinates<C>, PointAtInfinityError> {
        self.0.to_coordinates()
    }
}

impl<C: Curve> Random for Point<C> {
    #[inline]
    fn random(rng: impl CryptoRngCore) -> Self {
        Point(C::Point::random(rng))
    }
}

impl<C: Curve> Distribution<Point<C>> for Standard {
    #[inline]
    fn sample<R: RngCore + ?Sized>(&self, rng: &mut R) -> Point<C> {
        Point(C::Point::random(rng))
    }
}

// ------------------------------------
// | Curve Arithmetic Implementations |
// ------------------------------------

// === Addition === //

#[macros::op_variants(owned, borrowed, flipped_commutative)]
impl<C: Curve> Add<&Point<C>> for Point<C> {
    type Output = Point<C>;

    #[inline]
    fn add(mut self, rhs: &Point<C>) -> Self::Output {
        self.0 += rhs.0;
        self
    }
}

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

// === Subtraction === //

#[macros::op_variants(owned, borrowed, flipped)]
impl<C: Curve> Sub<&Point<C>> for Point<C> {
    type Output = Point<C>;

    #[inline]
    fn sub(mut self, rhs: &Point<C>) -> Self::Output {
        self.0 -= rhs.0;
        self
    }
}

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

// === Negation === //

#[macros::op_variants(borrowed)]
impl<C: Curve> Neg for Point<C> {
    type Output = Point<C>;

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

// === Scalar Multiplication === //

#[macros::op_variants(owned, borrowed, flipped)]
impl<C: Curve> Mul<&ScalarAsExtension<C>> for Point<C> {
    type Output = Point<C>;

    #[inline]
    fn mul(mut self, rhs: &ScalarAsExtension<C>) -> Self::Output {
        self.0 *= rhs.0;
        self
    }
}

#[macros::op_variants(owned, borrowed, flipped_commutative)]
impl<C: Curve> Mul<&Point<C>> for ScalarAsExtension<C> {
    type Output = Point<C>;

    #[inline]
    fn mul(self, rhs: &Point<C>) -> Self::Output {
        Point(rhs.0 * self.0)
    }
}

#[macros::op_variants(owned, borrowed, flipped)]
impl<C: Curve> Mul<&Scalar<C>> for Point<C> {
    type Output = Point<C>;

    #[inline]
    fn mul(self, rhs: &Scalar<C>) -> Self::Output {
        Point(self.0 * rhs.0)
    }
}

#[macros::op_variants(owned, borrowed, flipped_commutative)]
impl<C: Curve> Mul<&Point<C>> for Scalar<C> {
    type Output = Point<C>;

    #[inline]
    fn mul(self, rhs: &Point<C>) -> Self::Output {
        Point(rhs.0 * self.0)
    }
}

// === MulAssign === //

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

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

// === Equality === //

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

impl<C: Curve> ConditionallySelectable for Point<C> {
    #[inline]
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        let selected = C::Point::conditional_select(&a.0, &b.0, choice);
        Point(selected)
    }
}

// === Other === //

impl<C: Curve> AdditiveShares for Point<C> {}

// === Iterator traits === //

impl<C: Curve> Sum for Point<C> {
    #[inline]
    fn sum<I: Iterator<Item = Point<C>>>(iter: I) -> Self {
        iter.fold(Point::identity(), |acc, x| acc + x)
    }
}

impl<'a, C: Curve> Sum<&'a Point<C>> for Point<C> {
    #[inline]
    fn sum<I: Iterator<Item = &'a Point<C>>>(iter: I) -> Self {
        iter.fold(Point::identity(), |acc, x| acc + x)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{algebra::elliptic_curve::Curve25519Ristretto, utils::bincode_io};

    #[test]
    fn test_point_serialization() {
        let point = Point::<Curve25519Ristretto>::generator();
        let bytes = point.to_inplace_bytes();
        let deserialized_point = Point::<Curve25519Ristretto>::from_inplace_bytes(&bytes).unwrap();
        assert_eq!(point, deserialized_point);

        let bytes = bytes[1..].to_vec(); // Invalid length
        let result = Point::<Curve25519Ristretto>::from_inplace_bytes(&bytes);
        assert!(result.is_err());
    }

    #[test]
    fn test_point_serde_roundtrip() {
        let point = Point::<Curve25519Ristretto>::generator();
        let serialized = bincode_io::serialize(&point).unwrap();
        assert_eq!(serialized, point.to_inplace_bytes(), "no length prefix");
        let deserialized: Point<Curve25519Ristretto> =
            bincode_io::deserialize(&serialized).unwrap();
        assert_eq!(point, deserialized);
    }
}