arcium-primitives 0.4.4

use crate::{
    algebra::field::{FieldExtension, SubfieldElement},
    correlated_randomness::triples::{
        Triple,
        Triples,
        UnauthenticatedTriple,
        UnauthenticatedTriples,
    },
    izip_eq,
    random::{CryptoRngCore, Random, RandomWith},
    sharing::{unauthenticated::AdditiveShares, FieldShare, FieldShares, GlobalFieldKey},
    types::{heap_array::SubfieldElements, Positive},
    utils::IntoExactSizeIterator,
};
// -------------------------
// |   Random Generation   |
// -------------------------

impl<F: FieldExtension> Random for UnauthenticatedTriple<F> {
    fn random(_rng: impl CryptoRngCore) -> Self {
        unimplemented!("UnauthenticatedTriple is correlated randomness. Use `random_n` to establish correlation for n_parties")
    }

    /// Generates random triples for all the parties requested.
    fn random_n<Container: FromIterator<Self>>(
        mut rng: impl CryptoRngCore,
        n_parties: usize,
    ) -> Container {
        let a = SubfieldElement::random(&mut rng);
        let b = SubfieldElement::random(&mut rng);
        let c = a * b;

        let a_shares = a.to_additive_shares(n_parties, &mut rng);
        let b_shares = b.to_additive_shares(n_parties, &mut rng);
        let c_shares = c.to_additive_shares(n_parties, &mut rng);

        izip_eq!(a_shares, b_shares, c_shares)
            .map(|(a, b, c)| UnauthenticatedTriple { a, b, c })
            .collect()
    }
}

impl<F: FieldExtension, M: Positive> Random for UnauthenticatedTriples<F, M> {
    fn random(_rng: impl CryptoRngCore) -> Self {
        unimplemented!("UnauthenticatedTriples are correlated randomness. Use `random_n` to establish correlation for n_parties")
    }

    /// Generates random triples for all the parties requested.
    fn random_n<Container: FromIterator<Self>>(
        mut rng: impl CryptoRngCore,
        n_parties: usize,
    ) -> Container {
        let a = SubfieldElements::random(&mut rng);
        let b = SubfieldElements::random(&mut rng);
        let c = &a * &b;

        let a_shares = a.to_additive_shares(n_parties, &mut rng);
        let b_shares = b.to_additive_shares(n_parties, &mut rng);
        let c_shares = c.to_additive_shares(n_parties, &mut rng);

        izip_eq!(a_shares, b_shares, c_shares)
            .map(|(a, b, c)| UnauthenticatedTriples { a, b, c })
            .collect()
    }
}

impl<F: FieldExtension, M: Positive> Random for Triples<F, M> {
    fn random(_rng: impl CryptoRngCore) -> Self {
        unimplemented!("Triples are correlated randomness. Use `random_n` to establish correlation for n_parties")
    }

    /// Generates random triples for all the parties requested.
    fn random_n<Container: FromIterator<Self>>(
        mut rng: impl CryptoRngCore,
        n_parties: usize,
    ) -> Container {
        let a = SubfieldElements::random(&mut rng);
        let b = SubfieldElements::random(&mut rng);
        let c = &a * &b;

        let a_shares: Vec<_> = FieldShares::random_n_with(&mut rng, n_parties, a);
        let b_shares: Vec<_> = FieldShares::random_n_with(&mut rng, n_parties, b);
        let c_shares: Vec<_> = FieldShares::random_n_with(&mut rng, n_parties, c);

        izip_eq!(a_shares, b_shares, c_shares)
            .map(|(a, b, c)| Triples { a, b, c })
            .collect()
    }
}

impl<F: FieldExtension, M: Positive> RandomWith<Vec<GlobalFieldKey<F>>> for Triples<F, M> {
    /// Generate a single random triple share for the global keys of a single party.
    ///
    /// WARNING: This does not establish correlation with other parties! Use only to mock a single
    /// party's triple.
    #[allow(unused_mut)]
    fn random_with(mut rng: impl CryptoRngCore, alphas: Vec<GlobalFieldKey<F>>) -> Self {
        #[cfg(any(test, feature = "dev"))]
        {
            Triples {
                a: FieldShares::<F, M>::random_with(&mut rng, alphas.clone()),
                b: FieldShares::<F, M>::random_with(&mut rng, alphas.clone()),
                c: FieldShares::<F, M>::random_with(&mut rng, alphas),
            }
        }

        #[cfg(not(any(test, feature = "dev")))]
        unimplemented!(
            "Triples are correlated randomness. Use `random_n_with_each({}, Vec<{}>)` to establish correlation for n_parties", std::any::type_name_of_val(&rng),
            std::any::type_name_of_val(&alphas)
        );
    }

