move keys into different file

src/bool/keys.rs

@@ -0,0 +1,661 @@
+use std::{collections::HashMap, hash::Hash, marker::PhantomData};
+
+use crate::{
+    backend::{ModInit, VectorOps},
+    lwe::LweSecret,
+    random::{NewWithSeed, RandomFillUniformInModulus},
+    rgsw::RlweSecret,
+    utils::WithLocal,
+    Decryptor, Encryptor, Matrix, MatrixEntity, MatrixMut, MultiPartyDecryptor, RowEntity, RowMut,
+};
+
+use super::{parameters, BoolEvaluator, BoolParameters, CiphertextModulus};
+
+/// Client key with RLWE and LWE secrets
+#[derive(Clone)]
+pub struct ClientKey {
+    sk_rlwe: RlweSecret,
+    sk_lwe: LweSecret,
+}
+
+mod impl_ck {
+    use super::*;
+
+    // Client key
+    impl ClientKey {
+        pub(in super::super) fn random() -> Self {
+            let sk_rlwe = RlweSecret::random(0, 0);
+            let sk_lwe = LweSecret::random(0, 0);
+            Self { sk_rlwe, sk_lwe }
+        }
+
+        pub(in super::super) fn new(sk_rlwe: RlweSecret, sk_lwe: LweSecret) -> Self {
+            Self { sk_rlwe, sk_lwe }
+        }
+
+        pub(in super::super) fn sk_rlwe(&self) -> &RlweSecret {
+            &self.sk_rlwe
+        }
+
+        pub(in super::super) fn sk_lwe(&self) -> &LweSecret {
+            &self.sk_lwe
+        }
+    }
+
+    impl Encryptor<bool, Vec<u64>> for ClientKey {
+        fn encrypt(&self, m: &bool) -> Vec<u64> {
+            BoolEvaluator::with_local(|e| e.sk_encrypt(*m, self))
+        }
+    }
+
+    impl Decryptor<bool, Vec<u64>> for ClientKey {
+        fn decrypt(&self, c: &Vec<u64>) -> bool {
+            BoolEvaluator::with_local(|e| e.sk_decrypt(c, self))
+        }
+    }
+
+    impl MultiPartyDecryptor<bool, Vec<u64>> for ClientKey {
+        type DecryptionShare = u64;
+
+        fn gen_decryption_share(&self, c: &Vec<u64>) -> Self::DecryptionShare {
+            BoolEvaluator::with_local(|e| e.multi_party_decryption_share(c, &self))
+        }
+
+        fn aggregate_decryption_shares(
+            &self,
+            c: &Vec<u64>,
+            shares: &[Self::DecryptionShare],
+        ) -> bool {
+            BoolEvaluator::with_local(|e| e.multi_party_decrypt(shares, c))
+        }
+    }
+}
+
+/// Public key
+pub struct PublicKey<M, Rng, ModOp> {
+    key: M,
+    _phantom: PhantomData<(Rng, ModOp)>,
+}
+
+pub(super) mod impl_pk {
+    use super::*;
+
+    impl<M, R, Mo> PublicKey<M, R, Mo> {
+        pub(in super::super) fn key(&self) -> &M {
+            &self.key
+        }
+    }
+
+    impl<Rng, ModOp> Encryptor<bool, Vec<u64>> for PublicKey<Vec<Vec<u64>>, Rng, ModOp> {
+        fn encrypt(&self, m: &bool) -> Vec<u64> {
+            BoolEvaluator::with_local(|e| e.pk_encrypt(&self.key, *m))
+        }
+    }
+
+    impl<Rng, ModOp> Encryptor<[bool], Vec<Vec<u64>>> for PublicKey<Vec<Vec<u64>>, Rng, ModOp> {
+        fn encrypt(&self, m: &[bool]) -> Vec<Vec<u64>> {
+            BoolEvaluator::with_local(|e| e.pk_encrypt_batched(&self.key, m))
+        }
+    }
+
+    impl<
+            M: MatrixMut + MatrixEntity,
+            Rng: NewWithSeed
+                + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
+            ModOp,
+        > From<SeededPublicKey<M::R, Rng::Seed, BoolParameters<M::MatElement>, ModOp>>
+        for PublicKey<M, Rng, ModOp>
+    where
+        <M as Matrix>::R: RowMut,
+        M::MatElement: Copy,
+    {
+        fn from(
+            value: SeededPublicKey<M::R, Rng::Seed, BoolParameters<M::MatElement>, ModOp>,
+        ) -> Self {
+            let mut prng = Rng::new_with_seed(value.seed);
+
+            let mut key = M::zeros(2, value.parameters.rlwe_n().0);
+            // sample A
+            RandomFillUniformInModulus::random_fill(
+                &mut prng,
+                value.parameters.rlwe_q(),
+                key.get_row_mut(0),
+            );
+            // Copy over B
+            key.get_row_mut(1).copy_from_slice(value.part_b.as_ref());
+
+            PublicKey {
+                key,
+                _phantom: PhantomData,
+            }
+        }
+    }
+
+    impl<
+            M: MatrixMut + MatrixEntity,
+            Rng: NewWithSeed
+                + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
+            ModOp: VectorOps<Element = M::MatElement> + ModInit<M = CiphertextModulus<M::MatElement>>,
+        >
+        From<
+            &[CommonReferenceSeededCollectivePublicKeyShare<
+                M::R,
+                Rng::Seed,
+                BoolParameters<M::MatElement>,
+            >],
+        > for PublicKey<M, Rng, ModOp>
+    where
+        <M as Matrix>::R: RowMut,
+        Rng::Seed: Copy + PartialEq,
+        M::MatElement: PartialEq + Copy,
+    {
+        fn from(
+            value: &[CommonReferenceSeededCollectivePublicKeyShare<
+                M::R,
+                Rng::Seed,
+                BoolParameters<M::MatElement>,
+            >],
+        ) -> Self {
+            assert!(value.len() > 0);
+
+            let parameters = &value[0].parameters;
+            let mut key = M::zeros(2, parameters.rlwe_n().0);
+
+            // sample A
+            let seed = value[0].cr_seed;
+            let mut main_rng = Rng::new_with_seed(seed);
+            RandomFillUniformInModulus::random_fill(
+                &mut main_rng,
+                parameters.rlwe_q(),
+                key.get_row_mut(0),
+            );
+
+            // Sum all Bs
+            let rlweq_modop = ModOp::new(parameters.rlwe_q().clone());
+            value.iter().for_each(|share_i| {
+                assert!(share_i.cr_seed == seed);
+                assert!(&share_i.parameters == parameters);
+
+                rlweq_modop.elwise_add_mut(key.get_row_mut(1), share_i.share.as_ref());
+            });
+
+            PublicKey {
+                key,
+                _phantom: PhantomData,
+            }
+        }
+    }
+}
+
+/// Seeded public key
+struct SeededPublicKey<Ro, S, P, ModOp> {
+    part_b: Ro,
+    seed: S,
+    parameters: P,
+    _phantom: PhantomData<ModOp>,
+}
+
+mod impl_seeded_pk {
+    use super::*;
+
+    impl<R, S, ModOp>
+        From<&[CommonReferenceSeededCollectivePublicKeyShare<R, S, BoolParameters<R::Element>>]>
+        for SeededPublicKey<R, S, BoolParameters<R::Element>, ModOp>
+    where
+        ModOp: VectorOps<Element = R::Element> + ModInit<M = CiphertextModulus<R::Element>>,
+        S: PartialEq + Clone,
+        R: RowMut + RowEntity + Clone,
+        R::Element: Clone + PartialEq,
+    {
+        fn from(
+            value: &[CommonReferenceSeededCollectivePublicKeyShare<
+                R,
+                S,
+                BoolParameters<R::Element>,
+            >],
+        ) -> Self {
+            assert!(value.len() > 0);
+
+            let parameters = &value[0].parameters;
+            let cr_seed = value[0].cr_seed.clone();
+
+            // Sum all Bs
+            let rlweq_modop = ModOp::new(parameters.rlwe_q().clone());
+            let mut part_b = value[0].share.clone();
+            value.iter().skip(1).for_each(|share_i| {
+                assert!(&share_i.cr_seed == &cr_seed);
+                assert!(&share_i.parameters == parameters);
+
+                rlweq_modop.elwise_add_mut(part_b.as_mut(), share_i.share.as_ref());
+            });
+
+            Self {
+                part_b,
+                seed: cr_seed,
+                parameters: parameters.clone(),
+                _phantom: PhantomData,
+            }
+        }
+    }
+}
+
+/// CRS seeded collective public key share
+pub struct CommonReferenceSeededCollectivePublicKeyShare<Ro, S, P> {
+    share: Ro,
+    cr_seed: S,
+    parameters: P,
+}
+impl<Ro, S, P> CommonReferenceSeededCollectivePublicKeyShare<Ro, S, P> {
+    pub(super) fn new(share: Ro, cr_seed: S, parameters: P) -> Self {
+        CommonReferenceSeededCollectivePublicKeyShare {
+            share,
+            cr_seed,
+            parameters,
+        }
+    }
+}
+
+/// CRS seeded Multi-party server key share
+pub struct CommonReferenceSeededMultiPartyServerKeyShare<M: Matrix, P, S> {
+    rgsw_cts: Vec<M>,
+    /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
+    /// to -g is at 0
+    auto_keys: HashMap<usize, M>,
+    lwe_ksk: M::R,
+    /// Common reference seed
+    cr_seed: S,
+    parameters: P,
+}
+
+impl<M: Matrix, P, S> CommonReferenceSeededMultiPartyServerKeyShare<M, P, S> {
+    pub(super) fn new(
+        rgsw_cts: Vec<M>,
+        auto_keys: HashMap<usize, M>,
+        lwe_ksk: M::R,
+        cr_seed: S,
+        parameters: P,
+    ) -> Self {
+        CommonReferenceSeededMultiPartyServerKeyShare {
+            rgsw_cts,
+            auto_keys,
+            lwe_ksk,
+            cr_seed,
+            parameters,
+        }
+    }
+
+    pub(super) fn cr_seed(&self) -> &S {
+        &self.cr_seed
+    }
+
+    pub(super) fn parameters(&self) -> &P {
+        &self.parameters
+    }
+
+    pub(super) fn auto_keys(&self) -> &HashMap<usize, M> {
+        &self.auto_keys
+    }
+
+    pub(super) fn rgsw_cts(&self) -> &[M] {
+        &self.rgsw_cts
+    }
+
+    pub(super) fn lwe_ksk(&self) -> &M::R {
+        &self.lwe_ksk
+    }
+}
+
+/// CRS seeded MultiParty server key
+pub struct SeededMultiPartyServerKey<M: Matrix, S, P> {
+    rgsw_cts: Vec<M>,
+    /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
+    /// to -g is at 0
+    auto_keys: HashMap<usize, M>,
+    lwe_ksk: M::R,
+    cr_seed: S,
+    parameters: P,
+}
+
+impl<M: Matrix, S, P> SeededMultiPartyServerKey<M, S, P> {
+    pub(super) fn new(
+        rgsw_cts: Vec<M>,
+        auto_keys: HashMap<usize, M>,
+        lwe_ksk: M::R,
+        cr_seed: S,
+        parameters: P,
+    ) -> Self {
+        SeededMultiPartyServerKey {
+            rgsw_cts,
+            auto_keys,
+            lwe_ksk,
+            cr_seed,
+            parameters,
+        }
+    }
+
+    pub(super) fn rgsw_cts(&self) -> &[M] {
+        &self.rgsw_cts
+    }
+}
+
+/// Seeded single party server key
+pub struct SeededServerKey<M: Matrix, P, S> {
+    /// Rgsw cts of LWE secret elements
+    pub(crate) rgsw_cts: Vec<M>,
+    /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
+    /// to -g is at 0
+    pub(crate) auto_keys: HashMap<usize, M>,
+    /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
+    pub(crate) lwe_ksk: M::R,
+    /// Parameters
+    pub(crate) parameters: P,
+    /// Main seed
+    pub(crate) seed: S,
+}
+impl<M: Matrix, S> SeededServerKey<M, BoolParameters<M::MatElement>, S> {
+    pub(super) fn from_raw(
+        auto_keys: HashMap<usize, M>,
+        rgsw_cts: Vec<M>,
+        lwe_ksk: M::R,
+        parameters: BoolParameters<M::MatElement>,
+        seed: S,
+    ) -> Self {
+        // sanity checks
+        auto_keys.iter().for_each(|v| {
+            assert!(
+                v.1.dimension()
+                    == (
+                        parameters.auto_decomposition_count().0,
+                        parameters.rlwe_n().0
+                    )
+            )
+        });
+
+        let (part_a_d, part_b_d) = parameters.rlwe_rgsw_decomposition_count();
+        rgsw_cts.iter().for_each(|v| {
+            assert!(v.dimension() == (part_a_d.0 * 2 + part_b_d.0, parameters.rlwe_n().0))
+        });
+        assert!(
+            lwe_ksk.as_ref().len()
+                == (parameters.lwe_decomposition_count().0 * parameters.rlwe_n().0)
+        );
+
+        SeededServerKey {
+            rgsw_cts,
+            auto_keys,
+            lwe_ksk,
+            parameters,
+            seed,
+        }
+    }
+}
+
+/// Server key in evaluation domain
+pub(crate) struct ServerKeyEvaluationDomain<M, R, N> {
+    /// Rgsw cts of LWE secret elements
+    rgsw_cts: Vec<M>,
+    /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
+    /// to -g is at 0
+    galois_keys: HashMap<usize, M>,
+    /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
+    lwe_ksk: M,
+    _phanton: PhantomData<(R, N)>,
+}
+
+pub(super) mod impl_server_key_eval_domain {
+    use itertools::{izip, Itertools};
+
+    use crate::{
+        ntt::{Ntt, NttInit},
+        pbs::PbsKey,
+    };
+
+    use super::*;
+
+    impl<M, R, N> ServerKeyEvaluationDomain<M, R, N> {
+        pub(in super::super) fn rgsw_cts(&self) -> &[M] {
+            &self.rgsw_cts
+        }
+    }
+
+    impl<
+            M: MatrixMut + MatrixEntity,
+            R: RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>
+                + NewWithSeed,
+            N: NttInit<CiphertextModulus<M::MatElement>> + Ntt<Element = M::MatElement>,
+        > From<&SeededServerKey<M, BoolParameters<M::MatElement>, R::Seed>>
+        for ServerKeyEvaluationDomain<M, R, N>
+    where
+        <M as Matrix>::R: RowMut,
+        M::MatElement: Copy,
+        R::Seed: Clone,
+    {
+        fn from(value: &SeededServerKey<M, BoolParameters<M::MatElement>, R::Seed>) -> Self {
+            let mut main_prng = R::new_with_seed(value.seed.clone());
+            let parameters = &value.parameters;
+            let g = parameters.g() as isize;
+            let ring_size = value.parameters.rlwe_n().0;
+            let lwe_n = value.parameters.lwe_n().0;
+            let rlwe_q = value.parameters.rlwe_q();
+            let lwq_q = value.parameters.lwe_q();
+
+            let nttop = N::new(rlwe_q, ring_size);
+
+            // galois keys
+            let mut auto_keys = HashMap::new();
+            let auto_decomp_count = parameters.auto_decomposition_count().0;
+            let auto_element_dlogs = parameters.auto_element_dlogs();
+            for i in auto_element_dlogs.into_iter() {
+                let seeded_auto_key = value.auto_keys.get(&i).unwrap();
+                assert!(seeded_auto_key.dimension() == (auto_decomp_count, ring_size));
+
+                let mut data = M::zeros(auto_decomp_count * 2, ring_size);
+
+                // sample RLWE'_A(-s(X^k))
+                data.iter_rows_mut().take(auto_decomp_count).for_each(|ri| {
+                    RandomFillUniformInModulus::random_fill(&mut main_prng, &rlwe_q, ri.as_mut())
+                });
+
+                // copy over RLWE'B_(-s(X^k))
+                izip!(
+                    data.iter_rows_mut().skip(auto_decomp_count),
+                    seeded_auto_key.iter_rows()
+                )
+                .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
+
+                // Send to Evaluation domain
+                data.iter_rows_mut()
+                    .for_each(|ri| nttop.forward(ri.as_mut()));
+
+                auto_keys.insert(i, data);
+            }
+
+            // RGSW ciphertexts
+            let (rlrg_a_decomp, rlrg_b_decomp) = parameters.rlwe_rgsw_decomposition_count();
+            let rgsw_cts = value
+                .rgsw_cts
+                .iter()
+                .map(|seeded_rgsw_si| {
+                    assert!(
+                        seeded_rgsw_si.dimension()
+                            == (rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0, ring_size)
+                    );
+
+                    let mut data = M::zeros(rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0 * 2, ring_size);
+
+                    // copy over RLWE'(-sm)
+                    izip!(
+                        data.iter_rows_mut().take(rlrg_a_decomp.0 * 2),
+                        seeded_rgsw_si.iter_rows().take(rlrg_a_decomp.0 * 2)
+                    )
+                    .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
+
+                    // sample RLWE'_A(m)
+                    data.iter_rows_mut()
+                        .skip(rlrg_a_decomp.0 * 2)
+                        .take(rlrg_b_decomp.0)
+                        .for_each(|ri| {
+                            RandomFillUniformInModulus::random_fill(
+                                &mut main_prng,
+                                &rlwe_q,
+                                ri.as_mut(),
+                            )
+                        });
+
+                    // copy over RLWE'_B(m)
+                    izip!(
+                        data.iter_rows_mut()
+                            .skip(rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0),
+                        seeded_rgsw_si.iter_rows().skip(rlrg_a_decomp.0 * 2)
+                    )
+                    .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
+
+                    // send polynomials to evaluation domain
+                    data.iter_rows_mut()
+                        .for_each(|ri| nttop.forward(ri.as_mut()));
+
+                    data
+                })
+                .collect_vec();
+
+            // LWE ksk
+            let lwe_ksk = {
+                let d = parameters.lwe_decomposition_count().0;
+                assert!(value.lwe_ksk.as_ref().len() == d * ring_size);
+
+                let mut data = M::zeros(d * ring_size, lwe_n + 1);
+                izip!(data.iter_rows_mut(), value.lwe_ksk.as_ref().iter()).for_each(
+                    |(lwe_i, bi)| {
+                        RandomFillUniformInModulus::random_fill(
+                            &mut main_prng,
+                            &lwq_q,
+                            &mut lwe_i.as_mut()[1..],
+                        );
+                        lwe_i.as_mut()[0] = *bi;
+                    },
+                );
+
+                data
+            };
+
+            ServerKeyEvaluationDomain {
+                rgsw_cts,
+                galois_keys: auto_keys,
+                lwe_ksk,
+                _phanton: PhantomData,
+            }
+        }
+    }
+
+    impl<
+            M: MatrixMut + MatrixEntity,
+            Rng: NewWithSeed,
+            N: NttInit<CiphertextModulus<M::MatElement>> + Ntt<Element = M::MatElement>,
+        > From<&SeededMultiPartyServerKey<M, Rng::Seed, BoolParameters<M::MatElement>>>
+        for ServerKeyEvaluationDomain<M, Rng, N>
+    where
+        <M as Matrix>::R: RowMut,
+        Rng::Seed: Copy,
+        Rng: RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
+        M::MatElement: Copy,
+    {
+        fn from(
+            value: &SeededMultiPartyServerKey<M, Rng::Seed, BoolParameters<M::MatElement>>,
+        ) -> Self {
+            let g = value.parameters.g() as isize;
+            let rlwe_n = value.parameters.rlwe_n().0;
+            let lwe_n = value.parameters.lwe_n().0;
+            let rlwe_q = value.parameters.rlwe_q();
+            let lwe_q = value.parameters.lwe_q();
+
+            let mut main_prng = Rng::new_with_seed(value.cr_seed);
+
+            let rlwe_nttop = N::new(rlwe_q, rlwe_n);
+
+            // auto keys
+            let mut auto_keys = HashMap::new();
+            let auto_d_count = value.parameters.auto_decomposition_count().0;
+            let auto_element_dlogs = value.parameters.auto_element_dlogs();
+            for i in auto_element_dlogs.into_iter() {
+                let mut key = M::zeros(auto_d_count * 2, rlwe_n);
+
+                // sample a
+                key.iter_rows_mut().take(auto_d_count).for_each(|ri| {
+                    RandomFillUniformInModulus::random_fill(&mut main_prng, &rlwe_q, ri.as_mut())
+                });
+
+                let key_part_b = value.auto_keys.get(&i).unwrap();
+                assert!(key_part_b.dimension() == (auto_d_count, rlwe_n));
+                izip!(
+                    key.iter_rows_mut().skip(auto_d_count),
+                    key_part_b.iter_rows()
+                )
+                .for_each(|(to_ri, from_ri)| {
+                    to_ri.as_mut().copy_from_slice(from_ri.as_ref());
+                });
+
+                // send to evaluation domain
+                key.iter_rows_mut()
+                    .for_each(|ri| rlwe_nttop.forward(ri.as_mut()));
+
+                auto_keys.insert(i, key);
+            }
+
+            // rgsw cts
+            let (rlrg_d_a, rlrg_d_b) = value.parameters.rlwe_rgsw_decomposition_count();
+            let rgsw_ct_out = rlrg_d_a.0 * 2 + rlrg_d_b.0 * 2;
+            let rgsw_cts = value
+                .rgsw_cts
+                .iter()
+                .map(|ct_i_in| {
+                    assert!(ct_i_in.dimension() == (rgsw_ct_out, rlwe_n));
+                    let mut eval_ct_i_out = M::zeros(rgsw_ct_out, rlwe_n);
+
+                    izip!(eval_ct_i_out.iter_rows_mut(), ct_i_in.iter_rows()).for_each(
+                        |(to_ri, from_ri)| {
+                            to_ri.as_mut().copy_from_slice(from_ri.as_ref());
+                            rlwe_nttop.forward(to_ri.as_mut());
+                        },
+                    );
+
+                    eval_ct_i_out
+                })
+                .collect_vec();
+
+            // lwe ksk
+            let d_lwe = value.parameters.lwe_decomposition_count().0;
+            let mut lwe_ksk = M::zeros(rlwe_n * d_lwe, lwe_n + 1);
+            izip!(lwe_ksk.iter_rows_mut(), value.lwe_ksk.as_ref().iter()).for_each(
+                |(lwe_i, bi)| {
+                    RandomFillUniformInModulus::random_fill(
+                        &mut main_prng,
+                        &lwe_q,
+                        &mut lwe_i.as_mut()[1..],
+                    );
+                    lwe_i.as_mut()[0] = *bi;
+                },
+            );
+
+            ServerKeyEvaluationDomain {
+                rgsw_cts,
+                galois_keys: auto_keys,
+                lwe_ksk,
+                _phanton: PhantomData,
+            }
+        }
+    }
+
+    impl<M: Matrix, R, N> PbsKey for ServerKeyEvaluationDomain<M, R, N> {
+        type M = M;
+        fn galois_key_for_auto(&self, k: usize) -> &Self::M {
+            self.galois_keys.get(&k).unwrap()
+        }
+        fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::M {
+            &self.rgsw_cts[si]
+        }
+
+        fn lwe_ksk(&self) -> &Self::M {
+            &self.lwe_ksk
+        }
+    }
+}

