arcium-primitives 0.4.1

use std::{
    ops::{Add, AddAssign, Deref, Mul, MulAssign, Neg, Sub, SubAssign},
    sync::Arc,
};

use rand::{distributions::Standard, prelude::Distribution, Rng};
use serde::{Deserialize, Serialize};
use subtle::{Choice, ConstantTimeEq};
use wincode::{SchemaRead, SchemaWrite};

use super::OpenFieldShare;
use crate::{
    algebra::field::{FieldElement, FieldExtension, SubfieldElement},
    errors::PrimitiveError,
    random::{CryptoRngCore, Random, RandomWith},
    types::ConditionallySelectable,
};

/// α, 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 β).
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, SchemaRead, SchemaWrite)]
#[serde(bound = "F: FieldExtension")]
#[repr(C)]
pub struct GlobalFieldKey<F: FieldExtension>(pub Arc<FieldElement<F>>);

impl<F: FieldExtension> GlobalFieldKey<F> {
    #[inline]
    pub fn new(value: FieldElement<F>) -> Self {
        GlobalFieldKey(Arc::new(value))
    }
}

impl<F: FieldExtension> ConditionallySelectable for GlobalFieldKey<F> {
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        GlobalFieldKey(Arc::new(FieldElement::conditional_select(
            &a.0, &b.0, choice,
        )))
    }
}

impl<F: FieldExtension> Deref for GlobalFieldKey<F> {
    type Target = FieldElement<F>;

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.0.deref()
    }
}

impl<F: FieldExtension> From<Arc<FieldElement<F>>> for GlobalFieldKey<F> {
    #[inline]
    fn from(value: Arc<FieldElement<F>>) -> Self {
        GlobalFieldKey(value)
    }
}

impl<F: FieldExtension> From<FieldElement<F>> for GlobalFieldKey<F> {
    #[inline]
    fn from(value: FieldElement<F>) -> Self {
        GlobalFieldKey(Arc::new(value))
    }
}

// -------------------------
// |   Random Generation   |
// -------------------------

impl<F: FieldExtension> Distribution<GlobalFieldKey<F>> for Standard {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> GlobalFieldKey<F> {
        GlobalFieldKey::new(Standard.sample(rng))
    }
}

impl<F: FieldExtension> RandomWith<usize> for Vec<Vec<GlobalFieldKey<F>>> {
    /// Generates `n_parties` vectors of `n_parties - 1` global keys each.
    fn random_with(mut rng: impl CryptoRngCore, n_parties: usize) -> Self {
        (0..n_parties)
            .map(|_| GlobalFieldKey::<F>::random_n(&mut rng, n_parties - 1))
            .collect()
    }
}

// α and β, such that MAC(x) = α · x + β
// In the context of VOLE, this corresponds to w = Δ · u + v
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, SchemaRead, SchemaWrite)]
#[serde(bound = "F: FieldExtension")]
#[repr(C)]
pub struct FieldShareKey<F: FieldExtension> {
    pub(crate) alpha: GlobalFieldKey<F>, // α, global key
    pub(crate) beta: FieldElement<F>,    // β, a local key per shared value.
}

impl<F: FieldExtension> FieldShareKey<F> {
    #[inline]
    pub fn new(alpha: GlobalFieldKey<F>, beta: FieldElement<F>) -> Self {
        FieldShareKey { alpha, beta }
    }

    #[inline]
    pub fn get_alpha(&self) -> GlobalFieldKey<F> {
        self.alpha.clone()
    }

    #[inline]
    pub fn get_alpha_value(&self) -> FieldElement<F> {
        *self.alpha
    }

    #[inline]
    pub fn get_beta(&self) -> &FieldElement<F> {
        &self.beta
    }

    /* === Verification === */

    #[inline]
    pub fn compute_mac(&self, value: &SubfieldElement<F>) -> FieldElement<F> {
        self.beta + *self.alpha * value
    }

    #[inline]
    pub fn verify_mac(&self, open_share: &OpenFieldShare<F>) -> Result<(), PrimitiveError> {
        let expected_mac = self.compute_mac(&open_share.value);
        bool::from(FieldElement::ct_eq(&expected_mac, &open_share.mac))
            .then_some(())
            .ok_or_else(|| {
                PrimitiveError::WrongMAC(serde_json::to_string(&open_share.mac).unwrap())
            })
    }
}

// -------------------------
// |   Random Generation   |
// -------------------------

impl<F: FieldExtension> Random for FieldShareKey<F> {
    fn random(mut rng: impl CryptoRngCore) -> Self {
        FieldShareKey {
            alpha: GlobalFieldKey::random(&mut rng),
            beta: FieldElement::random(&mut rng),
        }
    }
}