    /// Generate random triples given the global keys (alphas) for all parties.
    fn random_n_with_each<Container: FromIterator<Self>>(
        mut rng: impl CryptoRngCore,
        all_alphas: impl IntoExactSizeIterator<Item = Vec<GlobalFieldKey<F>>>,
    ) -> Container {
        let all_alphas = all_alphas.into_iter().collect::<Vec<_>>();
        let a = SubfieldElements::random(&mut rng);
        let b = SubfieldElements::random(&mut rng);
        let c = &a * &b;

        let n_parties = all_alphas.len();
        let a_shares: Vec<_> =
            FieldShares::random_n_with(&mut rng, n_parties, (a, all_alphas.clone()));
        let b_shares: Vec<_> =
            FieldShares::random_n_with(&mut rng, n_parties, (b, all_alphas.clone()));
        let c_shares: Vec<_> = FieldShares::random_n_with(&mut rng, n_parties, (c, all_alphas));

        izip_eq!(a_shares, b_shares, c_shares)
            .map(|(a, b, c)| Triples { a, b, c })
            .collect()
    }
}

impl<F: FieldExtension> Random for Triple<F> {
    fn random(_rng: impl CryptoRngCore) -> Self {
        unimplemented!("Triples are correlated randomness. Use `random_n`/`random_n_with`` to establish correlation for n_parties")
    }

    /// Generates random triples for all the parties requested.
    fn random_n<Container: FromIterator<Self>>(
        mut rng: impl CryptoRngCore,
        n_parties: usize,
    ) -> Container {
        let a = SubfieldElement::random(&mut rng);
        let b = SubfieldElement::random(&mut rng);
        let c = a * b;

        let a_shares: Vec<_> = FieldShare::random_n_with(&mut rng, n_parties, a);
        let b_shares: Vec<_> = FieldShare::random_n_with(&mut rng, n_parties, b);
        let c_shares: Vec<_> = FieldShare::random_n_with(&mut rng, n_parties, c);

        izip_eq!(a_shares, b_shares, c_shares)
            .map(|(a, b, c)| Triple { a, b, c })
            .collect()
    }
}

impl<F: FieldExtension> RandomWith<Vec<GlobalFieldKey<F>>> for Triple<F> {
    /// Generate a single random triple share for the global keys of a single party.
    #[allow(unused_mut)]
    fn random_with(mut rng: impl CryptoRngCore, alphas: Vec<GlobalFieldKey<F>>) -> Self {
        #[cfg(any(test, feature = "dev"))]
        {
            Triple {
                a: FieldShare::random_with(&mut rng, alphas.clone()),
                b: FieldShare::random_with(&mut rng, alphas.clone()),
                c: FieldShare::random_with(&mut rng, alphas),
            }
        }

        #[cfg(not(any(test, feature = "dev")))]
        unimplemented!(
            "Triples are correlated randomness. Use `random_n_with_each({}, Vec<{}>)` to establish correlation for n_parties", std::any::type_name_of_val(&rng),
            std::any::type_name_of_val(&alphas)
        );
    }

    /// Generate random triples given the global keys (alphas) for all parties.
    fn random_n_with_each<Container: FromIterator<Self>>(
        mut rng: impl CryptoRngCore,
        all_alphas: impl IntoExactSizeIterator<Item = Vec<GlobalFieldKey<F>>>,
    ) -> Container {
        let all_alphas = all_alphas.into_iter().collect::<Vec<_>>();
        let a = SubfieldElement::random(&mut rng);
        let b = SubfieldElement::random(&mut rng);
        let c = a * b;

        let n_parties = all_alphas.len();
        let a_shares: Vec<_> =
            FieldShare::random_n_with(&mut rng, n_parties, (a, all_alphas.clone()));
        let b_shares: Vec<_> =
            FieldShare::random_n_with(&mut rng, n_parties, (b, all_alphas.clone()));
        let c_shares: Vec<_> = FieldShare::random_n_with(&mut rng, n_parties, (c, all_alphas));

        izip_eq!(a_shares, b_shares, c_shares)
            .map(|(a, b, c)| Triple { a, b, c })
            .collect()
    }
}