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