src/bool/mod.rs

@@ -1,4 +1,176 @@
 pub(crate) mod evaluator;
+pub(crate) mod keys;
 pub(crate) mod parameters;
 
 pub type FheBool = Vec<u64>;
+
+use std::{cell::RefCell, sync::OnceLock};
+
+use evaluator::*;
+use keys::*;
+use parameters::*;
+
+use crate::{
+    backend::ModularOpsU64,
+    ntt::NttBackendU64,
+    random::{DefaultSecureRng, NewWithSeed},
+    utils::{Global, WithLocal},
+};
+
+thread_local! {
+    static BOOL_EVALUATOR: RefCell<BoolEvaluator<Vec<Vec<u64>>, NttBackendU64, ModularOpsU64<CiphertextModulus<u64>>,  ModularOpsU64<CiphertextModulus<u64>>>> = RefCell::new(BoolEvaluator::new(MP_BOOL_PARAMS));
+
+}
+static BOOL_SERVER_KEY: OnceLock<
+    ServerKeyEvaluationDomain<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64>,
+> = OnceLock::new();
+
+static MULTI_PARTY_CRS: OnceLock<MultiPartyCrs<[u8; 32]>> = OnceLock::new();
+
+pub fn set_parameter_set(parameter: &BoolParameters<u64>) {
+    BoolEvaluator::with_local_mut(|e| *e = BoolEvaluator::new(parameter.clone()))
+}
+
+pub fn set_mp_seed(seed: [u8; 32]) {
+    assert!(
+        MULTI_PARTY_CRS.set(MultiPartyCrs { seed: seed }).is_ok(),
+        "Attempted to set MP SEED twice."
+    )
+}
+
+fn set_server_key(key: ServerKeyEvaluationDomain<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64>) {
+    assert!(
+        BOOL_SERVER_KEY.set(key).is_ok(),
+        "Attempted to set server key twice."
+    );
+}
+
+pub(crate) fn gen_keys() -> (
+    ClientKey,
+    SeededServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]>,
+) {
+    BoolEvaluator::with_local_mut(|e| {
+        let ck = e.client_key();
+        let sk = e.server_key(&ck);
+
+        (ck, sk)
+    })
+}
+
+pub fn gen_client_key() -> ClientKey {
+    BoolEvaluator::with_local(|e| e.client_key())
+}
+
+pub fn gen_mp_keys_phase1(
+    ck: &ClientKey,
+) -> CommonReferenceSeededCollectivePublicKeyShare<Vec<u64>, [u8; 32], BoolParameters<u64>> {
+    let seed = MultiPartyCrs::global().public_key_share_seed::<DefaultSecureRng>();
+    BoolEvaluator::with_local(|e| {
+        let pk_share = e.multi_party_public_key_share(seed, &ck);
+        pk_share
+    })
+}
+
+pub fn gen_mp_keys_phase2<R, ModOp>(
+    ck: &ClientKey,
+    pk: &PublicKey<Vec<Vec<u64>>, R, ModOp>,
+) -> CommonReferenceSeededMultiPartyServerKeyShare<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> {
+    let seed = MultiPartyCrs::global().server_key_share_seed::<DefaultSecureRng>();
+    BoolEvaluator::with_local_mut(|e| {
+        let server_key_share = e.multi_party_server_key_share(seed, pk.key(), ck);
+        server_key_share
+    })
+}
+
+pub fn aggregate_public_key_shares(
+    shares: &[CommonReferenceSeededCollectivePublicKeyShare<
+        Vec<u64>,
+        [u8; 32],
+        BoolParameters<u64>,
+    >],
+) -> PublicKey<Vec<Vec<u64>>, DefaultSecureRng, ModularOpsU64<CiphertextModulus<u64>>> {
+    PublicKey::from(shares)
+}
+
+pub fn aggregate_server_key_shares(
+    shares: &[CommonReferenceSeededMultiPartyServerKeyShare<
+        Vec<Vec<u64>>,
+        BoolParameters<u64>,
+        [u8; 32],
+    >],
+) -> SeededMultiPartyServerKey<Vec<Vec<u64>>, [u8; 32], BoolParameters<u64>> {
+    BoolEvaluator::with_local(|e| e.aggregate_multi_party_server_key_shares(shares))
+}
+
+// SERVER KEY EVAL DOMAIN //
+impl SeededServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> {
+    pub fn set_server_key(&self) {
+        set_server_key(ServerKeyEvaluationDomain::<
+            _,
+            DefaultSecureRng,
+            NttBackendU64,
+        >::from(self));
+    }
+}
+
+impl
+    SeededMultiPartyServerKey<
+        Vec<Vec<u64>>,
+        <DefaultSecureRng as NewWithSeed>::Seed,
+        BoolParameters<u64>,
+    >
+{
+    pub fn set_server_key(&self) {
+        set_server_key(ServerKeyEvaluationDomain::<
+            Vec<Vec<u64>>,
+            DefaultSecureRng,
+            NttBackendU64,
+        >::from(self))
+    }
+}
+
+impl Global for ServerKeyEvaluationDomain<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64> {
+    fn global() -> &'static Self {
+        BOOL_SERVER_KEY.get().unwrap()
+    }
+}
+
+//  MULTIPARTY CRS  //
+impl Global for MultiPartyCrs<[u8; 32]> {
+    fn global() -> &'static Self {
+        MULTI_PARTY_CRS
+            .get()
+            .expect("Multi Party Common Reference String not set")
+    }
+}
+
+// BOOL EVALUATOR //
+impl WithLocal
+    for BoolEvaluator<
+        Vec<Vec<u64>>,
+        NttBackendU64,
+        ModularOpsU64<CiphertextModulus<u64>>,
+        ModularOpsU64<CiphertextModulus<u64>>,
+    >
+{
+    fn with_local<F, R>(func: F) -> R
+    where
+        F: Fn(&Self) -> R,
+    {
+        BOOL_EVALUATOR.with_borrow(|s| func(s))
+    }
+
+    fn with_local_mut<F, R>(func: F) -> R
+    where
+        F: Fn(&mut Self) -> R,
+    {
+        BOOL_EVALUATOR.with_borrow_mut(|s| func(s))
+    }
+
+    fn with_local_mut_mut<F, R>(func: &mut F) -> R
+    where
+        F: FnMut(&mut Self) -> R,
+    {
+        BOOL_EVALUATOR.with_borrow_mut(|s| func(s))
+    }
+}

