arcium-primitives 0.4.2

use std::{
    iter::Sum as IterSum,
    mem::MaybeUninit,
    ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
    sync::Arc,
};

use serde::{Deserialize, Serialize};
use subtle::{Choice, ConstantTimeEq};
use typenum::{Prod, Sum, U1, U2, U3};
use wincode::{
    io::{Reader, Writer},
    ReadResult,
    SchemaRead,
    SchemaWrite,
    WriteResult,
};

use crate::{
    algebra::{
        elliptic_curve::{BaseField, Curve, Point, ScalarAsExtension, ScalarField},
        field::{FieldElement, FieldExtension},
    },
    random::{CryptoRngCore, Random, RandomWith},
    types::{
        heap_array::{CurvePoints, FieldElements},
        ConditionallySelectable,
        HeapArray,
        Positive,
    },
};

// ============================================================
// |                     GlobalKey                             |
// ============================================================

/// α, a global authentication key for field shares. Each party holds a
/// global key α for each peer, and uses that α to authenticate all its field shares
/// (alongside a local key β).
pub type GlobalKey<A> = Arc<A>;

/// Global authentication key for field shares. Alias for [`GlobalKey<FieldElement<F>>`].
pub type GlobalFieldKey<F> = Arc<FieldElement<F>>;
/// [`GlobalFieldKey`] for scalar field shares. Alias for [`GlobalFieldKey<ScalarField<C>>`].
pub type GlobalScalarKey<C> = GlobalFieldKey<ScalarField<C>>;
/// [`GlobalFieldKey`] for base field shares. Alias for [`GlobalFieldKey<BaseField<C>>`].
pub type GlobalBaseKey<C> = GlobalFieldKey<BaseField<C>>;
/// [`GlobalFieldKey`] for curve-point shares (MAC is computed over the scalar field).
/// Alias for [`GlobalFieldKey<ScalarField<C>>`].
pub type GlobalCurveKey<C> = GlobalFieldKey<ScalarField<C>>;

// ============================================================
// |               FieldShareKeyBase<A, B>                    |
// ============================================================

/// Generic authenticated key: a global key `alpha: A` and a local key `beta: B`,
/// satisfying `MAC(x) = α · x + β`.
///
/// This is the base type for both:
/// - [`FieldShareKey<F>`]    — single-value authenticated key
/// - [`FieldShareKeys<F,M>`] — batched authenticated key (M values, shared α)
// α and β, such that MAC(x) = α · x + β
// In the context of VOLE, this corresponds to w = Δ · u + v
#[derive(Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[repr(C)]
pub struct PairwiseAuthKey<A, B> {
    pub alpha: GlobalKey<A>, // α, global key
    pub beta: B,             // β, local key (single value or batch)
}

impl<A: std::fmt::Debug, B: std::fmt::Debug> std::fmt::Debug for PairwiseAuthKey<A, B> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PairwiseAuthKey")
            .field("alpha", &self.alpha)
            .field("beta", &self.beta)
            .finish()
    }
}

impl<A, B> PairwiseAuthKey<A, B> {
    /// Create a new [`PairwiseAuthKey`] with the given global key `alpha` and local key `beta`.
    #[inline]
    pub fn new(alpha: Arc<A>, beta: B) -> Self {
        PairwiseAuthKey { alpha, beta }
    }

    /// Return a reference to the local key `beta`.
    #[inline]
    pub fn get_beta(&self) -> &B {
        &self.beta
    }

    /// Return a reference to the global key `alpha`.
    #[inline]
    pub fn get_alpha(&self) -> &A {
        &self.alpha
    }

    /// Return the value of the global key `alpha`.
    #[inline]
    pub fn alpha(&self) -> Arc<A> {
        self.alpha.clone()
    }
}

// --- Aliases --- //

/// [`PairwiseAuthKey`] for a single field value.
///
/// See also [`FieldShareKeys<F, M>`] for the batched variant.
///
/// α and β, such that MAC(x) = α · x + β
pub type FieldKey<F> = PairwiseAuthKey<FieldElement<F>, FieldElement<F>>;

