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