src/bool/parameters.rs

@@ -1,9 +1,9 @@
-use num_traits::{ConstZero, FromPrimitive, PrimInt, ToPrimitive, Zero};
+use num_traits::{ConstZero, FromPrimitive, PrimInt};
 
 use crate::{backend::Modulus, decomposer::Decomposer};
 
 #[derive(Clone, PartialEq)]
-pub struct BoolParameters<El> {
+pub(crate) struct BoolParameters<El> {
     rlwe_q: CiphertextModulus<El>,
     lwe_q: CiphertextModulus<El>,
     br_q: usize,

src/pbs.rs

@@ -214,7 +214,7 @@ fn blind_rotation<
 >(
     trivial_rlwe_test_poly: &mut MT,
     scratch_matrix: &mut Mmut,
-    g: isize,
+    _g: isize,
     w: usize,
     q: usize,
     gk_to_si: &[Vec<usize>],

src/shortint/mod.rs

@@ -1,11 +1,8 @@
 use itertools::Itertools;
 
 use crate::{
-    bool::evaluator::{
-        BoolEvaluator, ClientKey, PublicKey, ServerKeyEvaluationDomain, BOOL_SERVER_KEY,
-    },
-    utils::{Global, WithLocal},
-    Decryptor, Encryptor, Matrix, MultiPartyDecryptor,
+    bool::keys::{ClientKey, PublicKey},
+    Decryptor, Encryptor, MultiPartyDecryptor,
 };
 
 mod ops;
@@ -100,7 +97,7 @@ mod frontend {
         eight_bit_mul,
     };
     use crate::{
-        bool::evaluator::{BoolEvaluator, ServerKeyEvaluationDomain},
+        bool::{evaluator::BoolEvaluator, keys::ServerKeyEvaluationDomain},
         utils::{Global, WithLocal},
     };
 
@@ -307,12 +304,10 @@ mod tests {
 
     use crate::{
         bool::{
-            evaluator::{
-                aggregate_public_key_shares, aggregate_server_key_shares, gen_client_key, gen_keys,
-                gen_mp_keys_phase1, gen_mp_keys_phase2, set_mp_seed, set_parameter_set,
-                BoolEvaluator, ClientKey,
-            },
+            aggregate_public_key_shares, aggregate_server_key_shares, gen_client_key, gen_keys,
+            gen_mp_keys_phase1, gen_mp_keys_phase2,
             parameters::{MP_BOOL_PARAMS, SP_BOOL_PARAMS},
+            set_mp_seed, set_parameter_set,
         },
         shortint::types::FheUint8,
         Decryptor, Encryptor, MultiPartyDecryptor,

src/shortint/ops.rs

@@ -1,18 +1,6 @@
-use std::mem::MaybeUninit;
-
 use itertools::{izip, Itertools};
-use num_traits::PrimInt;
-
-use crate::{
-    backend::ModularOpsU64,
-    bool::{
-        evaluator::{BoolEvaluator, BooleanGates, ClientKey, ServerKeyEvaluationDomain},
-        parameters::CiphertextModulus,
-    },
-    ntt::NttBackendU64,
-    random::DefaultSecureRng,
-    Decryptor,
-};
+
+use crate::bool::evaluator::BooleanGates;
 
 pub(super) fn half_adder<E: BooleanGates>(
     evaluator: &mut E,