1use std::{cell::RefCell, sync::OnceLock};
  2
  3use crate::{
  4    backend::{ModularOpsU64, ModulusPowerOf2},
  5    ntt::NttBackendU64,
  6    random::{DefaultSecureRng, NewWithSeed},
  7    utils::{Global, WithLocal},
  8};
  9
 10use super::{evaluator::*, keys::*, parameters::*};
 11
 12thread_local! {
 13    static BOOL_EVALUATOR: RefCell<Option<BoolEvaluator<Vec<Vec<u64>>, NttBackendU64, ModularOpsU64<CiphertextModulus<u64>>,  ModulusPowerOf2<CiphertextModulus<u64>>, ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>>>> = RefCell::new(None);
 14
 15}
 16static BOOL_SERVER_KEY: OnceLock<ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>> = OnceLock::new();
 17
 18static MULTI_PARTY_CRS: OnceLock<MultiPartyCrs<[u8; 32]>> = OnceLock::new();
 19
 20pub type ClientKey = super::keys::ClientKey<[u8; 32], u64>;
 21
 22pub enum ParameterSelector {
 23    MultiPartyLessThanOrEqualTo16,
 24}
 25
 26pub fn set_parameter_set(select: ParameterSelector) {
 27    match select {
 28        ParameterSelector::MultiPartyLessThanOrEqualTo16 => {
 29            BOOL_EVALUATOR.with_borrow_mut(|v| *v = Some(BoolEvaluator::new(SMALL_MP_BOOL_PARAMS)));
 30        }
 31    }
 32}
 33
 34pub fn set_mp_seed(seed: [u8; 32]) {
 35    assert!(
 36        MULTI_PARTY_CRS.set(MultiPartyCrs { seed: seed }).is_ok(),
 37        "Attempted to set MP SEED twice."
 38    )
 39}
 40
 41fn set_server_key(key: ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>) {
 42    assert!(
 43        BOOL_SERVER_KEY.set(key).is_ok(),
 44        "Attempted to set server key twice."
 45    );
 46}
 47
 48pub(crate) fn gen_keys() -> (
 49    ClientKey,
 50    SeededSinglePartyServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]>,
 51) {
 52    BoolEvaluator::with_local_mut(|e| {
 53        let ck = e.client_key();
 54        let sk = e.single_party_server_key(&ck);
 55
 56        (ck, sk)
 57    })
 58}
 59
 60pub fn gen_client_key() -> ClientKey {
 61    BoolEvaluator::with_local(|e| e.client_key())
 62}
 63
 64pub fn gen_mp_keys_phase1(
 65    ck: &ClientKey,
 66) -> CommonReferenceSeededCollectivePublicKeyShare<Vec<u64>, [u8; 32], BoolParameters<u64>> {
 67    let seed = MultiPartyCrs::global().public_key_share_seed::<DefaultSecureRng>();
 68    BoolEvaluator::with_local(|e| {
 69        let pk_share = e.multi_party_public_key_share(seed, ck);
 70        pk_share
 71    })
 72}
 73
 74pub fn gen_mp_keys_phase2<R, ModOp>(
 75    ck: &ClientKey,
 76    pk: &PublicKey<Vec<Vec<u64>>, R, ModOp>,
 77) -> CommonReferenceSeededMultiPartyServerKeyShare<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> {
 78    let seed = MultiPartyCrs::global().server_key_share_seed::<DefaultSecureRng>();
 79    BoolEvaluator::with_local_mut(|e| {
 80        let server_key_share = e.multi_party_server_key_share(seed, pk.key(), ck);
 81        server_key_share
 82    })
 83}
 84
 85pub fn aggregate_public_key_shares(
 86    shares: &[CommonReferenceSeededCollectivePublicKeyShare<
 87        Vec<u64>,
 88        [u8; 32],
 89        BoolParameters<u64>,
 90    >],
 91) -> PublicKey<Vec<Vec<u64>>, DefaultSecureRng, ModularOpsU64<CiphertextModulus<u64>>> {
 92    PublicKey::from(shares)
 93}
 94
 95pub fn aggregate_server_key_shares(
 96    shares: &[CommonReferenceSeededMultiPartyServerKeyShare<
 97        Vec<Vec<u64>>,
 98        BoolParameters<u64>,
 99        [u8; 32],
100    >],
101) -> SeededMultiPartyServerKey<Vec<Vec<u64>>, [u8; 32], BoolParameters<u64>> {
102    BoolEvaluator::with_local(|e| e.aggregate_multi_party_server_key_shares(shares))
103}
104
105// SERVER KEY EVAL (/SHOUP) DOMAIN //
106impl SeededSinglePartyServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> {
107    pub fn set_server_key(&self) {
108        let eval = ServerKeyEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(self);
109        set_server_key(ShoupServerKeyEvaluationDomain::from(eval));
110    }
111}
112
113impl
114    SeededMultiPartyServerKey<
115        Vec<Vec<u64>>,
116        <DefaultSecureRng as NewWithSeed>::Seed,
117        BoolParameters<u64>,
118    >
119{
120    pub fn set_server_key(&self) {
121        set_server_key(ShoupServerKeyEvaluationDomain::from(
122            ServerKeyEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(self),
123        ))
124    }
125}
126
127//  MULTIPARTY CRS  //
128impl Global for MultiPartyCrs<[u8; 32]> {
129    fn global() -> &'static Self {
130        MULTI_PARTY_CRS
131            .get()
132            .expect("Multi Party Common Reference String not set")
133    }
134}
135
136// BOOL EVALUATOR //
137impl WithLocal
138    for BoolEvaluator<
139        Vec<Vec<u64>>,
140        NttBackendU64,
141        ModularOpsU64<CiphertextModulus<u64>>,
142        ModulusPowerOf2<CiphertextModulus<u64>>,
143        ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>,
144    >
145{
146    fn with_local<F, R>(func: F) -> R
147    where
148        F: Fn(&Self) -> R,
149    {
150        BOOL_EVALUATOR.with_borrow(|s| func(s.as_ref().expect("Parameters not set")))
151    }
152
153    fn with_local_mut<F, R>(func: F) -> R
154    where
155        F: Fn(&mut Self) -> R,
156    {
157        BOOL_EVALUATOR.with_borrow_mut(|s| func(s.as_mut().expect("Parameters not set")))
158    }
159
160    fn with_local_mut_mut<F, R>(func: &mut F) -> R
161    where
162        F: FnMut(&mut Self) -> R,
163    {
164        BOOL_EVALUATOR.with_borrow_mut(|s| func(s.as_mut().expect("Parameters not set")))
165    }
166}
167
168pub(crate) type RuntimeServerKey = ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>;
169impl Global for RuntimeServerKey {
170    fn global() -> &'static Self {
171        BOOL_SERVER_KEY.get().expect("Server key not set!")
172    }
173}
174
175mod impl_enc_dec {
176    use crate::{
177        pbs::{sample_extract, PbsInfo},
178        rgsw::public_key_encrypt_rlwe,
179        Decryptor, Encryptor, Matrix, MatrixEntity, MultiPartyDecryptor, RowEntity,
180    };
181    use num_traits::Zero;
182
183    use super::*;
184
185    type Mat = Vec<Vec<u64>>;
186
187    impl<E> Encryptor<bool, Vec<u64>> for super::super::keys::ClientKey<[u8; 32], E> {
188        fn encrypt(&self, m: &bool) -> Vec<u64> {
189            BoolEvaluator::with_local(|e| e.sk_encrypt(*m, self))
190        }
191    }
192
193    impl<E> Decryptor<bool, Vec<u64>> for super::super::keys::ClientKey<[u8; 32], E> {
194        fn decrypt(&self, c: &Vec<u64>) -> bool {
195            BoolEvaluator::with_local(|e| e.sk_decrypt(c, self))
196        }
197    }
198
199    impl<E> MultiPartyDecryptor<bool, <Mat as Matrix>::R>
200        for super::super::keys::ClientKey<[u8; 32], E>
201    {
202        type DecryptionShare = <Mat as Matrix>::MatElement;
203
204        fn gen_decryption_share(&self, c: &<Mat as Matrix>::R) -> Self::DecryptionShare {
205            BoolEvaluator::with_local(|e| e.multi_party_decryption_share(c, self))
206        }
207
208        fn aggregate_decryption_shares(
209            &self,
210            c: &<Mat as Matrix>::R,
211            shares: &[Self::DecryptionShare],
212        ) -> bool {
213            BoolEvaluator::with_local(|e| e.multi_party_decrypt(shares, c))
214        }
215    }
216
217    impl<Rng, ModOp> Encryptor<[bool], Mat> for PublicKey<Mat, Rng, ModOp> {
218        fn encrypt(&self, m: &[bool]) -> Mat {
219            BoolEvaluator::with_local(|e| {
220                DefaultSecureRng::with_local_mut(|rng| {
221                    let parameters = e.parameters();
222                    let mut rlwe_out = <Mat as MatrixEntity>::zeros(2, parameters.rlwe_n().0);
223                    assert!(m.len() <= parameters.rlwe_n().0);
224
225                    let mut message =
226                        vec![<Mat as Matrix>::MatElement::zero(); parameters.rlwe_n().0];
227                    m.iter().enumerate().for_each(|(i, v)| {
228                        if *v {
229                            message[i] = parameters.rlwe_q().true_el()
230                        } else {
231                            message[i] = parameters.rlwe_q().false_el()
232                        }
233                    });
234
235                    // e.pk_encrypt_batched(self.key(), m)
236                    public_key_encrypt_rlwe::<_, _, _, _, i32, _>(
237                        &mut rlwe_out,
238                        self.key(),
239                        &message,
240                        e.pbs_info().modop_rlweq(),
241                        e.pbs_info().nttop_rlweq(),
242                        rng,
243                    );
244                    rlwe_out
245                })
246            })
247        }
248    }
249
250    impl<Rng, ModOp> Encryptor<bool, <Mat as Matrix>::R> for PublicKey<Mat, Rng, ModOp> {
251        fn encrypt(&self, m: &bool) -> <Mat as Matrix>::R {
252            let m = vec![*m];
253            let rlwe = self.encrypt(m.as_slice());
254            BoolEvaluator::with_local(|e| {
255                let mut lwe = <Mat as Matrix>::R::zeros(e.parameters().rlwe_n().0 + 1);
256                sample_extract(&mut lwe, &rlwe, e.pbs_info().modop_rlweq(), 0);
257                lwe
258            })
259        }
260    }
261}