1use crate::{
  2    backend::{ModularOpsU64, ModulusPowerOf2},
  3    ntt::NttBackendU64,
  4    random::{DefaultSecureRng, NewWithSeed},
  5    utils::{Global, WithLocal},
  6};
  7
  8use super::*;
  9
 10thread_local! {
 11    static BOOL_EVALUATOR: RefCell<Option<BoolEvaluator<Vec<Vec<u64>>, NttBackendU64, ModularOpsU64<CiphertextModulus<u64>>,  ModulusPowerOf2<CiphertextModulus<u64>>, ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>>>> = RefCell::new(None);
 12
 13}
 14static BOOL_SERVER_KEY: OnceLock<ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>> = OnceLock::new();
 15
 16static MULTI_PARTY_CRS: OnceLock<MultiPartyCrs<[u8; 32]>> = OnceLock::new();
 17
 18pub enum ParameterSelector {
 19    MultiPartyLessThanOrEqualTo16,
 20}
 21
 22pub fn set_parameter_set(select: ParameterSelector) {
 23    match select {
 24        ParameterSelector::MultiPartyLessThanOrEqualTo16 => {
 25            BOOL_EVALUATOR.with_borrow_mut(|v| *v = Some(BoolEvaluator::new(SMALL_MP_BOOL_PARAMS)));
 26        }
 27    }
 28}
 29
 30pub fn set_mp_seed(seed: [u8; 32]) {
 31    assert!(
 32        MULTI_PARTY_CRS.set(MultiPartyCrs { seed: seed }).is_ok(),
 33        "Attempted to set MP SEED twice."
 34    )
 35}
 36
 37fn set_server_key(key: ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>) {
 38    assert!(
 39        BOOL_SERVER_KEY.set(key).is_ok(),
 40        "Attempted to set server key twice."
 41    );
 42}
 43
 44pub(crate) fn gen_keys() -> (
 45    ClientKey,
 46    SeededSinglePartyServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]>,
 47) {
 48    BoolEvaluator::with_local_mut(|e| {
 49        let ck = e.client_key();
 50        let sk = e.single_party_server_key(&ck);
 51
 52        (ck, sk)
 53    })
 54}
 55
 56pub fn gen_client_key() -> ClientKey {
 57    BoolEvaluator::with_local(|e| e.client_key())
 58}
 59
 60pub fn gen_mp_keys_phase1(
 61    ck: &ClientKey,
 62) -> CommonReferenceSeededCollectivePublicKeyShare<Vec<u64>, [u8; 32], BoolParameters<u64>> {
 63    let seed = MultiPartyCrs::global().public_key_share_seed::<DefaultSecureRng>();
 64    BoolEvaluator::with_local(|e| {
 65        let pk_share = e.multi_party_public_key_share(seed, &ck);
 66        pk_share
 67    })
 68}
 69
 70pub fn gen_mp_keys_phase2<R, ModOp>(
 71    ck: &ClientKey,
 72    pk: &PublicKey<Vec<Vec<u64>>, R, ModOp>,
 73) -> CommonReferenceSeededMultiPartyServerKeyShare<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> {
 74    let seed = MultiPartyCrs::global().server_key_share_seed::<DefaultSecureRng>();
 75    BoolEvaluator::with_local_mut(|e| {
 76        let server_key_share = e.multi_party_server_key_share(seed, pk.key(), ck);
 77        server_key_share
 78    })
 79}
 80
 81pub fn aggregate_public_key_shares(
 82    shares: &[CommonReferenceSeededCollectivePublicKeyShare<
 83        Vec<u64>,
 84        [u8; 32],
 85        BoolParameters<u64>,
 86    >],
 87) -> PublicKey<Vec<Vec<u64>>, DefaultSecureRng, ModularOpsU64<CiphertextModulus<u64>>> {
 88    PublicKey::from(shares)
 89}
 90
 91pub fn aggregate_server_key_shares(
 92    shares: &[CommonReferenceSeededMultiPartyServerKeyShare<
 93        Vec<Vec<u64>>,
 94        BoolParameters<u64>,
 95        [u8; 32],
 96    >],
 97) -> SeededMultiPartyServerKey<Vec<Vec<u64>>, [u8; 32], BoolParameters<u64>> {
 98    BoolEvaluator::with_local(|e| e.aggregate_multi_party_server_key_shares(shares))
 99}
100
101// SERVER KEY EVAL (/SHOUP) DOMAIN //
102impl SeededSinglePartyServerKey<Vec<Vec<u64>>, BoolParameters<u64>, [u8; 32]> {
103    pub fn set_server_key(&self) {
104        let eval = ServerKeyEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(self);
105        set_server_key(ShoupServerKeyEvaluationDomain::from(eval));
106    }
107}
108
109impl
110    SeededMultiPartyServerKey<
111        Vec<Vec<u64>>,
112        <DefaultSecureRng as NewWithSeed>::Seed,
113        BoolParameters<u64>,
114    >
115{
116    pub fn set_server_key(&self) {
117        set_server_key(ShoupServerKeyEvaluationDomain::from(
118            ServerKeyEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(self),
119        ))
120    }
121}
122
123//  MULTIPARTY CRS  //
124impl Global for MultiPartyCrs<[u8; 32]> {
125    fn global() -> &'static Self {
126        MULTI_PARTY_CRS
127            .get()
128            .expect("Multi Party Common Reference String not set")
129    }
130}
131
132// BOOL EVALUATOR //
133impl WithLocal
134    for BoolEvaluator<
135        Vec<Vec<u64>>,
136        NttBackendU64,
137        ModularOpsU64<CiphertextModulus<u64>>,
138        ModulusPowerOf2<CiphertextModulus<u64>>,
139        ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>,
140    >
141{
142    fn with_local<F, R>(func: F) -> R
143    where
144        F: Fn(&Self) -> R,
145    {
146        BOOL_EVALUATOR.with_borrow(|s| func(s.as_ref().expect("Parameters not set")))
147    }
148
149    fn with_local_mut<F, R>(func: F) -> R
150    where
151        F: Fn(&mut Self) -> R,
152    {
153        BOOL_EVALUATOR.with_borrow_mut(|s| func(s.as_mut().expect("Parameters not set")))
154    }
155
156    fn with_local_mut_mut<F, R>(func: &mut F) -> R
157    where
158        F: FnMut(&mut Self) -> R,
159    {
160        BOOL_EVALUATOR.with_borrow_mut(|s| func(s.as_mut().expect("Parameters not set")))
161    }
162}
163
164pub(crate) type RuntimeServerKey = ShoupServerKeyEvaluationDomain<Vec<Vec<u64>>>;
165impl Global for RuntimeServerKey {
166    fn global() -> &'static Self {
167        BOOL_SERVER_KEY.get().expect("Server key not set!")
168    }
169}