separate ParameterSelector for non-interactive mp and interactive mp

9 months ago · d260910299
12 changed files with 114 additions and 318 deletions
--- a/examples/interactive_fheuint8.rs
+++ b/examples/interactive_fheuint8.rs
@ -35,7 +35,8 @@ fn main() {
    // multi-party key gen round 2
    let server_key_shares = client_keys
        .iter()
        .map(|k| gen_mp_keys_phase2(k, &public_key))
        .enumerate()
        .map(|(user_id, k)| gen_mp_keys_phase2(k, user_id, no_of_parties, &public_key))
        .collect_vec();

    // server aggregates server key shares and sets it
--- a/examples/non_interactive_fheuint8.rs
+++ b/examples/non_interactive_fheuint8.rs
@ -47,10 +47,15 @@ fn main() {
    server_key.set_server_key();

    // extract a and b from client0 inputs
    // let now = std::time::Instant::now();
    let (ct_c0_a, ct_c0_b) = {
        let ct = c0_batched_to_send.unseed::<Vec<Vec<u64>>>().key_switch(0);
        (ct.extract(0), ct.extract(1))
    };
    // println!(
    //     "Time to unseed, key switch, and extract 2 ciphertexts: {:?}",
    //     now.elapsed()
    // );

    // extract a and b from client1 inputs
    let (ct_c1_a, ct_c1_b) = {
--- a/src/bool/evaluator.rs
+++ b/src/bool/evaluator.rs
@ -6,7 +6,7 @@ use std::{
 };

 use itertools::{izip, Itertools};
 use num_traits::{FromPrimitive, Num, One, PrimInt, ToPrimitive, WrappingAdd, WrappingSub, Zero};
 use num_traits::{FromPrimitive, One, PrimInt, ToPrimitive, WrappingAdd, WrappingSub, Zero};
 use rand_distr::uniform::SampleUniform;

 use crate::{
@ -19,7 +19,7 @@ use crate::{
        public_key_share,
    },
    ntt::{Ntt, NttInit},
    pbs::{pbs, sample_extract, PbsInfo, PbsKey, WithShoupRepr},
    pbs::{pbs, PbsInfo, PbsKey, WithShoupRepr},
    random::{
        DefaultSecureRng, NewWithSeed, RandomFill, RandomFillGaussianInModulus,
        RandomFillUniformInModulus, RandomGaussianElementInModulus,
@ -30,8 +30,8 @@ use crate::{
        seeded_auto_key_gen, RgswCiphertextMutRef, RgswCiphertextRef, RuntimeScratchMutRef,
    },
    utils::{
        encode_x_pow_si_with_emebedding_factor, fill_random_ternary_secret_with_hamming_weight,
        mod_exponent, puncture_p_rng, TryConvertFrom1, WithLocal,
        encode_x_pow_si_with_emebedding_factor, mod_exponent, puncture_p_rng, TryConvertFrom1,
        WithLocal,
    },
    Encoder, Matrix, MatrixEntity, MatrixMut, RowEntity, RowMut,
 };
@ -61,11 +61,11 @@ use super::{
 ///             Puncture 1 -> Auto keys cipertexts seed
 ///             Puncture 2 -> LWE ksk seed
 #[derive(Clone, PartialEq)]
 pub struct MultiPartyCrs<S> {
 pub struct InteractiveMultiPartyCrs<S> {
    pub(super) seed: S,
 }

 impl MultiPartyCrs<[u8; 32]> {
 impl InteractiveMultiPartyCrs<[u8; 32]> {
    pub(super) fn random() -> Self {
        DefaultSecureRng::with_local_mut(|rng| {
            let mut seed = [0u8; 32];
@ -74,7 +74,8 @@ impl MultiPartyCrs<[u8; 32]> {
        })
    }
 }
 impl<S: Default + Copy> MultiPartyCrs<S> {

 impl<S: Default + Copy> InteractiveMultiPartyCrs<S> {
    /// Seed to generate public key share
    fn public_key_share_seed<Rng: NewWithSeed<Seed = S> + RandomFill<S>>(&self) -> S {
        let mut prng = Rng::new_with_seed(self.seed);
@ -110,7 +111,7 @@ impl MultiPartyCrs {
 ///         Puncture 3 -> Lwe key switching key seed
 ///     Puncture 2 -> user specific seed for u_j to s ksk
 ///         Punture j+1 -> user j's seed    
 #[derive(Clone)]
 #[derive(Clone, PartialEq)]
 pub struct NonInteractiveMultiPartyCrs<S> {
    pub(super) seed: S,
 }
@ -660,7 +661,7 @@ where
 /// Assigns user with user_id segement of LWE secret indices for which they
 /// generate RGSW(X^{s[i]}) as the leader (i.e. for RLWExRGSW). If returned
 /// tuple is (start, end), user's segment is [start, end)
 pub(super) fn interactive_mult_party_user_id_lwe_segment(
 pub(super) fn multi_party_user_id_lwe_segment(
    user_id: usize,
    total_users: usize,
    lwe_n: usize,
@ -946,17 +947,19 @@ where
        })
    }

    pub(super) fn multi_party_server_key_share<K: InteractiveMultiPartyClientKey<Element = i32>>(
    pub(super) fn gen_interactive_multi_party_server_key_share<
        K: InteractiveMultiPartyClientKey<Element = i32>,
    >(
        &self,
        user_id: usize,
        total_users: usize,
        cr_seed: &MultiPartyCrs<[u8; 32]>,
        cr_seed: &InteractiveMultiPartyCrs<[u8; 32]>,
        collective_pk: &M,
        client_key: &K,
    ) -> CommonReferenceSeededInteractiveMultiPartyServerKeyShare<
        M,
        BoolParameters<M::MatElement>,
        MultiPartyCrs<[u8; 32]>,
        InteractiveMultiPartyCrs<[u8; 32]>,
    > {
        assert_eq!(self.parameters().variant(), &ParameterVariant::MultiParty);
        assert!(user_id < total_users);
@ -987,11 +990,8 @@ where
            let mut self_leader_rgsw = vec![];
            let mut not_self_leader_rgsws = vec![];

            let (segment_start, segment_end) = interactive_mult_party_user_id_lwe_segment(
                user_id,
                total_users,
                self.pbs_info().lwe_n(),
            );
            let (segment_start, segment_end) =
                multi_party_user_id_lwe_segment(user_id, total_users, self.pbs_info().lwe_n());

            // self LWE secret indices
            {
@ -1095,14 +1095,18 @@ where
        )
    }

    pub(super) fn aggregate_multi_party_server_key_shares<S>(
    pub(super) fn aggregate_interactive_multi_party_server_key_shares<S>(
        &self,
        shares: &[CommonReferenceSeededInteractiveMultiPartyServerKeyShare<
            M,
            BoolParameters<M::MatElement>,
            MultiPartyCrs<S>,
            InteractiveMultiPartyCrs<S>,
        >],
    ) -> SeededInteractiveMultiPartyServerKey<M, MultiPartyCrs<S>, BoolParameters<M::MatElement>>
    ) -> SeededInteractiveMultiPartyServerKey<
        M,
        InteractiveMultiPartyCrs<S>,
        BoolParameters<M::MatElement>,
    >
    where
        S: PartialEq + Clone,
        M: Clone,
@ -1116,9 +1120,6 @@ where
        let cr_seed = shares[0].cr_seed();

        let rlwe_n = parameters.rlwe_n().0;
        let g = parameters.g() as isize;
        let rlwe_q = parameters.rlwe_q();
        let lwe_q = parameters.lwe_q();

        // sanity checks
        shares.iter().skip(1).for_each(|s| {
@ -1181,7 +1182,7 @@ where
                ..total_users)
                .map(|(user_id)| {
                    let (start_index, end_index) =
                        interactive_mult_party_user_id_lwe_segment(user_id, total_users, lwe_n);
                        multi_party_user_id_lwe_segment(user_id, total_users, lwe_n);
                    ((start_index, end_index), end_index - start_index)
                })
                .unzip();
@ -1263,10 +1264,9 @@ where
        )
    }

    pub(super) fn aggregate_non_interactive_multi_party_key_share(
    pub(super) fn aggregate_non_interactive_multi_party_server_key_shares(
        &self,
        cr_seed: &NonInteractiveMultiPartyCrs<[u8; 32]>,
        total_users: usize,
        key_shares: &[CommonReferenceSeededNonInteractiveMultiPartyServerKeyShare<
            M,
            BoolParameters<M::MatElement>,
@ -1286,8 +1286,6 @@ where
            &ParameterVariant::NonInteractiveMultiParty
        );

        // TODO: Check parameters are equivalent!

        let total_users = key_shares.len();
        let key_shares = (0..total_users)
            .map(|user_id| {
@ -1299,6 +1297,14 @@ where
            })
            .collect_vec();

        // check parameters and cr seed are equal
        {
            key_shares.iter().for_each(|k| {
                assert!(k.parameters() == self.parameters());
                assert!(k.cr_seed() == cr_seed);
            });
        }

        let rlwe_modop = &self.pbs_info().rlwe_modop;
        let nttop = &self.pbs_info().rlwe_nttop;
        let ring_size = self.parameters().rlwe_n().0;
@ -1410,7 +1416,7 @@ where
            // a_{i, l} * s + e = \sum_{j \in P} a_{i, l} * s_{j} + e
            let user_segments = (0..total_users)
                .map(|user_id| {
                    interactive_mult_party_user_id_lwe_segment(
                    multi_party_user_id_lwe_segment(
                        user_id,
                        total_users,
                        self.parameters().lwe_n().0,
@ -1660,7 +1666,7 @@ where
        )
    }

    pub(super) fn non_interactive_multi_party_key_share<
    pub(super) fn gen_non_interactive_multi_party_key_share<
        K: NonInteractiveMultiPartyClientKey<Element = i32>,
    >(
        &self,
@ -1822,7 +1828,7 @@ where

            // Generate non-interactive RGSW ciphertexts a_{i, l} u_j + e + \beta X^{s_j[l]}
            // for i \in (0, d_max]
            let (self_start_index, self_end_index) = interactive_mult_party_user_id_lwe_segment(
            let (self_start_index, self_end_index) = multi_party_user_id_lwe_segment(
                self_index,
                total_users,
                self.parameters().lwe_n().0,
@ -2086,7 +2092,7 @@ where

    pub(super) fn multi_party_public_key_share<K: InteractiveMultiPartyClientKey<Element = i32>>(
        &self,
        cr_seed: &MultiPartyCrs<[u8; 32]>,
        cr_seed: &InteractiveMultiPartyCrs<[u8; 32]>,
        client_key: &K,
    ) -> CommonReferenceSeededCollectivePublicKeyShare<
        <M as Matrix>::R,
@ -2150,94 +2156,6 @@ where
        self.pbs_info.parameters.rlwe_q().decode(encoded_m)
    }

    pub(crate) fn pk_encrypt(&self, pk: &M, m: bool) -> M::R {
        self.pk_encrypt_batched(pk, &vec![m]).remove(0)
    }

    /// Encrypts a batch booleans as multiple LWE ciphertexts.
    ///
    /// For public key encryption we first encrypt `m` as a RLWE ciphertext and
    /// then sample extract LWE samples at required indices.
    ///
    /// - TODO(Jay:) Communication can be improved by not sample exctracting and
    ///   instead just truncate degree 0 values (part Bs)
    pub(crate) fn pk_encrypt_batched(&self, pk: &M, m: &[bool]) -> Vec<M::R> {
        DefaultSecureRng::with_local_mut(|rng| {
            let ring_size = self.pbs_info.parameters.rlwe_n().0;
            assert!(
                m.len() <= ring_size,
                "Cannot batch encrypt > ring_size{ring_size} elements at once"
            );

            let modop = &self.pbs_info.rlwe_modop;
            let nttop = &self.pbs_info.rlwe_nttop;

            // RLWE(0)
            // sample ephemeral key u
            let mut u = vec![0i32; ring_size];
            fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
            let mut u = M::R::try_convert_from(&u, &self.pbs_info.parameters.rlwe_q());
            nttop.forward(u.as_mut());

            let mut ua = M::R::zeros(ring_size);
            ua.as_mut().copy_from_slice(pk.get_row_slice(0));
            let mut ub = M::R::zeros(ring_size);
            ub.as_mut().copy_from_slice(pk.get_row_slice(1));

            // a*u
            nttop.forward(ua.as_mut());
            modop.elwise_mul_mut(ua.as_mut(), u.as_ref());
            nttop.backward(ua.as_mut());

            // b*u
            nttop.forward(ub.as_mut());
            modop.elwise_mul_mut(ub.as_mut(), u.as_ref());
            nttop.backward(ub.as_mut());

            let mut rlwe = M::zeros(2, ring_size);
            // sample error
            rlwe.iter_rows_mut().for_each(|ri| {
                RandomFillGaussianInModulus::<[M::MatElement], CiphertextModulus<M::MatElement>>::random_fill(
                    rng,
                    &self.pbs_info.parameters.rlwe_q(),
                    ri.as_mut(),
                );
            });

            // a*u + e0
            modop.elwise_add_mut(rlwe.get_row_mut(0), ua.as_ref());
            // b*u + e1
            modop.elwise_add_mut(rlwe.get_row_mut(1), ub.as_ref());

            //FIXME(Jay): Figure out a way to get Q/8 form modulus
            let mut m_vec = M::R::zeros(ring_size);
            izip!(m_vec.as_mut().iter_mut(), m.iter()).for_each(|(m_el, m_bool)| {
                if *m_bool {
                    // Q/8
                    *m_el = self.pbs_info.rlwe_q().true_el()
                } else {
                    // -Q/8
                    *m_el = self.pbs_info.rlwe_q().false_el()
                }
            });

            // b*u + e1 + m
            modop.elwise_add_mut(rlwe.get_row_mut(1), m_vec.as_ref());
            // rlwe.set(1, 0, modop.add(rlwe.get(1, 0), &m));

            // sample extract index required indices
            let samples = m.len();
            (0..samples)
                .into_iter()
                .map(|i| {
                    let mut lwe_out = M::R::zeros(ring_size + 1);
                    sample_extract(&mut lwe_out, &rlwe, modop, i);
                    lwe_out
                })
                .collect_vec()
        })
    }

    pub fn sk_encrypt<K: SinglePartyClientKey<Element = i32>>(
        &self,
        m: bool,
--- a/src/bool/keys.rs
+++ b/src/bool/keys.rs
@ -497,7 +497,7 @@ pub(super) mod impl_server_key_eval_domain {
    use itertools::{izip, Itertools};

    use crate::{
        evaluator::MultiPartyCrs,
        evaluator::InteractiveMultiPartyCrs,
        ntt::{Ntt, NttInit},
        pbs::PbsKey,
        random::RandomFill,
@ -651,7 +651,7 @@ pub(super) mod impl_server_key_eval_domain {
        From<
            &SeededInteractiveMultiPartyServerKey<
                M,
                MultiPartyCrs<Rng::Seed>,
                InteractiveMultiPartyCrs<Rng::Seed>,
                BoolParameters<M::MatElement>,
            >,
        > for ServerKeyEvaluationDomain<M, BoolParameters<M::MatElement>, Rng, N>
@ -665,7 +665,7 @@ pub(super) mod impl_server_key_eval_domain {
        fn from(
            value: &SeededInteractiveMultiPartyServerKey<
                M,
                MultiPartyCrs<Rng::Seed>,
                InteractiveMultiPartyCrs<Rng::Seed>,
                BoolParameters<M::MatElement>,
            >,
        ) -> Self {
@ -1268,7 +1268,7 @@ pub struct CommonReferenceSeededNonInteractiveMultiPartyServerKeyShare
 }

 mod impl_common_ref_non_interactive_multi_party_server_share {
    use crate::evaluator::interactive_mult_party_user_id_lwe_segment;
    use crate::evaluator::multi_party_user_id_lwe_segment;

    use super::*;

@ -1307,11 +1307,8 @@ mod impl_common_ref_non_interactive_multi_party_server_share {
            &self,
            lwe_index: usize,
        ) -> &M {
            let self_segment = interactive_mult_party_user_id_lwe_segment(
                self.user_id,
                self.total_users,
                self.lwe_n,
            );
            let self_segment =
                multi_party_user_id_lwe_segment(self.user_id, self.total_users, self.lwe_n);
            assert!(lwe_index >= self_segment.0 && lwe_index < self_segment.1);
            &self.self_leader_ni_rgsw_cts[lwe_index - self_segment.0]
        }
@ -1320,11 +1317,8 @@ mod impl_common_ref_non_interactive_multi_party_server_share {
            &self,
            lwe_index: usize,
        ) -> &M {
            let self_segment = interactive_mult_party_user_id_lwe_segment(
                self.user_id,
                self.total_users,
                self.lwe_n,
            );
            let self_segment =
                multi_party_user_id_lwe_segment(self.user_id, self.total_users, self.lwe_n);
            // Non-interactive RGSW cts when self is not leader are stored in
            // sorted-order. For ex, if self is the leader for indices (5, 6]
            // then self stores NI-RGSW cts for rest of indices like [0, 1, 2,
@ -1366,6 +1360,14 @@ mod impl_common_ref_non_interactive_multi_party_server_share {
                &self.ksk_zero_encs_for_others[user_i - 1]
            }
        }

        pub(in super::super) fn cr_seed(&self) -> &S {
            &self.cr_seed
        }

        pub(in super::super) fn parameters(&self) -> &P {
            &self.parameters
        }
    }
 }

--- a/src/bool/mod.rs
+++ b/src/bool/mod.rs
@ -1,11 +1,7 @@
 pub(crate) mod evaluator;
 mod keys;
 mod noise;
 pub(crate) mod parameters;

 use keys::tests;
 pub(crate) use keys::PublicKey;

 #[cfg(feature = "interactive_mp")]
 mod mp_api;
 #[cfg(feature = "non_interactive_mp")]
@ -19,12 +15,6 @@ pub use mp_api::*;

 pub type ClientKey = keys::ClientKey<[u8; 32], u64>;

 pub enum ParameterSelector {
    InteractiveLTE2Party,
    NonInteractiveLTE2Party,
    NonInteractiveLTE4Party,
 }

 mod common_mp_enc_dec {
    use super::BoolEvaluator;
    use crate::{
--- a/src/bool/mp_api.rs
+++ b/src/bool/mp_api.rs
@ -7,7 +7,7 @@ use crate::{
    utils::{Global, WithLocal},
 };

 use super::{evaluator::MultiPartyCrs, keys::*, parameters::*, ClientKey, ParameterSelector};
 use super::{evaluator::InteractiveMultiPartyCrs, keys::*, parameters::*, ClientKey};

 pub(crate) type BoolEvaluator = super::evaluator::BoolEvaluator<
    Vec<Vec<u64>>,
@ -23,8 +23,13 @@ thread_local! {
 }
 static BOOL_SERVER_KEY: OnceLock<ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>> = OnceLock::new();

 static MULTI_PARTY_CRS: OnceLock<MultiPartyCrs<[u8; 32]>> = OnceLock::new();
 static MULTI_PARTY_CRS: OnceLock<InteractiveMultiPartyCrs<[u8; 32]>> = OnceLock::new();

 pub enum ParameterSelector {
    InteractiveLTE2Party,
 }

 /// Select Interactive multi-party parameter variant
 pub fn set_parameter_set(select: ParameterSelector) {
    match select {
        ParameterSelector::InteractiveLTE2Party => {
@ -32,31 +37,39 @@ pub fn set_parameter_set(select: ParameterSelector) {
        }

        _ => {
            panic!("Paramerters not supported")
            panic!("Paramerter not supported")
        }
    }
 }

 /// Set application specific interactive multi-party common reference string
 pub fn set_mp_seed(seed: [u8; 32]) {
    assert!(
        MULTI_PARTY_CRS.set(MultiPartyCrs { seed: seed }).is_ok(),
        MULTI_PARTY_CRS
            .set(InteractiveMultiPartyCrs { seed: seed })
            .is_ok(),
        "Attempted to set MP SEED twice."
    )
 }

 /// Generate client key for interactive multi-party protocol
 pub fn gen_client_key() -> ClientKey {
    BoolEvaluator::with_local(|e| e.client_key())
 }

 /// Generate client's share for collective public key, i.e round 1, of the
 /// protocol
 pub fn gen_mp_keys_phase1(
    ck: &ClientKey,
 ) -> CommonReferenceSeededCollectivePublicKeyShare<Vec<u64>, [u8; 32], BoolParameters<u64>> {
    BoolEvaluator::with_local(|e| {
        let pk_share = e.multi_party_public_key_share(MultiPartyCrs::global(), ck);
        let pk_share = e.multi_party_public_key_share(InteractiveMultiPartyCrs::global(), ck);
        pk_share
    })
 }

 /// Generate clients share for collective server key, i.e. round 2, of the
 /// protocol
 pub fn gen_mp_keys_phase2<R, ModOp>(
    ck: &ClientKey,
    user_id: usize,
@ -65,13 +78,13 @@ pub fn gen_mp_keys_phase2(
 ) -> CommonReferenceSeededInteractiveMultiPartyServerKeyShare<
    Vec<Vec<u64>>,
    BoolParameters<u64>,
    MultiPartyCrs<[u8; 32]>,
    InteractiveMultiPartyCrs<[u8; 32]>,
 > {
    BoolEvaluator::with_local_mut(|e| {
        let server_key_share = e.multi_party_server_key_share(
        let server_key_share = e.gen_interactive_multi_party_server_key_share(
            user_id,
            total_users,
            MultiPartyCrs::global(),
            InteractiveMultiPartyCrs::global(),
            pk.key(),
            ck,
        );
@ -79,6 +92,10 @@ pub fn gen_mp_keys_phase2(
    })
 }

 /// Aggregate public key shares from all parties.
 ///
 /// Public key shares are generated per client in round 1. Aggregation of public
 /// key shares marks the end of round 1.
 pub fn aggregate_public_key_shares(
    shares: &[CommonReferenceSeededCollectivePublicKeyShare<
        Vec<u64>,
@ -89,24 +106,29 @@ pub fn aggregate_public_key_shares(
    PublicKey::from(shares)
 }

 /// Aggregate server key shares
 pub fn aggregate_server_key_shares(
    shares: &[CommonReferenceSeededInteractiveMultiPartyServerKeyShare<
        Vec<Vec<u64>>,
        BoolParameters<u64>,
        MultiPartyCrs<[u8; 32]>,
        InteractiveMultiPartyCrs<[u8; 32]>,
    >],
 ) -> SeededInteractiveMultiPartyServerKey<Vec<Vec<u64>>, MultiPartyCrs<[u8; 32]>, BoolParameters<u64>>
 {
    BoolEvaluator::with_local(|e| e.aggregate_multi_party_server_key_shares(shares))
 ) -> SeededInteractiveMultiPartyServerKey<
    Vec<Vec<u64>>,
    InteractiveMultiPartyCrs<[u8; 32]>,
    BoolParameters<u64>,
 > {
    BoolEvaluator::with_local(|e| e.aggregate_interactive_multi_party_server_key_shares(shares))
 }

 impl
    SeededInteractiveMultiPartyServerKey<
        Vec<Vec<u64>>,
        MultiPartyCrs<<DefaultSecureRng as NewWithSeed>::Seed>,
        InteractiveMultiPartyCrs<<DefaultSecureRng as NewWithSeed>::Seed>,
        BoolParameters<u64>,
    >
 {
    /// Sets the server key as a global reference for circuit evaluation
    pub fn set_server_key(&self) {
        assert!(
            BOOL_SERVER_KEY
@ -120,7 +142,7 @@ impl
 }

 //  MULTIPARTY CRS  //
 impl Global for MultiPartyCrs<[u8; 32]> {
 impl Global for InteractiveMultiPartyCrs<[u8; 32]> {
    fn global() -> &'static Self {
        MULTI_PARTY_CRS
            .get()
@ -164,7 +186,7 @@ mod impl_enc_dec {
        bool::evaluator::BoolEncoding,
        pbs::{sample_extract, PbsInfo},
        rgsw::public_key_encrypt_rlwe,
        Encryptor, Matrix, MatrixEntity, MultiPartyDecryptor, RowEntity,
        Encryptor, Matrix, MatrixEntity, RowEntity,
    };
    use itertools::Itertools;
    use num_traits::{ToPrimitive, Zero};
--- a/src/bool/ni_mp_api.rs
+++ b/src/bool/ni_mp_api.rs
@ -17,7 +17,7 @@ use super::{
        ShoupNonInteractiveServerKeyEvaluationDomain,
    },
    parameters::{BoolParameters, CiphertextModulus, NI_2P},
    ClientKey, ParameterSelector,
    ClientKey,
 };

 pub(crate) type BoolEvaluator = super::evaluator::BoolEvaluator<
@ -37,6 +37,11 @@ static BOOL_SERVER_KEY: OnceLock

 static MULTI_PARTY_CRS: OnceLock<NonInteractiveMultiPartyCrs<[u8; 32]>> = OnceLock::new();

 pub enum ParameterSelector {
    NonInteractiveLTE2Party,
    NonInteractiveLTE4Party,
 }

 pub fn set_parameter_set(select: ParameterSelector) {
    match select {
        ParameterSelector::NonInteractiveLTE2Party => {
@ -84,7 +89,7 @@ pub fn gen_server_key_share(
 > {
    BoolEvaluator::with_local(|e| {
        let cr_seed = NonInteractiveMultiPartyCrs::global();
        e.non_interactive_multi_party_key_share(cr_seed, user_id, total_users, client_key)
        e.gen_non_interactive_multi_party_key_share(cr_seed, user_id, total_users, client_key)
    })
 }

@ -101,7 +106,7 @@ pub fn aggregate_server_key_shares(
 > {
    BoolEvaluator::with_local(|e| {
        let cr_seed = NonInteractiveMultiPartyCrs::global();
        e.aggregate_non_interactive_multi_party_key_share(cr_seed, shares.len(), shares)
        e.aggregate_non_interactive_multi_party_server_key_shares(cr_seed, shares)
    })
 }

@ -411,12 +416,12 @@ mod tests {

    #[test]
    fn non_interactive_mp_bool_nand() {
        set_parameter_set(ParameterSelector::NonInteractiveLTE2Party);
        set_parameter_set(ParameterSelector::NonInteractiveLTE4Party);
        let mut seed = [0u8; 32];
        thread_rng().fill_bytes(&mut seed);
        set_common_reference_seed(seed);

        let parties = 2;
        let parties = 3;

        let cks = (0..parties).map(|_| gen_client_key()).collect_vec();

--- a/src/bool/noise.rs
+++ b/src/bool/noise.rs
@ -1,144 +0,0 @@
 mod test {
    use itertools::{izip, Itertools};

    use crate::{
        backend::{ModularOpsU64, ModulusPowerOf2},
        bool::{
            evaluator::{BoolEncoding, BoolEvaluator, BooleanGates},
            keys::{
                InteractiveMultiPartyClientKey, PublicKey, ServerKeyEvaluationDomain,
                ShoupServerKeyEvaluationDomain,
            },
            parameters::{CiphertextModulus, SMALL_MP_BOOL_PARAMS},
        },
        evaluator::MultiPartyCrs,
        ntt::NttBackendU64,
        parameters::I_2P,
        random::DefaultSecureRng,
    };

    #[test]
    fn mp_noise() {
        // set_parameter_set(&SMALL_MP_BOOL_PARAMS);
        let mut evaluator = BoolEvaluator::<
            Vec<Vec<u64>>,
            NttBackendU64,
            ModularOpsU64<CiphertextModulus<u64>>,
            ModulusPowerOf2<CiphertextModulus<u64>>,
            ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>,
        >::new(I_2P);

        let parties = 2;

        let cr_seed = MultiPartyCrs::random();

        let cks = (0..parties)
            .into_iter()
            .map(|_| evaluator.client_key())
            .collect_vec();

        // construct ideal rlwe sk for meauring noise
        let mut ideal_rlwe_sk = vec![0i32; evaluator.parameters().rlwe_n().0];
        cks.iter().for_each(|k| {
            izip!(
                ideal_rlwe_sk.iter_mut(),
                InteractiveMultiPartyClientKey::sk_rlwe(k)
            )
            .for_each(|(ideal_i, s_i)| {
                *ideal_i = *ideal_i + s_i;
            });
        });
        let mut ideal_lwe_sk = vec![0i32; evaluator.parameters().lwe_n().0];
        cks.iter().for_each(|k| {
            izip!(
                ideal_lwe_sk.iter_mut(),
                InteractiveMultiPartyClientKey::sk_lwe(k)
            )
            .for_each(|(ideal_i, s_i)| {
                *ideal_i = *ideal_i + s_i;
            });
        });

        // round 1
        let pk_shares = cks
            .iter()
            .map(|c| evaluator.multi_party_public_key_share(&cr_seed, c))
            .collect_vec();

        // public key
        let pk = PublicKey::<Vec<Vec<u64>>, DefaultSecureRng, ModularOpsU64<CiphertextModulus<u64>>>::from(
            pk_shares.as_slice(),
        );

        // round 2
        let server_key_shares = cks
            .iter()
            .enumerate()
            .map(|(index, c)| {
                evaluator.multi_party_server_key_share(index, parties, &cr_seed, &pk.key(), c)
            })
            .collect_vec();

        let server_key = evaluator.aggregate_multi_party_server_key_shares(&server_key_shares);
        let runtime_server_key = ShoupServerKeyEvaluationDomain::from(ServerKeyEvaluationDomain::<
            _,
            _,
            DefaultSecureRng,
            NttBackendU64,
        >::from(&server_key));

        let mut m0 = false;
        let mut m1 = true;

        let mut c_m0 = evaluator.pk_encrypt(pk.key(), m0);
        let mut c_m1 = evaluator.pk_encrypt(pk.key(), m1);

        // let mut stats = Stats::new();

        for _ in 0..1000 {
            let now = std::time::Instant::now();
            let c_out = evaluator.xor(&c_m0, &c_m1, &runtime_server_key);
            println!("Gate time: {:?}", now.elapsed());

            // mp decrypt
            let decryption_shares = cks
                .iter()
                .map(|c| evaluator.multi_party_decryption_share(&c_out, c))
                .collect_vec();
            let m_out = evaluator.multi_party_decrypt(&decryption_shares, &c_out);
            let m_expected = (m0 ^ m1);
            assert_eq!(m_expected, m_out, "Expected {m_expected} but got {m_out}");

            // // find noise update
            // {
            //     let out = decrypt_lwe(
            //         &c_out,
            //         ideal_client_key.sk_rlwe().values(),
            //         evaluator.pbs_info().modop_rlweq(),
            //     );

            //     let out_want = {
            //         if m_expected == true {
            //             true_el_encoded
            //         } else {
            //             false_el_encoded
            //         }
            //     };
            //     let diff = evaluator.pbs_info().modop_rlweq().sub(&out, &out_want);

            //     stats.add_more(&vec![evaluator
            //         .pbs_info()
            //         .rlwe_q()
            //         .map_element_to_i64(&diff)]);
            // }

            m1 = m0;
            m0 = m_expected;

            c_m1 = c_m0;
            c_m0 = c_out;
        }

        // println!("log2 std dev {}", stats.std_dev().abs().log2());
    }
 }
--- a/src/bool/print_noise.rs
+++ b/src/bool/print_noise.rs
@ -371,7 +371,7 @@ mod tests {
        use crate::{
            aggregate_public_key_shares, aggregate_server_key_shares,
            bool::keys::{key_size::KeySize, ServerKeyEvaluationDomain},
            evaluator::MultiPartyCrs,
            evaluator::InteractiveMultiPartyCrs,
            gen_client_key, gen_mp_keys_phase1, gen_mp_keys_phase2,
            parameters::CiphertextModulus,
            random::DefaultSecureRng,
@ -381,7 +381,7 @@ mod tests {
        };

        set_parameter_set(crate::ParameterSelector::InteractiveLTE2Party);
        set_mp_seed(MultiPartyCrs::random().seed);
        set_mp_seed(InteractiveMultiPartyCrs::random().seed);
        let parties = 2;
        let cks = (0..parties).map(|_| gen_client_key()).collect_vec();
        let pk_shares = cks.iter().map(|k| gen_mp_keys_phase1(k)).collect_vec();
--- a/src/lib.rs
+++ b/src/lib.rs
@ -1,7 +1,4 @@
 use std::{iter::Once, sync::OnceLock};

 use itertools::{izip, Itertools};
 use num_traits::{abs, Zero};
 use num_traits::Zero;

 mod backend;
 mod bool;
--- a/src/lwe.rs
+++ b/src/lwe.rs
@ -43,7 +43,7 @@ pub(crate) fn lwe_key_switch<
    lwe_out.as_mut()[0] = out_b;
 }

 pub fn seeded_lwe_ksk_keygen<
 pub(crate) fn seeded_lwe_ksk_keygen<
    Ro: RowMut + RowEntity,
    S,
    Op: VectorOps<Element = Ro::Element>
@ -101,7 +101,7 @@ where
 }

 /// Encrypts encoded message m as LWE ciphertext
 pub fn encrypt_lwe<
 pub(crate) fn encrypt_lwe<
    Ro: RowMut + RowEntity,
    Op: ArithmeticOps<Element = Ro::Element> + GetModulus<Element = Ro::Element>,
    R: RandomGaussianElementInModulus<Ro::Element, Op::M>
--- a/src/shortint/ops.rs
+++ b/src/shortint/ops.rs
@ -196,7 +196,7 @@ where
    let n = a.len();
    let neg_b = b.iter().map(|v| evaluator.not(v)).collect_vec();

    // Both remainder and quotient are initially stored in Big-endian in contract to
    // Both remainder and quotient are initially stored in Big-endian in contrast to
    // the usual little endian we use. This is more friendly to vec pushes in
    // division. After computing remainder and quotient, we simply reverse the
    // vectors.