/// [`PairwiseAuthKey`] for a single scalar field element. Alias for
/// [`FieldKey<ScalarField<C>>`].
pub type ScalarKey<C> = FieldKey<ScalarField<C>>;
/// [`PairwiseAuthKey`] for a single base field element. Alias for [`FieldKey<BaseField<C>>`].
pub type BaseFieldKey<C> = FieldKey<BaseField<C>>;

/// [`PairwiseAuthKey`] for a batch of `M` field values sharing a single global key.
///
/// See also [`FieldKey<F>`] for the single-value variant.
///
/// α and β₁…βₘ, such that MAC(xₜ) = α · xₜ + βₜ
pub type FieldKeys<F, M> = PairwiseAuthKey<FieldElement<F>, FieldElements<F, M>>;

/// [`PairwiseAuthKey`] for a batch of scalar field elements.
pub type ScalarKeys<C, M> = FieldKeys<ScalarField<C>, M>;
/// [`PairwiseAuthKey`] for a batch of base field elements.
pub type BaseFieldKeys<C, M> = FieldKeys<BaseField<C>, M>;

/// [`PairwiseAuthKey`] for a single curve point.
pub type CurveKey<C> = PairwiseAuthKey<ScalarAsExtension<C>, Point<C>>;
/// [`PairwiseAuthKey`] for `M` curve points.
pub type CurveKeys<C, M> = PairwiseAuthKey<ScalarAsExtension<C>, CurvePoints<C, M>>;

// ========================
// |       Addition       |
// ========================

#[macros::op_variants(owned, borrowed, flipped)]
impl<A: PartialEq, B> Add<&PairwiseAuthKey<A, B>> for PairwiseAuthKey<A, B>
where
    for<'a> B: Add<&'a B, Output = B>,
{
    type Output = PairwiseAuthKey<A, B>;

    #[inline]
    fn add(mut self, other: &PairwiseAuthKey<A, B>) -> Self::Output {
        assert!(self.alpha == other.alpha, "alpha mismatch");
        self.beta = self.beta + &other.beta;
        self
    }
}

#[macros::op_variants(owned)]
impl<'a, A: PartialEq, B> AddAssign<&'a PairwiseAuthKey<A, B>> for PairwiseAuthKey<A, B>
where
    for<'b> B: AddAssign<&'b B>,
{
    #[inline]
    fn add_assign(&mut self, other: &'a PairwiseAuthKey<A, B>) {
        assert!(self.alpha == other.alpha, "alpha mismatch");
        self.beta += &other.beta;
    }
}

#[macros::op_variants(owned, borrowed, flipped)]
impl<A: PartialEq, B> Sub<&PairwiseAuthKey<A, B>> for PairwiseAuthKey<A, B>
where
    for<'a> B: Sub<&'a B, Output = B>,
{
    type Output = PairwiseAuthKey<A, B>;

    #[inline]
    fn sub(self, other: &PairwiseAuthKey<A, B>) -> Self::Output {
        assert!(self.alpha == other.alpha, "alpha mismatch");
        PairwiseAuthKey {
            alpha: self.alpha,
            beta: self.beta - &other.beta,
        }
    }
}

#[macros::op_variants(owned)]
impl<'a, A: PartialEq, B> SubAssign<&'a PairwiseAuthKey<A, B>> for PairwiseAuthKey<A, B>
where
    for<'b> B: SubAssign<&'b B>,
{
    #[inline]
    fn sub_assign(&mut self, other: &'a PairwiseAuthKey<A, B>) {
        assert!(self.alpha == other.alpha, "alpha mismatch");
        self.beta -= &other.beta;
    }
}

#[macros::op_variants(borrowed)]
impl<A, B: Neg<Output = B>> Neg for PairwiseAuthKey<A, B> {
    type Output = PairwiseAuthKey<A, B>;

