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