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