    #[inline]
    fn neg(mut self) -> Self::Output {
        self.beta = -self.beta;
        self
    }
}

impl<A: Default + PartialEq, B: Default> IterSum for PairwiseAuthKey<A, B>
where
    for<'a> B: Add<&'a B, Output = B>,
{
    fn sum<I: Iterator<Item = Self>>(mut iter: I) -> Self {
        let first = iter.next().unwrap_or_default();
        iter.fold(first, |acc, x| acc + &x)
    }
}

// ========================
// |   Multiplication     |
// ========================

impl<A, B, RHS, BPrime> Mul<RHS> for PairwiseAuthKey<A, B>
where
    B: Mul<RHS, Output = BPrime>,
{
    type Output = PairwiseAuthKey<A, BPrime>;
    #[inline]
    fn mul(self, other: RHS) -> Self::Output {
        PairwiseAuthKey {
            alpha: self.alpha,
            beta: self.beta * other,
        }
    }
}

impl<A: Clone, B, RHS, BPrime> Mul<RHS> for &PairwiseAuthKey<A, B>
where
    for<'b> &'b B: Mul<RHS, Output = BPrime>,
{
    type Output = PairwiseAuthKey<A, BPrime>;
    #[inline]
    fn mul(self, other: RHS) -> Self::Output {
        PairwiseAuthKey {
            alpha: self.alpha.clone(),
            beta: &self.beta * other,
        }
    }
}

impl<A, B, RHS> MulAssign<RHS> for PairwiseAuthKey<A, B>
where
    B: MulAssign<RHS>,
{
    #[inline]
    fn mul_assign(&mut self, other: RHS) {
        self.beta *= other;
    }
}

// ========================
// |   Random Generation  |
// ========================

impl<A: Random, B: Random> Random for PairwiseAuthKey<A, B> {
    fn random(mut rng: impl CryptoRngCore) -> Self {
        PairwiseAuthKey {
            alpha: A::random(&mut rng).into(),
            beta: B::random(&mut rng),
        }
    }
}

impl<A: Clone, B: Random> RandomWith<GlobalKey<A>> for PairwiseAuthKey<A, B> {
    fn random_with(mut rng: impl CryptoRngCore, alpha: GlobalKey<A>) -> Self {
        PairwiseAuthKey {
            alpha,
            beta: B::random(&mut rng),
        }
    }
}

// ---------------------------------------
// |  Constant time Selection / Equality |
// ---------------------------------------

impl<A: ConstantTimeEq, B: ConstantTimeEq> ConstantTimeEq for PairwiseAuthKey<A, B> {
    #[inline]
    fn ct_eq(&self, other: &Self) -> Choice {
        self.alpha.ct_eq(&other.alpha) & self.beta.ct_eq(&other.beta)
    }
}

impl<A: ConditionallySelectable, B: ConditionallySelectable> ConditionallySelectable
    for PairwiseAuthKey<A, B>
{
    #[inline]
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        PairwiseAuthKey {
            alpha: A::conditional_select(&a.alpha, &b.alpha, choice).into(),
            beta: B::conditional_select(&a.beta, &b.beta, choice),
        }
    }
}

// ========================
// |  SchemaWrite / Read  |
// ========================

impl<A: SchemaWrite<Src = A>, B: SchemaWrite<Src = B>> SchemaWrite for PairwiseAuthKey<A, B> {
    type Src = Self;

    fn size_of(src: &Self::Src) -> WriteResult<usize> {
        Ok(A::size_of(&src.alpha)? + B::size_of(&src.beta)?)
    }

    fn write(writer: &mut impl Writer, src: &Self::Src) -> WriteResult<()> {
        A::write(writer, &src.alpha)?;
        B::write(writer, &src.beta)
    }
}

