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