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