impl<'de, A: SchemaRead<'de, Dst = A>, B: SchemaRead<'de, Dst = B>> SchemaRead<'de>
    for PairwiseAuthKey<A, B>
{
    type Dst = Self;

    fn read(reader: &mut impl Reader<'de>, dst: &mut MaybeUninit<Self::Dst>) -> ReadResult<()> {
        let mut alpha_uninit = MaybeUninit::<A>::uninit();
        let mut beta_uninit = MaybeUninit::<B>::uninit();
        A::read(reader, &mut alpha_uninit)?;
        B::read(reader, &mut beta_uninit)?;
        // SAFETY: both fields were initialized by the reads above
        let alpha = unsafe { alpha_uninit.assume_init() }.into();
        let beta = unsafe { beta_uninit.assume_init() };
        dst.write(PairwiseAuthKey { alpha, beta });
        Ok(())
    }
}

// -----------------------
// |   Split and Merge   |
// -----------------------

/// Generic split and merge operations on any batched key `PairwiseAuthKey<A, HeapArray<B, M>>`.
/// Covers both `FieldShareKeys<F, M>` and `CurveKeys<C, M>`.
impl<A: Clone, B: Copy, M: Positive> PairwiseAuthKey<A, HeapArray<B, M>> {
    #[allow(clippy::type_complexity)]
    /// Split a batched key into two smaller batched keys with the same global key `alpha` and
    /// disjoint local keys `beta` (first `M1` values in the first key, next `M2` values in the
    /// second key).
    pub fn split<M1, M2>(
        self,
    ) -> (
        PairwiseAuthKey<A, HeapArray<B, M1>>,
        PairwiseAuthKey<A, HeapArray<B, M2>>,
    )
    where
        M1: Positive,
        M2: Positive + Add<M1, Output = M>,
    {
        let PairwiseAuthKey { alpha, beta } = self;
        let (betas1, betas2) = beta.split::<M1, M2>();
        (
            PairwiseAuthKey::new(alpha.clone(), betas1),
            PairwiseAuthKey::new(alpha, betas2),
        )
    }
    #[allow(clippy::type_complexity)]
    /// Split a batched key into two smaller batched keys with the same global key `alpha` and
    /// disjoint local keys `beta` (first `M/2` values in the first key, next `M/2` values in the
    /// second key).
    pub fn split_halves<MDiv2>(
        self,
    ) -> (
        PairwiseAuthKey<A, HeapArray<B, MDiv2>>,
        PairwiseAuthKey<A, HeapArray<B, MDiv2>>,
    )
    where
        MDiv2: Positive + Mul<U2, Output = M>,
    {
        let PairwiseAuthKey { alpha, beta } = self;
        let (betas1, betas2) = beta.split_halves::<MDiv2>();
        (
            PairwiseAuthKey::new(alpha.clone(), betas1),
            PairwiseAuthKey::new(alpha, betas2),
        )
    }

    /// Merge two batched keys with the same global key `alpha` and disjoint local keys `beta`
    /// (first `M/2` values in the first key, next `M/2` values in the second key) into a single
    /// batched key with `M` values in the local key.
    pub fn merge_halves(this: Self, other: Self) -> PairwiseAuthKey<A, HeapArray<B, Prod<M, U2>>>
    where
        M: Positive + Mul<U2, Output: Positive>,
        A: PartialEq,
    {
        assert!(this.alpha == other.alpha, "alpha mismatch in merge_halves");
        PairwiseAuthKey {
            alpha: this.alpha,
            beta: HeapArray::merge_halves(this.beta, other.beta),
        }
    }

