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