impl<F: FieldExtension> RandomWith<GlobalFieldKey<F>> for FieldShareKey<F> {
    fn random_with(mut rng: impl CryptoRngCore, alpha: GlobalFieldKey<F>) -> Self {
        FieldShareKey {
            alpha,
            beta: FieldElement::random(&mut rng),
        }
    }
}

// --------------
// | Arithmetic |
// --------------

// === Addition === //

#[macros::op_variants(owned, borrowed, flipped_commutative)]
impl<F: FieldExtension> Add<&FieldShareKey<F>> for FieldShareKey<F> {
    type Output = FieldShareKey<F>;

    #[inline]
    fn add(mut self, other: &FieldShareKey<F>) -> Self::Output {
        assert_eq!(self.alpha, other.alpha);
        self.beta += &other.beta;
        self
    }
}

#[macros::op_variants(owned)]
impl<'a, F: FieldExtension> AddAssign<&'a FieldShareKey<F>> for FieldShareKey<F> {
    #[inline]
    fn add_assign(&mut self, other: &'a FieldShareKey<F>) {
        assert_eq!(self.alpha, other.alpha);
        self.beta += &other.beta;
    }
}

// === Subtraction === //

#[macros::op_variants(owned, borrowed, flipped)]
impl<F: FieldExtension> Sub<&FieldShareKey<F>> for FieldShareKey<F> {
    type Output = FieldShareKey<F>;

    #[inline]
    fn sub(mut self, other: &FieldShareKey<F>) -> Self::Output {
        assert_eq!(self.alpha, other.alpha);
        self.beta -= other.beta;
        self
    }
}

#[macros::op_variants(owned)]
impl<'a, F: FieldExtension> SubAssign<&'a FieldShareKey<F>> for FieldShareKey<F> {
    #[inline]
    fn sub_assign(&mut self, other: &'a FieldShareKey<F>) {
        assert_eq!(self.alpha, other.alpha);
        self.beta -= &other.beta;
    }
}

// === One-to-many multiplication === //

#[macros::op_variants(owned, borrowed, flipped)]
impl<'a, F: FieldExtension> Mul<&'a FieldElement<F>> for FieldShareKey<F> {
    type Output = FieldShareKey<F>;

    #[inline]
    fn mul(mut self, other: &'a FieldElement<F>) -> Self::Output {
        self.beta *= other;
        self
    }
}

#[macros::op_variants(owned, borrowed, flipped)]
impl<'a, F: FieldExtension> Mul<&'a SubfieldElement<F>> for FieldShareKey<F> {
    type Output = FieldShareKey<F>;

    #[inline]
    fn mul(mut self, other: &'a SubfieldElement<F>) -> Self::Output {
        self.beta *= other;
        self
    }
}

#[macros::op_variants(owned)]
impl<'a, F: FieldExtension> MulAssign<&'a FieldElement<F>> for FieldShareKey<F> {
    #[inline]
    fn mul_assign(&mut self, other: &'a FieldElement<F>) {
        self.beta *= other;
    }
}

#[macros::op_variants(owned)]
impl<'a, F: FieldExtension> MulAssign<&'a SubfieldElement<F>> for FieldShareKey<F> {
    #[inline]
    fn mul_assign(&mut self, other: &'a SubfieldElement<F>) {
        self.beta *= other;
    }
}

// === Negation === //

#[macros::op_variants(borrowed)]
impl<F: FieldExtension> Neg for FieldShareKey<F> {
    type Output = FieldShareKey<F>;

    #[inline]
    fn neg(self) -> Self::Output {
        FieldShareKey {
            alpha: self.alpha,
            beta: -self.beta,
        }
    }
}

// === Constant Addition === //

impl<F: FieldExtension> FieldShareKey<F> {
    #[inline]
    pub fn add_secret_owned(mut self, constant: FieldElement<F>) -> Result<Self, PrimitiveError> {
        self.beta -= constant * *self.alpha;
        Ok(self)
    }

    #[inline]
    pub fn add_secret(&self, constant: FieldElement<F>) -> Result<Self, PrimitiveError> {
        let mut result = self.clone();
        result.beta -= constant * *result.alpha;
        Ok(result)
    }
}

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

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

impl<F: FieldExtension> FieldShareKey<F> {
    #[inline]
    pub fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        FieldShareKey {
            alpha: GlobalFieldKey::conditional_select(&a.alpha, &b.alpha, choice),
            beta: FieldElement::conditional_select(&a.beta, &b.beta, choice),
        }
    }
}

#[cfg(test)]
mod tests {

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

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

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

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

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