    #[allow(clippy::type_complexity)]
    /// Split a batched key into three smaller batched keys with the same global key `alpha` and
    /// disjoint local keys `beta` (first `M1` values in the first key, next `M2` values in the
    /// second key, last `M3` values in the third key).
    pub fn split3<M1, M2, M3>(
        self,
    ) -> (
        PairwiseAuthKey<A, HeapArray<B, M1>>,
        PairwiseAuthKey<A, HeapArray<B, M2>>,
        PairwiseAuthKey<A, HeapArray<B, M3>>,
    )
    where
        M1: Positive,
        M2: Positive + Add<M1>,
        M3: Positive + Add<Sum<M2, M1>, Output = M>,
    {
        let PairwiseAuthKey { alpha, beta } = self;
        let (betas1, betas2, betas3) = beta.split3::<M1, M2, M3>();
        (
            PairwiseAuthKey::new(alpha.clone(), betas1),
            PairwiseAuthKey::new(alpha.clone(), betas2),
            PairwiseAuthKey::new(alpha, betas3),
        )
    }

    #[allow(clippy::type_complexity)]
    /// Split a batched key into three smaller batched keys with the same global key `alpha` and
    /// disjoint local keys `beta` (first `M/3` values in the first key, next `M/3` values in the
    /// second key, last `M/3` values in the third key).
    pub fn split_thirds<MDiv3>(
        self,
    ) -> (
        PairwiseAuthKey<A, HeapArray<B, MDiv3>>,
        PairwiseAuthKey<A, HeapArray<B, MDiv3>>,
        PairwiseAuthKey<A, HeapArray<B, MDiv3>>,
    )
    where
        MDiv3: Positive + Mul<U3, Output = M>,
    {
        let PairwiseAuthKey { alpha, beta } = self;
        let (betas1, betas2, betas3) = beta.split_thirds::<MDiv3>();
        (
            PairwiseAuthKey::new(alpha.clone(), betas1),
            PairwiseAuthKey::new(alpha.clone(), betas2),
            PairwiseAuthKey::new(alpha, betas3),
        )
    }

    /// Merge three batched keys with the same global key `alpha` and disjoint local keys `beta`
    /// (first `M/3` values in the first key, next `M/3` values in the second key, last `M/3` values
    /// in the third key) into a single batched key with `M` values in the local key.
    pub fn merge_thirds(
        first: Self,
        second: Self,
        third: Self,
    ) -> PairwiseAuthKey<A, HeapArray<B, Prod<M, U3>>>
    where
        M: Positive + Mul<U3, Output: Positive>,
        A: PartialEq,
    {
        assert!(
            first.alpha == second.alpha,
            "alpha mismatch in merge_thirds"
        );
        assert!(first.alpha == third.alpha, "alpha mismatch in merge_thirds");
        PairwiseAuthKey {
            alpha: first.alpha,
            beta: HeapArray::merge_thirds(first.beta, second.beta, third.beta),
        }
    }
}

// ------------------------
// |   Iterate and Cast   |
// ------------------------

/// Convert a single-element key into a one-element batched key.
/// Covers `From<FieldShareKey<F>> for FieldShareKeys<F, U1>` and
/// `From<CurveKey<C>> for CurveKeys<C, U1>`.
impl<A, T> From<PairwiseAuthKey<A, T>> for PairwiseAuthKey<A, HeapArray<T, U1>> {
    fn from(key: PairwiseAuthKey<A, T>) -> Self {
        PairwiseAuthKey {
            alpha: key.alpha,
            beta: HeapArray::from(key.beta),
        }
    }
}

pub struct FieldShareKeysIterator<F: FieldExtension, M: Positive> {
    keys: FieldKeys<F, M>,
    index: usize,
}

impl<F: FieldExtension, M: Positive> Iterator for FieldShareKeysIterator<F, M> {
    type Item = FieldKey<F>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index < M::to_usize() {
            let key = PairwiseAuthKey {
                alpha: self.keys.alpha.clone(),
                beta: self.keys.beta[self.index],
            };
            self.index += 1;
            Some(key)
        } else {
            None
        }
    }
}

impl<F: FieldExtension, M: Positive> ExactSizeIterator for FieldShareKeysIterator<F, M> {
    fn len(&self) -> usize {
        M::to_usize()
    }
}

impl<F: FieldExtension, M: Positive> IntoIterator for FieldKeys<F, M> {
    type Item = FieldKey<F>;
    type IntoIter = FieldShareKeysIterator<F, M>;

    fn into_iter(self) -> Self::IntoIter {
        FieldShareKeysIterator::<F, M> {
            keys: self,
            index: 0,
        }
    }
}

pub struct FieldShareKeyRef<'a, F>
where
    F: FieldExtension,
{
    pub alpha: GlobalFieldKey<F>,
    pub beta: &'a FieldElement<F>,
}

impl<'a, F: FieldExtension> From<FieldShareKeyRef<'a, F>> for FieldKey<F> {
    fn from(val: FieldShareKeyRef<'a, F>) -> Self {
        PairwiseAuthKey {
            alpha: val.alpha,
            beta: *val.beta,
        }
    }
}

pub struct FieldShareKeysRefIterator<'a, F, M>
where
    F: FieldExtension,
    M: Positive,
{
    keys: &'a FieldKeys<F, M>,
    index: usize,
}

impl<F: FieldExtension, M: Positive> ExactSizeIterator for FieldShareKeysRefIterator<'_, F, M> {
    fn len(&self) -> usize {
        M::to_usize()
    }
}

impl<'a, F: FieldExtension, M: Positive> Iterator for FieldShareKeysRefIterator<'a, F, M> {
    type Item = FieldShareKeyRef<'a, F>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index < M::to_usize() {
            let key = FieldShareKeyRef {
                alpha: self.keys.alpha.clone(),
                beta: &self.keys.beta[self.index],
            };
            self.index += 1;
            Some(key)
        } else {
            None
        }
    }
}

impl<'a, F: FieldExtension, M: Positive> IntoIterator for &'a FieldKeys<F, M> {
    type Item = FieldShareKeyRef<'a, F>;
    type IntoIter = FieldShareKeysRefIterator<'a, F, M>;

    fn into_iter(self) -> Self::IntoIter {
        FieldShareKeysRefIterator {
            keys: self,
            index: 0,
        }
    }
}

// --- Type conversions --- //

impl<C: Curve> From<ScalarKey<C>> for CurveKey<C> {
    #[inline]
    fn from(scalar_key: ScalarKey<C>) -> Self {
        PairwiseAuthKey {
            alpha: scalar_key.alpha,
            beta: scalar_key.beta * Point::<C>::generator(),
        }
    }
}

impl<C: Curve, M: Positive> From<ScalarKeys<C, M>> for CurveKeys<C, M> {
    #[inline]
    fn from(scalar_key: ScalarKeys<C, M>) -> Self {
        PairwiseAuthKey {
            alpha: scalar_key.alpha,
            beta: scalar_key.beta * &Point::<C>::generator(),
        }
    }
}

#[cfg(test)]
mod tests {
    use typenum::{U10, U8};

    use super::*;
    use crate::algebra::elliptic_curve::{Curve25519Ristretto as C, ScalarAsExtension};

    type Fq = ScalarAsExtension<C>;

    #[test]
    fn test_addition() {
        let alpha = GlobalFieldKey::new(Fq::from(3u32));
        let key1 = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: Fq::from(10u32),
        };
        let key2 = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: Fq::from(7u32),
        };
        let expected_result = PairwiseAuthKey {
            alpha,
            beta: Fq::from(17u32),
        };
        assert_eq!(key1 + key2, expected_result);
    }

    #[test]
    fn test_subtraction() {
        let alpha = GlobalFieldKey::new(Fq::from(3u32));
        let key1 = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: Fq::from(10u32),
        };
        let key2 = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: Fq::from(7u32),
        };
        let expected_result = PairwiseAuthKey {
            alpha,
            beta: Fq::from(3u32),
        };
        assert_eq!(key1 - key2, expected_result);
    }

    #[test]
    fn test_multiplication() {
        let alpha = GlobalFieldKey::new(Fq::from(3u32));
        let key = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: Fq::from(10u32),
        };
        let scalar = Fq::from(3u32);
        let expected_result = PairwiseAuthKey {
            alpha,
            beta: Fq::from(30u32),
        };
        assert_eq!(key * scalar, expected_result);
    }

    #[test]
    fn test_negation() {
        let key = PairwiseAuthKey {
            alpha: GlobalFieldKey::new(Fq::from(5u32)),
            beta: Fq::from(10u32),
        };
        let expected_result = PairwiseAuthKey {
            alpha: GlobalFieldKey::new(Fq::from(5u32)),
            beta: -Fq::from(10u32),
        };
        assert_eq!(-key, expected_result);
    }

    #[test]
    fn test_batched_addition() {
        let alpha = GlobalFieldKey::new(Fq::from(3u32));
        let key1 = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: HeapArray::<_, U10>::from_fn(|_| Fq::from(10u32)),
        };
        let key2 = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: HeapArray::<_, U10>::from_fn(|_| Fq::from(7u32)),
        };
        let expected_result = PairwiseAuthKey {
            alpha,
            beta: HeapArray::<_, U10>::from_fn(|_| Fq::from(17u32)),
        };
        assert_eq!(key1 + key2, expected_result);
    }

    #[test]
    fn test_batched_subtraction() {
        let alpha = GlobalFieldKey::new(Fq::from(3u32));
        let key1 = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: HeapArray::<_, U10>::from_fn(|_| Fq::from(10u32)),
        };
        let key2 = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: HeapArray::<_, U10>::from_fn(|_| Fq::from(7u32)),
        };
        let expected_result = PairwiseAuthKey {
            alpha,
            beta: HeapArray::<_, U10>::from_fn(|_| Fq::from(3u32)),
        };
        assert_eq!(key1 - key2, expected_result);
    }

    #[test]
    fn test_batched_multiplication() {
        let alpha = GlobalFieldKey::new(Fq::from(3u32));
        let key = PairwiseAuthKey {
            alpha: alpha.clone(),
            beta: HeapArray::<_, U10>::from_fn(|_| Fq::from(10u32)),
        };
        let scalar = Fq::from(3u32);
        let expected_result = PairwiseAuthKey {
            alpha,
            beta: HeapArray::<_, U10>::from_fn(|_| Fq::from(30u32)),
        };
        assert_eq!(key * &scalar, expected_result);
    }

    #[test]
    fn test_batched_negation() {
        let key = PairwiseAuthKey {
            alpha: GlobalFieldKey::new(Fq::from(5u32)),
            beta: HeapArray::<_, U8>::from_fn(|_| Fq::from(10u32)),
        };
        let expected_result = PairwiseAuthKey {
            alpha: GlobalFieldKey::new(Fq::from(5u32)),
            beta: -HeapArray::<_, U8>::from_fn(|_| Fq::from(10u32)),
        };
        assert_eq!(-key, expected_result);
    }

    // --- CurveKey (single point) tests --- //

    mod curve_single {
        use super::*;
        use crate::algebra::elliptic_curve::Curve25519Ristretto as C;

        type P = Point<C>;
        type FrExt = ScalarAsExtension<C>;

        #[test]
        fn test_addition() {
            let mut rng = rand::thread_rng();
            let alpha = GlobalCurveKey::<C>::new(FrExt::random(&mut rng));
            let beta1 = P::random(&mut rng);
            let beta2 = P::random(&mut rng);
            let key1 = CurveKey {
                alpha: alpha.clone(),
                beta: beta1,
            };
            let key2 = CurveKey {
                alpha: alpha.clone(),
                beta: beta2,
            };
            let expected = CurveKey {
                alpha,
                beta: beta1 + beta2,
            };
            assert_eq!(key1 + key2, expected);
        }

        #[test]
        fn test_subtraction() {
            let mut rng = rand::thread_rng();
            let alpha = GlobalCurveKey::<C>::new(FrExt::random(&mut rng));
            let beta1 = P::random(&mut rng);
            let beta2 = P::random(&mut rng);
            let key1 = CurveKey {
                alpha: alpha.clone(),
                beta: beta1,
            };
            let key2 = CurveKey {
                alpha: alpha.clone(),
                beta: beta2,
            };
            let expected = CurveKey {
                alpha,
                beta: beta1 - beta2,
            };
            assert_eq!(key1 - key2, expected);
        }

        #[test]
        fn test_multiplication() {
            let mut rng = rand::thread_rng();
            let alpha = GlobalCurveKey::<C>::new(FrExt::random(&mut rng));
            let beta1 = P::random(&mut rng);
            let key = CurveKey {
                alpha: alpha.clone(),
                beta: beta1,
            };
            let scalar = FrExt::from(3u32);
            let expected = CurveKey {
                alpha,
                beta: beta1 * scalar,
            };
            assert_eq!(key * scalar, expected);
        }

        #[test]
        fn test_negation() {
            let mut rng = rand::thread_rng();
            let alpha = GlobalCurveKey::<C>::new(FrExt::random(&mut rng));
            let beta1 = P::random(&mut rng);
            let key = CurveKey {
                alpha: alpha.clone(),
                beta: beta1,
            };
            let expected = CurveKey {
                alpha,
                beta: -beta1,
            };
            assert_eq!(-key, expected);
        }
    }

    // --- CurveKeys (batched) tests --- //

    mod curve_batched {
        use typenum::U8;

        use super::*;
        use crate::{
            algebra::elliptic_curve::Curve25519Ristretto as C,
            random::{self, Random},
            types::heap_array::CurvePoints,
        };

        type FrExt = ScalarAsExtension<C>;
        type Ps = CurvePoints<C, U8>;

        #[test]
        fn test_addition() {
            let mut rng = random::test_rng();
            let alpha = GlobalCurveKey::<C>::random(&mut rng);
            let beta1 = Ps::random(&mut rng);
            let beta2 = Ps::random(&mut rng);
            let key1 = CurveKeys {
                alpha: alpha.clone(),
                beta: beta1.clone(),
            };
            let key2 = CurveKeys {
                alpha: alpha.clone(),
                beta: beta2.clone(),
            };
            let expected = CurveKeys {
                alpha,
                beta: beta1 + beta2,
            };
            assert_eq!(key1 + key2, expected);
        }

        #[test]
        fn test_subtraction() {
            let mut rng = random::test_rng();
            let alpha = GlobalCurveKey::<C>::random(&mut rng);
            let beta1 = Ps::random(&mut rng);
            let beta2 = Ps::random(&mut rng);
            let key1 = CurveKeys {
                alpha: alpha.clone(),
                beta: beta1.clone(),
            };
            let key2 = CurveKeys {
                alpha: alpha.clone(),
                beta: beta2.clone(),
            };
            let expected = CurveKeys {
                alpha,
                beta: beta1 - beta2,
            };
            assert_eq!(key1 - key2, expected);
        }

        #[test]
        fn test_multiplication() {
            let mut rng = random::test_rng();
            let alpha = GlobalCurveKey::<C>::random(&mut rng);
            let beta1 = Ps::random(&mut rng);
            let key = CurveKeys {
                alpha: alpha.clone(),
                beta: beta1.clone(),
            };
            let scalar = FrExt::from(3u32);
            let expected = CurveKeys {
                alpha,
                beta: beta1 * &scalar,
            };
            assert_eq!(key * &scalar, expected);
        }

        #[test]
        fn test_negation() {
            let mut rng = random::test_rng();
            let alpha = GlobalCurveKey::<C>::random(&mut rng);
            let beta1 = Ps::random(&mut rng);
            let key = CurveKeys {
                alpha: alpha.clone(),
                beta: beta1.clone(),
            };
            let expected = CurveKeys {
                alpha,
                beta: -beta1,
            };
            assert_eq!(-key, expected);
        }
    }
}