arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/bool/evaluator.rs
src/bool/evaluator.rsRaw
   1use std::{
   2    cell::{OnceCell, RefCell},
   3    collections::HashMap,
   4    fmt::{Debug, Display},
   5    marker::PhantomData,
   6    ops::Shr,
   7};
   8
   9use itertools::{izip, partition, Itertools};
  10use num_traits::{FromPrimitive, Num, One, PrimInt, ToPrimitive, WrappingSub, Zero};
  11
  12use crate::{
  13    backend::{ArithmeticOps, GetModulus, ModInit, ModularOpsU64, Modulus, VectorOps},
  14    bool::parameters::{MP_BOOL_PARAMS, SP_BOOL_PARAMS},
  15    decomposer::{Decomposer, DefaultDecomposer, NumInfo, RlweDecomposer},
  16    lwe::{decrypt_lwe, encrypt_lwe, lwe_key_switch, lwe_ksk_keygen, measure_noise_lwe, LweSecret},
  17    multi_party::public_key_share,
  18    ntt::{self, Ntt, NttBackendU64, NttInit},
  19    random::{
  20        DefaultSecureRng, NewWithSeed, RandomFillGaussianInModulus, RandomFillUniformInModulus,
  21        RandomGaussianElementInModulus,
  22    },
  23    rgsw::{
  24        decrypt_rlwe, galois_auto, galois_key_gen, generate_auto_map, public_key_encrypt_rgsw,
  25        rgsw_by_rgsw_inplace, rlwe_by_rgsw, secret_key_encrypt_rgsw, IsTrivial, RgswCiphertext,
  26        RlweCiphertext, RlweSecret,
  27    },
  28    utils::{
  29        fill_random_ternary_secret_with_hamming_weight, generate_prime, mod_exponent,
  30        TryConvertFrom1, WithLocal,
  31    },
  32    Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret,
  33};
  34
  35use super::parameters::{BoolParameters, CiphertextModulus};
  36
  37thread_local! {
  38    static BOOL_EVALUATOR: RefCell<BoolEvaluator<Vec<Vec<u64>>, NttBackendU64, ModularOpsU64<CiphertextModulus<u64>>,  ModularOpsU64<CiphertextModulus<u64>>>> = RefCell::new(BoolEvaluator::new(MP_BOOL_PARAMS));
  39}
  40
  41pub fn set_parameter_set(parameter: &BoolParameters<u64>) {
  42    BoolEvaluator::with_local_mut(|e| *e = BoolEvaluator::new(parameter.clone()))
  43}
  44
  45impl WithLocal
  46    for BoolEvaluator<
  47        Vec<Vec<u64>>,
  48        NttBackendU64,
  49        ModularOpsU64<CiphertextModulus<u64>>,
  50        ModularOpsU64<CiphertextModulus<u64>>,
  51    >
  52{
  53    fn with_local<F, R>(func: F) -> R
  54    where
  55        F: Fn(&Self) -> R,
  56    {
  57        BOOL_EVALUATOR.with_borrow(|s| func(s))
  58    }
  59
  60    fn with_local_mut<F, R>(func: F) -> R
  61    where
  62        F: Fn(&mut Self) -> R,
  63    {
  64        BOOL_EVALUATOR.with_borrow_mut(|s| func(s))
  65    }
  66}
  67
  68struct ScratchMemory<M>
  69where
  70    M: Matrix,
  71{
  72    lwe_vector: M::R,
  73    decomposition_matrix: M,
  74}
  75
  76impl<M: MatrixEntity> ScratchMemory<M>
  77where
  78    M::R: RowEntity,
  79{
  80    fn new(parameters: &BoolParameters<M::MatElement>) -> Self {
  81        // Vector to store LWE ciphertext with LWE dimesnion n
  82        let lwe_vector = M::R::zeros(parameters.lwe_n().0 + 1);
  83
  84        // Matrix to store decomposed polynomials
  85        // Max decompistion count + space for temporary RLWE
  86        let d = std::cmp::max(
  87            parameters.auto_decomposition_count().0,
  88            std::cmp::max(
  89                parameters.rlwe_rgsw_decomposition_count().0 .0,
  90                parameters.rlwe_rgsw_decomposition_count().1 .0,
  91            ),
  92        ) + 2;
  93        let decomposition_matrix = M::zeros(d, parameters.rlwe_n().0);
  94
  95        Self {
  96            lwe_vector,
  97            decomposition_matrix,
  98        }
  99    }
 100}
 101
 102// thread_local! {
 103//     pub(crate) static CLIENT_KEY: RefCell<ClientKey> =
 104// RefCell::new(ClientKey::random()); }
 105
 106trait BoolEncoding {
 107    type Element;
 108    fn true_el(&self) -> Self::Element;
 109    fn false_el(&self) -> Self::Element;
 110    fn qby4(&self) -> Self::Element;
 111    fn decode(&self, m: Self::Element) -> bool;
 112}
 113
 114impl<T> BoolEncoding for CiphertextModulus<T>
 115where
 116    CiphertextModulus<T>: Modulus<Element = T>,
 117    T: PrimInt,
 118{
 119    type Element = T;
 120
 121    fn qby4(&self) -> Self::Element {
 122        if self.is_native() {
 123            T::one() << (CiphertextModulus::<T>::_bits() - 2)
 124        } else {
 125            self.q().unwrap() >> 2
 126        }
 127    }
 128    /// Q/8
 129    fn true_el(&self) -> Self::Element {
 130        if self.is_native() {
 131            T::one() << (CiphertextModulus::<T>::_bits() - 3)
 132        } else {
 133            self.q().unwrap() >> 3
 134        }
 135    }
 136    /// -Q/8
 137    fn false_el(&self) -> Self::Element {
 138        self.largest_unsigned_value() - self.true_el() + T::one()
 139    }
 140    fn decode(&self, m: Self::Element) -> bool {
 141        let qby8 = self.true_el();
 142        let m = (((m + qby8).to_f64().unwrap() * 4.0f64) / self.q_as_f64().unwrap()).round()
 143            as usize
 144            % 4usize;
 145
 146        if m == 0 {
 147            return false;
 148        } else if m == 1 {
 149            return true;
 150        } else {
 151            panic!("Incorrect bool decryption. Got m={m} but expected m to be 0 or 1")
 152        }
 153    }
 154}
 155
 156trait PbsKey {
 157    type M: Matrix;
 158
 159    /// RGSW ciphertext of LWE secret elements
 160    fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::M;
 161    /// Key for automorphism
 162    fn galois_key_for_auto(&self, k: isize) -> &Self::M;
 163    /// LWE ksk to key switch from RLWE secret to LWE secret
 164    fn lwe_ksk(&self) -> &Self::M;
 165}
 166
 167trait PbsInfo {
 168    type Element;
 169    type Modulus: Modulus<Element = Self::Element>;
 170    type NttOp: Ntt<Element = Self::Element>;
 171    type D: Decomposer<Element = Self::Element>;
 172
 173    // Although both types have same bounds, they can be different types. For ex,
 174    // type RlweModOp may only support native modulus, where LweModOp may only
 175    // support prime modulus, etc.
 176    type RlweModOp: VectorOps<Element = Self::Element> + ArithmeticOps<Element = Self::Element>;
 177    type LweModOp: VectorOps<Element = Self::Element> + ArithmeticOps<Element = Self::Element>;
 178
 179    fn rlwe_q(&self) -> &Self::Modulus;
 180    fn lwe_q(&self) -> &Self::Modulus;
 181    fn br_q(&self) -> usize;
 182    fn rlwe_n(&self) -> usize;
 183    fn lwe_n(&self) -> usize;
 184    /// Embedding fator for ring X^{q}+1 inside
 185    fn embedding_factor(&self) -> usize;
 186    /// generator g
 187    fn g(&self) -> isize;
 188    /// Decomposers
 189    fn lwe_decomposer(&self) -> &Self::D;
 190    fn rlwe_rgsw_decomposer(&self) -> &(Self::D, Self::D);
 191    fn auto_decomposer(&self) -> &Self::D;
 192
 193    /// Modulus operators
 194    fn modop_lweq(&self) -> &Self::LweModOp;
 195    fn modop_rlweq(&self) -> &Self::RlweModOp;
 196
 197    /// Ntt operators
 198    fn nttop_rlweq(&self) -> &Self::NttOp;
 199
 200    /// Maps a \in Z^*_{q} to discrete log k, with generator g (i.e. g^k =
 201    /// a). Returned vector is of size q that stores dlog of a at `vec[a]`.
 202    /// For any a, if k is s.t. a = g^{k}, then k is expressed as k. If k is s.t
 203    /// a = -g^{k}, then k is expressed as k=k+q/2
 204    fn g_k_dlog_map(&self) -> &[usize];
 205    fn rlwe_auto_map(&self, k: isize) -> &(Vec<usize>, Vec<bool>);
 206}
 207
 208#[derive(Clone)]
 209struct ClientKey {
 210    sk_rlwe: RlweSecret,
 211    sk_lwe: LweSecret,
 212}
 213
 214impl ClientKey {
 215    fn random() -> Self {
 216        let sk_rlwe = RlweSecret::random(0, 0);
 217        let sk_lwe = LweSecret::random(0, 0);
 218        Self { sk_rlwe, sk_lwe }
 219    }
 220}
 221
 222// impl WithLocal for ClientKey {
 223//     fn with_local<F, R>(func: F) -> R
 224//     where
 225//         F: Fn(&Self) -> R,
 226//     {
 227//         CLIENT_KEY.with_borrow(|client_key| func(client_key))
 228//     }
 229
 230//     fn with_local_mut<F, R>(func: F) -> R
 231//     where
 232//         F: Fn(&mut Self) -> R,
 233//     {
 234//         CLIENT_KEY.with_borrow_mut(|client_key| func(client_key))
 235//     }
 236// }
 237
 238// fn set_client_key(key: &ClientKey) {
 239//     ClientKey::with_local_mut(|k| *k = key.clone())
 240// }
 241
 242struct MultiPartyDecryptionShare<E> {
 243    share: E,
 244}
 245
 246struct CommonReferenceSeededCollectivePublicKeyShare<R, S, P> {
 247    share: R,
 248    cr_seed: S,
 249    parameters: P,
 250}
 251
 252struct PublicKey<M, R, O> {
 253    key: M,
 254    _phantom: PhantomData<(R, O)>,
 255}
 256
 257impl<
 258        M: MatrixMut + MatrixEntity,
 259        Rng: NewWithSeed + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
 260        ModOp: VectorOps<Element = M::MatElement> + ModInit<M = CiphertextModulus<M::MatElement>>,
 261    >
 262    From<
 263        &[CommonReferenceSeededCollectivePublicKeyShare<
 264            M::R,
 265            Rng::Seed,
 266            BoolParameters<M::MatElement>,
 267        >],
 268    > for PublicKey<M, Rng, ModOp>
 269where
 270    <M as Matrix>::R: RowMut,
 271    Rng::Seed: Copy + PartialEq,
 272    M::MatElement: PartialEq + Copy,
 273{
 274    fn from(
 275        value: &[CommonReferenceSeededCollectivePublicKeyShare<
 276            M::R,
 277            Rng::Seed,
 278            BoolParameters<M::MatElement>,
 279        >],
 280    ) -> Self {
 281        assert!(value.len() > 0);
 282
 283        let parameters = &value[0].parameters;
 284        let mut key = M::zeros(2, parameters.rlwe_n().0);
 285
 286        // sample A
 287        let seed = value[0].cr_seed;
 288        let mut main_rng = Rng::new_with_seed(seed);
 289        RandomFillUniformInModulus::random_fill(
 290            &mut main_rng,
 291            parameters.rlwe_q(),
 292            key.get_row_mut(0),
 293        );
 294
 295        // Sum all Bs
 296        let rlweq_modop = ModOp::new(parameters.rlwe_q().clone());
 297        value.iter().for_each(|share_i| {
 298            assert!(share_i.cr_seed == seed);
 299            assert!(&share_i.parameters == parameters);
 300
 301            rlweq_modop.elwise_add_mut(key.get_row_mut(1), share_i.share.as_ref());
 302        });
 303
 304        PublicKey {
 305            key,
 306            _phantom: PhantomData,
 307        }
 308    }
 309}
 310
 311struct CommonReferenceSeededMultiPartyServerKeyShare<M: Matrix, P, S> {
 312    rgsw_cts: Vec<M>,
 313    auto_keys: HashMap<isize, M>,
 314    lwe_ksk: M::R,
 315    /// Common reference seed
 316    cr_seed: S,
 317    parameters: P,
 318}
 319struct SeededMultiPartyServerKey<M: Matrix, S, P> {
 320    rgsw_cts: Vec<M>,
 321    auto_keys: HashMap<isize, M>,
 322    lwe_ksk: M::R,
 323    cr_seed: S,
 324    parameters: P,
 325}
 326
 327/// Seeded single party server key
 328struct SeededServerKey<M: Matrix, P, S> {
 329    /// Rgsw cts of LWE secret elements
 330    pub(crate) rgsw_cts: Vec<M>,
 331    /// Auto keys
 332    pub(crate) auto_keys: HashMap<isize, M>,
 333    /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
 334    pub(crate) lwe_ksk: M::R,
 335    /// Parameters
 336    pub(crate) parameters: P,
 337    /// Main seed
 338    pub(crate) seed: S,
 339}
 340
 341impl<M: Matrix, S> SeededServerKey<M, BoolParameters<M::MatElement>, S> {
 342    pub(crate) fn from_raw(
 343        auto_keys: HashMap<isize, M>,
 344        rgsw_cts: Vec<M>,
 345        lwe_ksk: M::R,
 346        parameters: BoolParameters<M::MatElement>,
 347        seed: S,
 348    ) -> Self {
 349        // sanity checks
 350        auto_keys.iter().for_each(|v| {
 351            assert!(
 352                v.1.dimension()
 353                    == (
 354                        parameters.auto_decomposition_count().0,
 355                        parameters.rlwe_n().0
 356                    )
 357            )
 358        });
 359
 360        let (part_a_d, part_b_d) = parameters.rlwe_rgsw_decomposition_count();
 361        rgsw_cts.iter().for_each(|v| {
 362            assert!(v.dimension() == (part_a_d.0 * 2 + part_b_d.0, parameters.rlwe_n().0))
 363        });
 364        assert!(
 365            lwe_ksk.as_ref().len()
 366                == (parameters.lwe_decomposition_count().0 * parameters.rlwe_n().0)
 367        );
 368
 369        SeededServerKey {
 370            rgsw_cts,
 371            auto_keys,
 372            lwe_ksk,
 373            parameters,
 374            seed,
 375        }
 376    }
 377}
 378
 379/// Server key in evaluation domain
 380struct ServerKeyEvaluationDomain<M, R, N> {
 381    /// Rgsw cts of LWE secret elements
 382    rgsw_cts: Vec<M>,
 383    /// Galois keys
 384    galois_keys: HashMap<isize, M>,
 385    /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
 386    lwe_ksk: M,
 387    _phanton: PhantomData<(R, N)>,
 388}
 389
 390impl<
 391        M: MatrixMut + MatrixEntity,
 392        R: RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>> + NewWithSeed,
 393        N: NttInit<CiphertextModulus<M::MatElement>> + Ntt<Element = M::MatElement>,
 394    > From<&SeededServerKey<M, BoolParameters<M::MatElement>, R::Seed>>
 395    for ServerKeyEvaluationDomain<M, R, N>
 396where
 397    <M as Matrix>::R: RowMut,
 398    M::MatElement: Copy,
 399    R::Seed: Clone,
 400{
 401    fn from(value: &SeededServerKey<M, BoolParameters<M::MatElement>, R::Seed>) -> Self {
 402        let mut main_prng = R::new_with_seed(value.seed.clone());
 403        let parameters = &value.parameters;
 404        let g = parameters.g() as isize;
 405        let ring_size = value.parameters.rlwe_n().0;
 406        let lwe_n = value.parameters.lwe_n().0;
 407        let rlwe_q = value.parameters.rlwe_q();
 408        let lwq_q = value.parameters.lwe_q();
 409
 410        let nttop = N::new(rlwe_q, ring_size);
 411
 412        // galois keys
 413        let mut auto_keys = HashMap::new();
 414        let auto_decomp_count = parameters.auto_decomposition_count().0;
 415        for i in [g, -g] {
 416            let seeded_auto_key = value.auto_keys.get(&i).unwrap();
 417            assert!(seeded_auto_key.dimension() == (auto_decomp_count, ring_size));
 418
 419            let mut data = M::zeros(auto_decomp_count * 2, ring_size);
 420
 421            // sample RLWE'_A(-s(X^k))
 422            data.iter_rows_mut().take(auto_decomp_count).for_each(|ri| {
 423                RandomFillUniformInModulus::random_fill(&mut main_prng, &rlwe_q, ri.as_mut())
 424            });
 425
 426            // copy over RLWE'B_(-s(X^k))
 427            izip!(
 428                data.iter_rows_mut().skip(auto_decomp_count),
 429                seeded_auto_key.iter_rows()
 430            )
 431            .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
 432
 433            // Send to Evaluation domain
 434            data.iter_rows_mut()
 435                .for_each(|ri| nttop.forward(ri.as_mut()));
 436
 437            auto_keys.insert(i, data);
 438        }
 439
 440        // RGSW ciphertexts
 441        let (rlrg_a_decomp, rlrg_b_decomp) = parameters.rlwe_rgsw_decomposition_count();
 442        let rgsw_cts = value
 443            .rgsw_cts
 444            .iter()
 445            .map(|seeded_rgsw_si| {
 446                assert!(
 447                    seeded_rgsw_si.dimension()
 448                        == (rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0, ring_size)
 449                );
 450
 451                let mut data = M::zeros(rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0 * 2, ring_size);
 452
 453                // copy over RLWE'(-sm)
 454                izip!(
 455                    data.iter_rows_mut().take(rlrg_a_decomp.0 * 2),
 456                    seeded_rgsw_si.iter_rows().take(rlrg_a_decomp.0 * 2)
 457                )
 458                .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
 459
 460                // sample RLWE'_A(m)
 461                data.iter_rows_mut()
 462                    .skip(rlrg_a_decomp.0 * 2)
 463                    .take(rlrg_b_decomp.0)
 464                    .for_each(|ri| {
 465                        RandomFillUniformInModulus::random_fill(
 466                            &mut main_prng,
 467                            &rlwe_q,
 468                            ri.as_mut(),
 469                        )
 470                    });
 471
 472                // copy over RLWE'_B(m)
 473                izip!(
 474                    data.iter_rows_mut()
 475                        .skip(rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0),
 476                    seeded_rgsw_si.iter_rows().skip(rlrg_a_decomp.0 * 2)
 477                )
 478                .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
 479
 480                // send polynomials to evaluation domain
 481                data.iter_rows_mut()
 482                    .for_each(|ri| nttop.forward(ri.as_mut()));
 483
 484                data
 485            })
 486            .collect_vec();
 487
 488        // LWE ksk
 489        let lwe_ksk = {
 490            let d = parameters.lwe_decomposition_count().0;
 491            assert!(value.lwe_ksk.as_ref().len() == d * ring_size);
 492
 493            let mut data = M::zeros(d * ring_size, lwe_n + 1);
 494            izip!(data.iter_rows_mut(), value.lwe_ksk.as_ref().iter()).for_each(|(lwe_i, bi)| {
 495                RandomFillUniformInModulus::random_fill(
 496                    &mut main_prng,
 497                    &lwq_q,
 498                    &mut lwe_i.as_mut()[1..],
 499                );
 500                lwe_i.as_mut()[0] = *bi;
 501            });
 502
 503            data
 504        };
 505
 506        ServerKeyEvaluationDomain {
 507            rgsw_cts,
 508            galois_keys: auto_keys,
 509            lwe_ksk,
 510            _phanton: PhantomData,
 511        }
 512    }
 513}
 514
 515impl<
 516        M: MatrixMut + MatrixEntity,
 517        Rng: NewWithSeed,
 518        N: NttInit<CiphertextModulus<M::MatElement>> + Ntt<Element = M::MatElement>,
 519    > From<&SeededMultiPartyServerKey<M, Rng::Seed, BoolParameters<M::MatElement>>>
 520    for ServerKeyEvaluationDomain<M, Rng, N>
 521where
 522    <M as Matrix>::R: RowMut,
 523    Rng::Seed: Copy,
 524    Rng: RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
 525    M::MatElement: Copy,
 526{
 527    fn from(
 528        value: &SeededMultiPartyServerKey<M, Rng::Seed, BoolParameters<M::MatElement>>,
 529    ) -> Self {
 530        let g = value.parameters.g() as isize;
 531        let rlwe_n = value.parameters.rlwe_n().0;
 532        let lwe_n = value.parameters.lwe_n().0;
 533        let rlwe_q = value.parameters.rlwe_q();
 534        let lwe_q = value.parameters.lwe_q();
 535
 536        let mut main_prng = Rng::new_with_seed(value.cr_seed);
 537
 538        let rlwe_nttop = N::new(rlwe_q, rlwe_n);
 539
 540        // auto keys
 541        let mut auto_keys = HashMap::new();
 542        let auto_d_count = value.parameters.auto_decomposition_count().0;
 543        for i in [g, -g] {
 544            let mut key = M::zeros(auto_d_count * 2, rlwe_n);
 545
 546            // sample a
 547            key.iter_rows_mut().take(auto_d_count).for_each(|ri| {
 548                RandomFillUniformInModulus::random_fill(&mut main_prng, &rlwe_q, ri.as_mut())
 549            });
 550
 551            let key_part_b = value.auto_keys.get(&i).unwrap();
 552            assert!(key_part_b.dimension() == (auto_d_count, rlwe_n));
 553            izip!(
 554                key.iter_rows_mut().skip(auto_d_count),
 555                key_part_b.iter_rows()
 556            )
 557            .for_each(|(to_ri, from_ri)| {
 558                to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 559            });
 560
 561            // send to evaluation domain
 562            key.iter_rows_mut()
 563                .for_each(|ri| rlwe_nttop.forward(ri.as_mut()));
 564
 565            auto_keys.insert(i, key);
 566        }
 567
 568        // rgsw cts
 569        let (rgswrgsw_d_a, rgswrgsw_d_b) = value.parameters.rgsw_rgsw_decomposition_count();
 570        let (rlrg_d_a, rlrg_d_b) = value.parameters.rlwe_rgsw_decomposition_count();
 571        // d_a and d_b may differ for RGSWxRGSW multiplication and RLWExRGSW
 572        // multiplication. After this point RGSW ciphertexts will only be used for
 573        // RLWExRGSW multiplication (in blind rotation). Thus we drop any additional
 574        // RLWE ciphertexts in incoming collective RGSW ciphertexts, which are
 575        // result of RGSWxRGSW multiplications.
 576        let rgsw_ct_rows_in = rgswrgsw_d_a.0 * 2 + rgswrgsw_d_b.0 * 2;
 577        let rgsw_ct_rows_out = rlrg_d_a.0 * 2 + rlrg_d_b.0 * 2;
 578        let rgsw_cts = value
 579            .rgsw_cts
 580            .iter()
 581            .map(|ct_i_in| {
 582                assert!(ct_i_in.dimension() == (rgsw_ct_rows_in, rlwe_n));
 583                let mut eval_ct_i_out = M::zeros(rgsw_ct_rows_out, rlwe_n);
 584
 585                // RLWE'(-sm) part A
 586                izip!(
 587                    eval_ct_i_out.iter_rows_mut().take(rlrg_d_a.0),
 588                    ct_i_in.iter_rows().take(rlrg_d_a.0)
 589                )
 590                .for_each(|(to_ri, from_ri)| {
 591                    to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 592                    rlwe_nttop.forward(to_ri.as_mut());
 593                });
 594
 595                // RLWE'(-sm) part B
 596                izip!(
 597                    eval_ct_i_out
 598                        .iter_rows_mut()
 599                        .skip(rlrg_d_a.0)
 600                        .take(rlrg_d_a.0),
 601                    ct_i_in.iter_rows().skip(rgswrgsw_d_a.0).take(rlrg_d_a.0)
 602                )
 603                .for_each(|(to_ri, from_ri)| {
 604                    to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 605                    rlwe_nttop.forward(to_ri.as_mut());
 606                });
 607
 608                // RLWE'(m) Part A
 609                izip!(
 610                    eval_ct_i_out
 611                        .iter_rows_mut()
 612                        .skip(rlrg_d_a.0 * 2)
 613                        .take(rlrg_d_b.0),
 614                    ct_i_in
 615                        .iter_rows()
 616                        .skip(rgswrgsw_d_a.0 * 2)
 617                        .take(rlrg_d_b.0)
 618                )
 619                .for_each(|(to_ri, from_ri)| {
 620                    to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 621                    rlwe_nttop.forward(to_ri.as_mut());
 622                });
 623
 624                // RLWE'(m) Part B
 625                izip!(
 626                    eval_ct_i_out
 627                        .iter_rows_mut()
 628                        .skip(rlrg_d_a.0 * 2 + rlrg_d_b.0)
 629                        .take(rlrg_d_b.0),
 630                    ct_i_in
 631                        .iter_rows()
 632                        .skip(rgswrgsw_d_a.0 * 2 + rgswrgsw_d_b.0)
 633                        .take(rlrg_d_b.0)
 634                )
 635                .for_each(|(to_ri, from_ri)| {
 636                    to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 637                    rlwe_nttop.forward(to_ri.as_mut());
 638                });
 639
 640                eval_ct_i_out
 641            })
 642            .collect_vec();
 643
 644        // lwe ksk
 645        let d_lwe = value.parameters.lwe_decomposition_count().0;
 646        let mut lwe_ksk = M::zeros(rlwe_n * d_lwe, lwe_n + 1);
 647        izip!(lwe_ksk.iter_rows_mut(), value.lwe_ksk.as_ref().iter()).for_each(|(lwe_i, bi)| {
 648            RandomFillUniformInModulus::random_fill(
 649                &mut main_prng,
 650                &lwe_q,
 651                &mut lwe_i.as_mut()[1..],
 652            );
 653            lwe_i.as_mut()[0] = *bi;
 654        });
 655
 656        ServerKeyEvaluationDomain {
 657            rgsw_cts,
 658            galois_keys: auto_keys,
 659            lwe_ksk,
 660            _phanton: PhantomData,
 661        }
 662    }
 663}
 664
 665//FIXME(Jay): Figure out a way for BoolEvaluator to have access to ServerKey
 666// via a pointer and implement PbsKey for BoolEvaluator instead of ServerKey
 667// directly
 668impl<M: Matrix, R, N> PbsKey for ServerKeyEvaluationDomain<M, R, N> {
 669    type M = M;
 670    fn galois_key_for_auto(&self, k: isize) -> &Self::M {
 671        self.galois_keys.get(&k).unwrap()
 672    }
 673    fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::M {
 674        &self.rgsw_cts[si]
 675    }
 676
 677    fn lwe_ksk(&self) -> &Self::M {
 678        &self.lwe_ksk
 679    }
 680}
 681
 682struct BoolPbsInfo<M: Matrix, Ntt, RlweModOp, LweModOp> {
 683    auto_decomposer: DefaultDecomposer<M::MatElement>,
 684    rlwe_rgsw_decomposer: (
 685        DefaultDecomposer<M::MatElement>,
 686        DefaultDecomposer<M::MatElement>,
 687    ),
 688    lwe_decomposer: DefaultDecomposer<M::MatElement>,
 689    g_k_dlog_map: Vec<usize>,
 690    rlwe_nttop: Ntt,
 691    rlwe_modop: RlweModOp,
 692    lwe_modop: LweModOp,
 693    embedding_factor: usize,
 694    nand_test_vec: M::R,
 695    rlwe_qby4: M::MatElement,
 696    rlwe_auto_maps: Vec<(Vec<usize>, Vec<bool>)>,
 697    parameters: BoolParameters<M::MatElement>,
 698}
 699
 700impl<M: Matrix, NttOp, RlweModOp, LweModOp> PbsInfo for BoolPbsInfo<M, NttOp, RlweModOp, LweModOp>
 701where
 702    M::MatElement: PrimInt + WrappingSub + NumInfo + Debug + FromPrimitive,
 703    RlweModOp: ArithmeticOps<Element = M::MatElement> + VectorOps<Element = M::MatElement>,
 704    LweModOp: ArithmeticOps<Element = M::MatElement> + VectorOps<Element = M::MatElement>,
 705    NttOp: Ntt<Element = M::MatElement>,
 706{
 707    type Modulus = CiphertextModulus<M::MatElement>;
 708    type Element = M::MatElement;
 709    type D = DefaultDecomposer<M::MatElement>;
 710    type RlweModOp = RlweModOp;
 711    type LweModOp = LweModOp;
 712    type NttOp = NttOp;
 713    fn rlwe_auto_map(&self, k: isize) -> &(Vec<usize>, Vec<bool>) {
 714        let g = self.parameters.g() as isize;
 715        if k == g {
 716            &self.rlwe_auto_maps[0]
 717        } else if k == -g {
 718            &self.rlwe_auto_maps[1]
 719        } else {
 720            panic!("RLWE auto map only supports k in [-g, g], but got k={k}");
 721        }
 722    }
 723    fn br_q(&self) -> usize {
 724        *self.parameters.br_q()
 725    }
 726    fn lwe_decomposer(&self) -> &Self::D {
 727        &self.lwe_decomposer
 728    }
 729    fn rlwe_rgsw_decomposer(&self) -> &(Self::D, Self::D) {
 730        &self.rlwe_rgsw_decomposer
 731    }
 732    fn auto_decomposer(&self) -> &Self::D {
 733        &self.auto_decomposer
 734    }
 735    fn embedding_factor(&self) -> usize {
 736        self.embedding_factor
 737    }
 738    fn g(&self) -> isize {
 739        self.parameters.g() as isize
 740    }
 741    fn g_k_dlog_map(&self) -> &[usize] {
 742        &self.g_k_dlog_map
 743    }
 744    fn lwe_n(&self) -> usize {
 745        self.parameters.lwe_n().0
 746    }
 747    fn lwe_q(&self) -> &Self::Modulus {
 748        self.parameters.lwe_q()
 749    }
 750    fn rlwe_n(&self) -> usize {
 751        self.parameters.rlwe_n().0
 752    }
 753    fn rlwe_q(&self) -> &Self::Modulus {
 754        self.parameters.rlwe_q()
 755    }
 756    fn modop_lweq(&self) -> &Self::LweModOp {
 757        &self.lwe_modop
 758    }
 759    fn modop_rlweq(&self) -> &Self::RlweModOp {
 760        &self.rlwe_modop
 761    }
 762    fn nttop_rlweq(&self) -> &Self::NttOp {
 763        &self.rlwe_nttop
 764    }
 765}
 766
 767struct BoolEvaluator<M, Ntt, RlweModOp, LweModOp>
 768where
 769    M: Matrix,
 770{
 771    pbs_info: BoolPbsInfo<M, Ntt, RlweModOp, LweModOp>,
 772    scratch_memory: ScratchMemory<M>,
 773    _phantom: PhantomData<M>,
 774}
 775
 776impl<M: Matrix, NttOp, RlweModOp, LweModOp> BoolEvaluator<M, NttOp, RlweModOp, LweModOp> {}
 777
 778impl<M: Matrix, NttOp, RlweModOp, LweModOp> BoolEvaluator<M, NttOp, RlweModOp, LweModOp>
 779where
 780    M: MatrixEntity + MatrixMut,
 781    M::MatElement: PrimInt + Debug + Display + NumInfo + FromPrimitive + WrappingSub,
 782    NttOp: Ntt<Element = M::MatElement>,
 783    RlweModOp: ArithmeticOps<Element = M::MatElement>
 784        + VectorOps<Element = M::MatElement>
 785        + GetModulus<Element = M::MatElement, M = CiphertextModulus<M::MatElement>>,
 786    LweModOp: ArithmeticOps<Element = M::MatElement>
 787        + VectorOps<Element = M::MatElement>
 788        + GetModulus<Element = M::MatElement, M = CiphertextModulus<M::MatElement>>,
 789    M::R: TryConvertFrom1<[i32], CiphertextModulus<M::MatElement>> + RowEntity + Debug,
 790    <M as Matrix>::R: RowMut,
 791    DefaultSecureRng: RandomFillGaussianInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>
 792        + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>
 793        + RandomGaussianElementInModulus<M::MatElement, CiphertextModulus<M::MatElement>>
 794        + NewWithSeed,
 795{
 796    fn new(parameters: BoolParameters<M::MatElement>) -> Self
 797    where
 798        RlweModOp: ModInit<M = CiphertextModulus<M::MatElement>>,
 799        LweModOp: ModInit<M = CiphertextModulus<M::MatElement>>,
 800        NttOp: NttInit<CiphertextModulus<M::MatElement>>,
 801    {
 802        //TODO(Jay): Run sanity checks for modulus values in parameters
 803
 804        // generatr dlog map s.t. g^{k} % q = a, for all a \in Z*_{q}
 805        let g = parameters.g();
 806        let q = *parameters.br_q();
 807        let mut g_k_dlog_map = vec![0usize; q];
 808        for i in 0..q / 2 {
 809            let v = mod_exponent(g as u64, i as u64, q as u64) as usize;
 810            // g^i
 811            g_k_dlog_map[v] = i;
 812            // -(g^i)
 813            g_k_dlog_map[q - v] = i + (q / 2);
 814        }
 815
 816        let embedding_factor = (2 * parameters.rlwe_n().0) / q;
 817
 818        let rlwe_nttop = NttOp::new(parameters.rlwe_q(), parameters.rlwe_n().0);
 819        let rlwe_modop = RlweModOp::new(*parameters.rlwe_q());
 820        let lwe_modop = LweModOp::new(*parameters.lwe_q());
 821
 822        // set test vectors
 823        let q = *parameters.br_q();
 824        let qby2 = q >> 1;
 825        let qby8 = q >> 3;
 826        let mut nand_test_vec = M::R::zeros(qby2);
 827        // Q/8 (Q: rlwe_q)
 828        let true_m_el = parameters.rlwe_q().true_el();
 829        // -Q/8
 830        let false_m_el = parameters.rlwe_q().false_el();
 831        for i in 0..qby2 {
 832            if i < (3 * qby8) {
 833                nand_test_vec.as_mut()[i] = true_m_el;
 834            } else {
 835                nand_test_vec.as_mut()[i] = false_m_el;
 836            }
 837        }
 838
 839        // v(X) -> v(X^{-g})
 840        let (auto_map_index, auto_map_sign) = generate_auto_map(qby2, -(g as isize));
 841        let mut nand_test_vec_autog = M::R::zeros(qby2);
 842        izip!(
 843            nand_test_vec.as_ref().iter(),
 844            auto_map_index.iter(),
 845            auto_map_sign.iter()
 846        )
 847        .for_each(|(v, to_index, to_sign)| {
 848            if !to_sign {
 849                // negate
 850                nand_test_vec_autog.as_mut()[*to_index] = rlwe_modop.neg(v);
 851            } else {
 852                nand_test_vec_autog.as_mut()[*to_index] = *v;
 853            }
 854        });
 855
 856        // auto map indices and sign
 857        let mut rlwe_auto_maps = vec![];
 858        let ring_size = parameters.rlwe_n().0;
 859        let g = parameters.g() as isize;
 860        for i in [g, -g] {
 861            rlwe_auto_maps.push(generate_auto_map(ring_size, i))
 862        }
 863
 864        let rlwe_qby4 = parameters.rlwe_q().qby4();
 865
 866        let scratch_memory = ScratchMemory::new(&parameters);
 867
 868        let pbs_info = BoolPbsInfo {
 869            auto_decomposer: parameters.auto_decomposer(),
 870            lwe_decomposer: parameters.lwe_decomposer(),
 871            rlwe_rgsw_decomposer: parameters.rlwe_rgsw_decomposer(),
 872            g_k_dlog_map,
 873            embedding_factor,
 874            lwe_modop,
 875            rlwe_modop,
 876            rlwe_nttop,
 877            nand_test_vec: nand_test_vec_autog,
 878            rlwe_qby4,
 879            rlwe_auto_maps,
 880            parameters: parameters,
 881        };
 882
 883        BoolEvaluator {
 884            pbs_info,
 885            scratch_memory,
 886            _phantom: PhantomData,
 887        }
 888    }
 889
 890    fn client_key(&self) -> ClientKey {
 891        let sk_lwe = LweSecret::random(
 892            self.pbs_info.parameters.lwe_n().0 >> 1,
 893            self.pbs_info.parameters.lwe_n().0,
 894        );
 895        let sk_rlwe = RlweSecret::random(
 896            self.pbs_info.parameters.rlwe_n().0 >> 1,
 897            self.pbs_info.parameters.rlwe_n().0,
 898        );
 899        ClientKey { sk_rlwe, sk_lwe }
 900    }
 901
 902    fn server_key(
 903        &self,
 904        client_key: &ClientKey,
 905    ) -> SeededServerKey<M, BoolParameters<M::MatElement>, [u8; 32]> {
 906        DefaultSecureRng::with_local_mut(|rng| {
 907            let mut main_seed = [0u8; 32];
 908            rng.fill_bytes(&mut main_seed);
 909
 910            let mut main_prng = DefaultSecureRng::new_seeded(main_seed);
 911
 912            let rlwe_n = self.pbs_info.parameters.rlwe_n().0;
 913            let sk_rlwe = &client_key.sk_rlwe;
 914            let sk_lwe = &client_key.sk_lwe;
 915
 916            // generate auto keys -g, g
 917            let mut auto_keys = HashMap::new();
 918            let auto_gadget = self.pbs_info.auto_decomposer.gadget_vector();
 919            let g = self.pbs_info.parameters.g() as isize;
 920            for i in [g, -g] {
 921                let mut gk = M::zeros(self.pbs_info.auto_decomposer.decomposition_count(), rlwe_n);
 922                galois_key_gen(
 923                    &mut gk,
 924                    sk_rlwe.values(),
 925                    i,
 926                    &auto_gadget,
 927                    &self.pbs_info.rlwe_modop,
 928                    &self.pbs_info.rlwe_nttop,
 929                    &mut main_prng,
 930                    rng,
 931                );
 932                auto_keys.insert(i, gk);
 933            }
 934
 935            // generate rgsw ciphertexts RGSW(si) where si is i^th LWE secret element
 936            let ring_size = self.pbs_info.parameters.rlwe_n().0;
 937            let rlwe_q = self.pbs_info.parameters.rlwe_q();
 938            let (rlrg_d_a, rlrg_d_b) = (
 939                self.pbs_info.rlwe_rgsw_decomposer.0.decomposition_count(),
 940                self.pbs_info.rlwe_rgsw_decomposer.1.decomposition_count(),
 941            );
 942            let rlrg_gadget_a = self.pbs_info.rlwe_rgsw_decomposer.0.gadget_vector();
 943            let rlrg_gadget_b = self.pbs_info.rlwe_rgsw_decomposer.1.gadget_vector();
 944            let rgsw_cts = sk_lwe
 945                .values()
 946                .iter()
 947                .map(|si| {
 948                    // X^{si}; assume |emebedding_factor * si| < N
 949                    let mut m = M::R::zeros(ring_size);
 950                    let si = (self.pbs_info.embedding_factor as i32) * si;
 951                    // dbg!(si);
 952                    if si < 0 {
 953                        // X^{-i} = X^{2N - i} = -X^{N-i}
 954                        m.as_mut()[ring_size - (si.abs() as usize)] = rlwe_q.neg_one();
 955                    } else {
 956                        // X^{i}
 957                        m.as_mut()[si.abs() as usize] = M::MatElement::one();
 958                    }
 959
 960                    let mut rgsw_si = M::zeros(rlrg_d_a * 2 + rlrg_d_b, ring_size);
 961                    secret_key_encrypt_rgsw(
 962                        &mut rgsw_si,
 963                        m.as_ref(),
 964                        &rlrg_gadget_a,
 965                        &rlrg_gadget_b,
 966                        sk_rlwe.values(),
 967                        &self.pbs_info.rlwe_modop,
 968                        &self.pbs_info.rlwe_nttop,
 969                        &mut main_prng,
 970                        rng,
 971                    );
 972
 973                    rgsw_si
 974                })
 975                .collect_vec();
 976
 977            // LWE KSK from RLWE secret s -> LWE secret z
 978            let d_lwe_gadget = self.pbs_info.lwe_decomposer.gadget_vector();
 979            let mut lwe_ksk =
 980                M::R::zeros(self.pbs_info.lwe_decomposer.decomposition_count() * ring_size);
 981            lwe_ksk_keygen(
 982                &sk_rlwe.values(),
 983                &sk_lwe.values(),
 984                &mut lwe_ksk,
 985                &d_lwe_gadget,
 986                &self.pbs_info.lwe_modop,
 987                &mut main_prng,
 988                rng,
 989            );
 990
 991            SeededServerKey::from_raw(
 992                auto_keys,
 993                rgsw_cts,
 994                lwe_ksk,
 995                self.pbs_info.parameters.clone(),
 996                main_seed,
 997            )
 998        })
 999    }
1000
1001    fn multi_party_server_key_share(
1002        &self,
1003        cr_seed: [u8; 32],
1004        collective_pk: &M,
1005        client_key: &ClientKey,
1006    ) -> CommonReferenceSeededMultiPartyServerKeyShare<M, BoolParameters<M::MatElement>, [u8; 32]>
1007    {
1008        DefaultSecureRng::with_local_mut(|rng| {
1009            let mut main_prng = DefaultSecureRng::new_seeded(cr_seed);
1010
1011            let sk_rlwe = &client_key.sk_rlwe;
1012            let sk_lwe = &client_key.sk_lwe;
1013
1014            let g = self.pbs_info.parameters.g() as isize;
1015            let ring_size = self.pbs_info.parameters.rlwe_n().0;
1016            let rlwe_q = self.pbs_info.parameters.rlwe_q();
1017            let lwe_q = self.pbs_info.parameters.lwe_q();
1018
1019            let rlweq_modop = &self.pbs_info.rlwe_modop;
1020            let rlweq_nttop = &self.pbs_info.rlwe_nttop;
1021
1022            // sanity check
1023            assert!(sk_rlwe.values().len() == ring_size);
1024            assert!(sk_lwe.values().len() == self.pbs_info.parameters.lwe_n().0);
1025
1026            // auto keys
1027            let mut auto_keys = HashMap::new();
1028            let auto_gadget = self.pbs_info.auto_decomposer.gadget_vector();
1029            for i in [g, -g] {
1030                let mut ksk_out = M::zeros(
1031                    self.pbs_info.auto_decomposer.decomposition_count(),
1032                    ring_size,
1033                );
1034                galois_key_gen(
1035                    &mut ksk_out,
1036                    sk_rlwe.values(),
1037                    i,
1038                    &auto_gadget,
1039                    rlweq_modop,
1040                    rlweq_nttop,
1041                    &mut main_prng,
1042                    rng,
1043                );
1044                auto_keys.insert(i, ksk_out);
1045            }
1046
1047            // rgsw ciphertexts of lwe secret elements
1048            let rgsw_rgsw_decomposer = self
1049                .pbs_info
1050                .parameters
1051                .rgsw_rgsw_decomposer::<DefaultDecomposer<M::MatElement>>();
1052            let (rgrg_d_a, rgrg_d_b) = (
1053                rgsw_rgsw_decomposer.0.decomposition_count(),
1054                rgsw_rgsw_decomposer.1.decomposition_count(),
1055            );
1056            let (rgrg_gadget_a, rgrg_gadget_b) = (
1057                rgsw_rgsw_decomposer.0.gadget_vector(),
1058                rgsw_rgsw_decomposer.1.gadget_vector(),
1059            );
1060            let rgsw_cts = sk_lwe
1061                .values()
1062                .iter()
1063                .map(|si| {
1064                    let mut m = M::R::zeros(ring_size);
1065                    //TODO(Jay): It will be nice to have a function that returns polynomial
1066                    // (monomial infact!) corresponding to secret element embedded in ring X^{2N+1}.
1067                    // Save lots of mistakes where one forgest to emebed si in bigger ring.
1068                    let si = *si * (self.pbs_info.embedding_factor as i32);
1069                    if si < 0 {
1070                        // X^{-si} = X^{2N-si} = -X^{N-si}, assuming abs(si) < N
1071                        // (which it is given si is secret element)
1072                        m.as_mut()[ring_size - (si.abs() as usize)] = rlwe_q.neg_one();
1073                    } else {
1074                        m.as_mut()[si as usize] = M::MatElement::one();
1075                    }
1076
1077                    // public key RGSW encryption has no part that can be seeded, unlike secret key
1078                    // RGSW encryption where RLWE'_A(m) is seeded
1079                    let mut out_rgsw = M::zeros(rgrg_d_a * 2 + rgrg_d_b * 2, ring_size);
1080                    public_key_encrypt_rgsw(
1081                        &mut out_rgsw,
1082                        &m.as_ref(),
1083                        collective_pk,
1084                        &rgrg_gadget_a,
1085                        &rgrg_gadget_b,
1086                        rlweq_modop,
1087                        rlweq_nttop,
1088                        rng,
1089                    );
1090
1091                    out_rgsw
1092                })
1093                .collect_vec();
1094
1095            // LWE ksk
1096            let mut lwe_ksk =
1097                M::R::zeros(self.pbs_info.lwe_decomposer.decomposition_count() * ring_size);
1098            let lwe_modop = &self.pbs_info.lwe_modop;
1099            let d_lwe_gadget_vec = self.pbs_info.lwe_decomposer.gadget_vector();
1100            lwe_ksk_keygen(
1101                sk_rlwe.values(),
1102                sk_lwe.values(),
1103                &mut lwe_ksk,
1104                &d_lwe_gadget_vec,
1105                lwe_modop,
1106                &mut main_prng,
1107                rng,
1108            );
1109
1110            CommonReferenceSeededMultiPartyServerKeyShare {
1111                auto_keys,
1112                rgsw_cts,
1113                lwe_ksk,
1114                cr_seed,
1115                parameters: self.pbs_info.parameters.clone(),
1116            }
1117        })
1118    }
1119
1120    fn multi_party_public_key_share(
1121        &self,
1122        cr_seed: [u8; 32],
1123        client_key: &ClientKey,
1124    ) -> CommonReferenceSeededCollectivePublicKeyShare<
1125        <M as Matrix>::R,
1126        [u8; 32],
1127        BoolParameters<<M as Matrix>::MatElement>,
1128    > {
1129        DefaultSecureRng::with_local_mut(|rng| {
1130            let mut share_out = M::R::zeros(self.pbs_info.parameters.rlwe_n().0);
1131            let modop = &self.pbs_info.rlwe_modop;
1132            let nttop = &self.pbs_info.rlwe_nttop;
1133            let mut main_prng = DefaultSecureRng::new_seeded(cr_seed);
1134            public_key_share(
1135                &mut share_out,
1136                client_key.sk_rlwe.values(),
1137                modop,
1138                nttop,
1139                &mut main_prng,
1140                rng,
1141            );
1142
1143            CommonReferenceSeededCollectivePublicKeyShare {
1144                share: share_out,
1145                cr_seed: cr_seed,
1146                parameters: self.pbs_info.parameters.clone(),
1147            }
1148        })
1149    }
1150
1151    fn multi_party_decryption_share(
1152        &self,
1153        lwe_ct: &M::R,
1154        client_key: &ClientKey,
1155    ) -> MultiPartyDecryptionShare<<M as Matrix>::MatElement> {
1156        assert!(lwe_ct.as_ref().len() == self.pbs_info.parameters.rlwe_n().0 + 1);
1157        let modop = &self.pbs_info.rlwe_modop;
1158        let mut neg_s = M::R::try_convert_from(
1159            client_key.sk_rlwe.values(),
1160            &self.pbs_info.parameters.rlwe_q(),
1161        );
1162        modop.elwise_neg_mut(neg_s.as_mut());
1163
1164        let mut neg_sa = M::MatElement::zero();
1165        izip!(lwe_ct.as_ref().iter().skip(1), neg_s.as_ref().iter()).for_each(|(ai, nsi)| {
1166            neg_sa = modop.add(&neg_sa, &modop.mul(ai, nsi));
1167        });
1168
1169        let e = DefaultSecureRng::with_local_mut(|rng| {
1170            let mut e =
1171                RandomGaussianElementInModulus::random(rng, self.pbs_info.parameters.rlwe_q());
1172            e
1173        });
1174        let share = modop.add(&neg_sa, &e);
1175
1176        MultiPartyDecryptionShare { share }
1177    }
1178
1179    pub(crate) fn multi_party_decrypt(
1180        &self,
1181        shares: &[MultiPartyDecryptionShare<M::MatElement>],
1182        lwe_ct: &M::R,
1183    ) -> bool {
1184        let modop = &self.pbs_info.rlwe_modop;
1185        let mut sum_a = M::MatElement::zero();
1186        shares
1187            .iter()
1188            .for_each(|share_i| sum_a = modop.add(&sum_a, &share_i.share));
1189
1190        let encoded_m = modop.add(&lwe_ct.as_ref()[0], &sum_a);
1191        self.pbs_info.parameters.rlwe_q().decode(encoded_m)
1192    }
1193
1194    /// First encrypt as RLWE(m) with m as constant polynomial and extract it as
1195    /// LWE ciphertext
1196    pub(crate) fn pk_encrypt(&self, pk: &M, m: bool) -> M::R {
1197        DefaultSecureRng::with_local_mut(|rng| {
1198            let modop = &self.pbs_info.rlwe_modop;
1199            let nttop = &self.pbs_info.rlwe_nttop;
1200
1201            // RLWE(0)
1202            // sample ephemeral key u
1203            let ring_size = self.pbs_info.parameters.rlwe_n().0;
1204            let mut u = vec![0i32; ring_size];
1205            fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
1206            let mut u = M::R::try_convert_from(&u, &self.pbs_info.parameters.rlwe_q());
1207            nttop.forward(u.as_mut());
1208
1209            let mut ua = M::R::zeros(ring_size);
1210            ua.as_mut().copy_from_slice(pk.get_row_slice(0));
1211            let mut ub = M::R::zeros(ring_size);
1212            ub.as_mut().copy_from_slice(pk.get_row_slice(1));
1213
1214            // a*u
1215            nttop.forward(ua.as_mut());
1216            modop.elwise_mul_mut(ua.as_mut(), u.as_ref());
1217            nttop.backward(ua.as_mut());
1218
1219            // b*u
1220            nttop.forward(ub.as_mut());
1221            modop.elwise_mul_mut(ub.as_mut(), u.as_ref());
1222            nttop.backward(ub.as_mut());
1223
1224            let mut rlwe = M::zeros(2, ring_size);
1225            // sample error
1226            rlwe.iter_rows_mut().for_each(|ri| {
1227                RandomFillGaussianInModulus::random_fill(
1228                    rng,
1229                    &self.pbs_info.parameters.rlwe_q(),
1230                    ri.as_mut(),
1231                );
1232            });
1233
1234            // a*u + e0
1235            modop.elwise_add_mut(rlwe.get_row_mut(0), ua.as_ref());
1236            // b*u + e1
1237            modop.elwise_add_mut(rlwe.get_row_mut(1), ub.as_ref());
1238
1239            //FIXME(Jay): Figure out a way to get Q/8 form modulus
1240            let m = if m {
1241                // Q/8
1242                self.pbs_info.rlwe_q().true_el()
1243            } else {
1244                // -Q/8
1245                self.pbs_info.rlwe_q().false_el()
1246            };
1247
1248            // b*u + e1 + m, where m is constant polynomial
1249            rlwe.set(1, 0, modop.add(rlwe.get(1, 0), &m));
1250
1251            // sample extract index 0
1252            let mut lwe_out = M::R::zeros(ring_size + 1);
1253            sample_extract(&mut lwe_out, &rlwe, modop, 0);
1254
1255            lwe_out
1256        })
1257    }
1258
1259    /// TODO(Jay): Fetch client key from thread local
1260    pub fn sk_encrypt(&self, m: bool, client_key: &ClientKey) -> M::R {
1261        //FIXME(Jay): Figure out a way to get Q/8 form modulus
1262        let m = if m {
1263            // Q/8
1264            self.pbs_info.rlwe_q().true_el()
1265        } else {
1266            // -Q/8
1267            self.pbs_info.rlwe_q().false_el()
1268        };
1269
1270        DefaultSecureRng::with_local_mut(|rng| {
1271            let mut lwe_out = M::R::zeros(self.pbs_info.parameters.rlwe_n().0 + 1);
1272            encrypt_lwe(
1273                &mut lwe_out,
1274                &m,
1275                client_key.sk_rlwe.values(),
1276                &self.pbs_info.rlwe_modop,
1277                rng,
1278            );
1279            lwe_out
1280        })
1281    }
1282
1283    pub fn sk_decrypt(&self, lwe_ct: &M::R, client_key: &ClientKey) -> bool {
1284        let m = decrypt_lwe(
1285            lwe_ct,
1286            client_key.sk_rlwe.values(),
1287            &self.pbs_info.rlwe_modop,
1288        );
1289        self.pbs_info.rlwe_q().decode(m)
1290    }
1291
1292    fn aggregate_multi_party_server_key_shares<S>(
1293        &self,
1294        shares: &[CommonReferenceSeededMultiPartyServerKeyShare<
1295            M,
1296            BoolParameters<M::MatElement>,
1297            S,
1298        >],
1299    ) -> SeededMultiPartyServerKey<M, S, BoolParameters<M::MatElement>>
1300    where
1301        S: PartialEq + Clone,
1302        M: Clone,
1303    {
1304        assert!(shares.len() > 0);
1305        let parameters = shares[0].parameters.clone();
1306        let cr_seed = &shares[0].cr_seed;
1307
1308        let rlwe_n = parameters.rlwe_n().0;
1309        let g = parameters.g() as isize;
1310        let rlwe_q = parameters.rlwe_q();
1311        let lwe_q = parameters.lwe_q();
1312
1313        // sanity checks
1314        shares.iter().skip(1).for_each(|s| {
1315            assert!(s.parameters == parameters);
1316            assert!(&s.cr_seed == cr_seed);
1317        });
1318
1319        let rlweq_modop = &self.pbs_info.rlwe_modop;
1320        let rlweq_nttop = &self.pbs_info.rlwe_nttop;
1321
1322        // auto keys
1323        let mut auto_keys = HashMap::new();
1324        for i in [g, -g] {
1325            let mut key = M::zeros(parameters.auto_decomposition_count().0, rlwe_n);
1326
1327            shares.iter().for_each(|s| {
1328                let auto_key_share_i = s.auto_keys.get(&i).expect("Auto key {i} missing");
1329                assert!(
1330                    auto_key_share_i.dimension()
1331                        == (parameters.auto_decomposition_count().0, rlwe_n)
1332                );
1333                izip!(key.iter_rows_mut(), auto_key_share_i.iter_rows()).for_each(
1334                    |(partb_out, partb_share)| {
1335                        rlweq_modop.elwise_add_mut(partb_out.as_mut(), partb_share.as_ref());
1336                    },
1337                );
1338            });
1339
1340            auto_keys.insert(i, key);
1341        }
1342
1343        // rgsw ciphertext (most expensive part!)
1344        let lwe_n = parameters.lwe_n().0;
1345        let rgsw_by_rgsw_decomposer =
1346            parameters.rgsw_rgsw_decomposer::<DefaultDecomposer<M::MatElement>>();
1347        let mut scratch_matrix = M::zeros(
1348            std::cmp::max(
1349                rgsw_by_rgsw_decomposer.a().decomposition_count(),
1350                rgsw_by_rgsw_decomposer.b().decomposition_count(),
1351            ) + (rgsw_by_rgsw_decomposer.a().decomposition_count() * 2
1352                + rgsw_by_rgsw_decomposer.b().decomposition_count() * 2),
1353            rlwe_n,
1354        );
1355
1356        let mut tmp_rgsw =
1357            RgswCiphertext::<M, _>::empty(rlwe_n, &rgsw_by_rgsw_decomposer, rlwe_q.clone()).data;
1358        let rgsw_cts = (0..lwe_n)
1359            .into_iter()
1360            .map(|index| {
1361                // copy over rgsw ciphertext for index^th secret element from first share and
1362                // treat it as accumulating rgsw ciphertext
1363                let mut rgsw_i = shares[0].rgsw_cts[index].clone();
1364
1365                shares.iter().skip(1).for_each(|si| {
1366                    // copy over si's RGSW[index] ciphertext and send to evaluation domain
1367                    izip!(tmp_rgsw.iter_rows_mut(), si.rgsw_cts[index].iter_rows()).for_each(
1368                        |(to_ri, from_ri)| {
1369                            to_ri.as_mut().copy_from_slice(from_ri.as_ref());
1370                            rlweq_nttop.forward(to_ri.as_mut())
1371                        },
1372                    );
1373
1374                    rgsw_by_rgsw_inplace(
1375                        &mut rgsw_i,
1376                        &tmp_rgsw,
1377                        &rgsw_by_rgsw_decomposer,
1378                        &mut scratch_matrix,
1379                        rlweq_nttop,
1380                        rlweq_modop,
1381                    );
1382                });
1383
1384                rgsw_i
1385            })
1386            .collect_vec();
1387
1388        // LWE ksks
1389        let mut lwe_ksk = M::R::zeros(rlwe_n * parameters.lwe_decomposition_count().0);
1390        let lweq_modop = &self.pbs_info.lwe_modop;
1391        shares.iter().for_each(|si| {
1392            assert!(si.lwe_ksk.as_ref().len() == rlwe_n * parameters.lwe_decomposition_count().0);
1393            lweq_modop.elwise_add_mut(lwe_ksk.as_mut(), si.lwe_ksk.as_ref())
1394        });
1395
1396        SeededMultiPartyServerKey {
1397            rgsw_cts,
1398            auto_keys,
1399            lwe_ksk,
1400            cr_seed: cr_seed.clone(),
1401            parameters: parameters,
1402        }
1403    }
1404
1405    // TODO(Jay): scratch spaces must be thread local. Don't pass them as arguments
1406    pub fn nand(
1407        &mut self,
1408        c0: &M::R,
1409        c1: &M::R,
1410        server_key: &ServerKeyEvaluationDomain<M, DefaultSecureRng, NttOp>,
1411    ) -> M::R {
1412        let mut c_out = M::R::zeros(c0.as_ref().len());
1413        let modop = &self.pbs_info.rlwe_modop;
1414        izip!(
1415            c_out.as_mut().iter_mut(),
1416            c0.as_ref().iter(),
1417            c1.as_ref().iter()
1418        )
1419        .for_each(|(o, i0, i1)| {
1420            *o = modop.add(i0, i1);
1421        });
1422        // +Q/4
1423        c_out.as_mut()[0] = modop.add(&c_out.as_ref()[0], &self.pbs_info.rlwe_qby4);
1424
1425        // PBS
1426        pbs(
1427            &self.pbs_info,
1428            &self.pbs_info.nand_test_vec,
1429            &mut c_out,
1430            server_key,
1431            &mut self.scratch_memory.lwe_vector,
1432            &mut self.scratch_memory.decomposition_matrix,
1433        );
1434
1435        c_out
1436    }
1437}
1438
1439/// LMKCY+ Blind rotation
1440///
1441/// gk_to_si: [g^0, ..., g^{q/2-1}, -g^0, -g^1, .., -g^{q/2-1}]
1442fn blind_rotation<
1443    MT: IsTrivial + MatrixMut,
1444    Mmut: MatrixMut<MatElement = MT::MatElement> + Matrix,
1445    D: Decomposer<Element = MT::MatElement>,
1446    NttOp: Ntt<Element = MT::MatElement>,
1447    ModOp: ArithmeticOps<Element = MT::MatElement> + VectorOps<Element = MT::MatElement>,
1448    K: PbsKey<M = Mmut>,
1449    P: PbsInfo<Element = MT::MatElement>,
1450>(
1451    trivial_rlwe_test_poly: &mut MT,
1452    scratch_matrix: &mut Mmut,
1453    g: isize,
1454    w: usize,
1455    q: usize,
1456    gk_to_si: &[Vec<usize>],
1457    rlwe_rgsw_decomposer: &(D, D),
1458    auto_decomposer: &D,
1459    ntt_op: &NttOp,
1460    mod_op: &ModOp,
1461    parameters: &P,
1462    pbs_key: &K,
1463) where
1464    <Mmut as Matrix>::R: RowMut,
1465    Mmut::MatElement: Copy + Zero,
1466    <MT as Matrix>::R: RowMut,
1467{
1468    let q_by_2 = q / 2;
1469
1470    // -(g^k)
1471    for i in (1..q_by_2).rev() {
1472        gk_to_si[q_by_2 + i].iter().for_each(|s_index| {
1473            rlwe_by_rgsw(
1474                trivial_rlwe_test_poly,
1475                pbs_key.rgsw_ct_lwe_si(*s_index),
1476                scratch_matrix,
1477                rlwe_rgsw_decomposer,
1478                ntt_op,
1479                mod_op,
1480            );
1481        });
1482
1483        let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(g);
1484        galois_auto(
1485            trivial_rlwe_test_poly,
1486            pbs_key.galois_key_for_auto(g),
1487            scratch_matrix,
1488            &auto_map_index,
1489            &auto_map_sign,
1490            mod_op,
1491            ntt_op,
1492            auto_decomposer,
1493        );
1494    }
1495
1496    // -(g^0)
1497    gk_to_si[q_by_2].iter().for_each(|s_index| {
1498        rlwe_by_rgsw(
1499            trivial_rlwe_test_poly,
1500            pbs_key.rgsw_ct_lwe_si(*s_index),
1501            scratch_matrix,
1502            rlwe_rgsw_decomposer,
1503            ntt_op,
1504            mod_op,
1505        );
1506    });
1507    let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(-g);
1508    galois_auto(
1509        trivial_rlwe_test_poly,
1510        pbs_key.galois_key_for_auto(-g),
1511        scratch_matrix,
1512        &auto_map_index,
1513        &auto_map_sign,
1514        mod_op,
1515        ntt_op,
1516        auto_decomposer,
1517    );
1518
1519    // +(g^k)
1520    for i in (1..q_by_2).rev() {
1521        gk_to_si[i].iter().for_each(|s_index| {
1522            rlwe_by_rgsw(
1523                trivial_rlwe_test_poly,
1524                pbs_key.rgsw_ct_lwe_si(*s_index),
1525                scratch_matrix,
1526                rlwe_rgsw_decomposer,
1527                ntt_op,
1528                mod_op,
1529            );
1530        });
1531
1532        let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(g);
1533        galois_auto(
1534            trivial_rlwe_test_poly,
1535            pbs_key.galois_key_for_auto(g),
1536            scratch_matrix,
1537            &auto_map_index,
1538            &auto_map_sign,
1539            mod_op,
1540            ntt_op,
1541            auto_decomposer,
1542        );
1543    }
1544
1545    // +(g^0)
1546    gk_to_si[0].iter().for_each(|s_index| {
1547        rlwe_by_rgsw(
1548            trivial_rlwe_test_poly,
1549            pbs_key.rgsw_ct_lwe_si(gk_to_si[q_by_2][*s_index]),
1550            scratch_matrix,
1551            rlwe_rgsw_decomposer,
1552            ntt_op,
1553            mod_op,
1554        );
1555    });
1556}
1557
1558/// - Mod down
1559/// - key switching
1560/// - mod down
1561/// - blind rotate
1562fn pbs<
1563    M: Matrix + MatrixMut + MatrixEntity,
1564    P: PbsInfo<Element = M::MatElement>,
1565    K: PbsKey<M = M>,
1566>(
1567    pbs_info: &P,
1568    test_vec: &M::R,
1569    lwe_in: &mut M::R,
1570    pbs_key: &K,
1571    scratch_lwe_vec: &mut M::R,
1572    scratch_blind_rotate_matrix: &mut M,
1573) where
1574    <M as Matrix>::R: RowMut,
1575    M::MatElement: PrimInt + ToPrimitive + FromPrimitive + One + Copy + Zero + Display,
1576{
1577    let rlwe_q = pbs_info.rlwe_q();
1578    let lwe_q = pbs_info.lwe_q();
1579    let br_q = pbs_info.br_q();
1580    let rlwe_qf64 = rlwe_q.q_as_f64().unwrap();
1581    let lwe_qf64 = lwe_q.q_as_f64().unwrap();
1582    let br_qf64 = br_q.to_f64().unwrap();
1583    let rlwe_n = pbs_info.rlwe_n();
1584
1585    // PBSTracer::with_local_mut(|t| {
1586    //     let out = lwe_in
1587    //         .as_ref()
1588    //         .iter()
1589    //         .map(|v| v.to_u64().unwrap())
1590    //         .collect_vec();
1591    //     t.ct_rlwe_q_mod = out;
1592    // });
1593
1594    // moddown Q -> Q_ks
1595    lwe_in.as_mut().iter_mut().for_each(|v| {
1596        *v =
1597            M::MatElement::from_f64(((v.to_f64().unwrap() * lwe_qf64) / rlwe_qf64).round()).unwrap()
1598    });
1599
1600    // PBSTracer::with_local_mut(|t| {
1601    //     let out = lwe_in
1602    //         .as_ref()
1603    //         .iter()
1604    //         .map(|v| v.to_u64().unwrap())
1605    //         .collect_vec();
1606    //     t.ct_lwe_q_mod = out;
1607    // });
1608
1609    // key switch RLWE secret to LWE secret
1610    scratch_lwe_vec.as_mut().fill(M::MatElement::zero());
1611    lwe_key_switch(
1612        scratch_lwe_vec,
1613        lwe_in,
1614        pbs_key.lwe_ksk(),
1615        pbs_info.modop_lweq(),
1616        pbs_info.lwe_decomposer(),
1617    );
1618
1619    // PBSTracer::with_local_mut(|t| {
1620    //     let out = scratch_lwe_vec
1621    //         .as_ref()
1622    //         .iter()
1623    //         .map(|v| v.to_u64().unwrap())
1624    //         .collect_vec();
1625    //     t.ct_lwe_q_mod_after_ksk = out;
1626    // });
1627
1628    // odd mowdown Q_ks -> q
1629    let g_k_dlog_map = pbs_info.g_k_dlog_map();
1630    let mut g_k_si = vec![vec![]; br_q];
1631    scratch_lwe_vec
1632        .as_ref()
1633        .iter()
1634        .skip(1)
1635        .enumerate()
1636        .for_each(|(index, v)| {
1637            let odd_v = mod_switch_odd(v.to_f64().unwrap(), lwe_qf64, br_qf64);
1638            let k = g_k_dlog_map[odd_v];
1639            g_k_si[k].push(index);
1640        });
1641
1642    // PBSTracer::with_local_mut(|t| {
1643    //     let out = scratch_lwe_vec
1644    //         .as_ref()
1645    //         .iter()
1646    //         .map(|v| mod_switch_odd(v.to_f64().unwrap(), lwe_qf64, br_qf64) as
1647    // u64)         .collect_vec();
1648    //     t.ct_br_q_mod = out;
1649    // });
1650
1651    // handle b and set trivial test RLWE
1652    let g = pbs_info.g() as usize;
1653    let g_times_b = (g * mod_switch_odd(
1654        scratch_lwe_vec.as_ref()[0].to_f64().unwrap(),
1655        lwe_qf64,
1656        br_qf64,
1657    )) % (br_q);
1658    // v = (v(X) * X^{g*b}) mod X^{q/2}+1
1659    let br_qby2 = br_q / 2;
1660    let mut gb_monomial_sign = true;
1661    let mut gb_monomial_exp = g_times_b;
1662    // X^{g*b} mod X^{q/2}+1
1663    if gb_monomial_exp > br_qby2 {
1664        gb_monomial_exp -= br_qby2;
1665        gb_monomial_sign = false
1666    }
1667    // monomial mul
1668    let mut trivial_rlwe_test_poly = RlweCiphertext::<_, DefaultSecureRng> {
1669        data: M::zeros(2, rlwe_n),
1670        is_trivial: true,
1671        _phatom: PhantomData,
1672    };
1673    if pbs_info.embedding_factor() == 1 {
1674        monomial_mul(
1675            test_vec.as_ref(),
1676            trivial_rlwe_test_poly.get_row_mut(1).as_mut(),
1677            gb_monomial_exp,
1678            gb_monomial_sign,
1679            br_qby2,
1680            pbs_info.modop_rlweq(),
1681        );
1682    } else {
1683        // use lwe_in to store the `t = v(X) * X^{g*2} mod X^{q/2}+1` temporarily. This
1684        // works because q/2 <= N (where N is lwe_in LWE dimension) always.
1685        monomial_mul(
1686            test_vec.as_ref(),
1687            &mut lwe_in.as_mut()[..br_qby2],
1688            gb_monomial_exp,
1689            gb_monomial_sign,
1690            br_qby2,
1691            pbs_info.modop_rlweq(),
1692        );
1693
1694        // emebed poly `t` in ring X^{q/2}+1 inside the bigger ring X^{N}+1
1695        let embed_factor = pbs_info.embedding_factor();
1696        let partb_trivial_rlwe = trivial_rlwe_test_poly.get_row_mut(1);
1697        lwe_in.as_ref()[..br_qby2]
1698            .iter()
1699            .enumerate()
1700            .for_each(|(index, v)| {
1701                partb_trivial_rlwe[embed_factor * index] = *v;
1702            });
1703    }
1704
1705    // blind rotate
1706    blind_rotation(
1707        &mut trivial_rlwe_test_poly,
1708        scratch_blind_rotate_matrix,
1709        pbs_info.g(),
1710        1,
1711        br_q,
1712        &g_k_si,
1713        pbs_info.rlwe_rgsw_decomposer(),
1714        pbs_info.auto_decomposer(),
1715        pbs_info.nttop_rlweq(),
1716        pbs_info.modop_rlweq(),
1717        pbs_info,
1718        pbs_key,
1719    );
1720
1721    // ClientKey::with_local(|ck| {
1722    //     let ring_size = parameters.rlwe_n();
1723    //     let mut rlwe_ct = vec![vec![0u64; ring_size]; 2];
1724    //     izip!(
1725    //         rlwe_ct[0].iter_mut(),
1726    //         trivial_rlwe_test_poly.0.get_row_slice(0)
1727    //     )
1728    //     .for_each(|(t, f)| {
1729    //         *t = f.to_u64().unwrap();
1730    //     });
1731    //     izip!(
1732    //         rlwe_ct[1].iter_mut(),
1733    //         trivial_rlwe_test_poly.0.get_row_slice(1)
1734    //     )
1735    //     .for_each(|(t, f)| {
1736    //         *t = f.to_u64().unwrap();
1737    //     });
1738    //     let mut m_out = vec![vec![0u64; ring_size]];
1739    //     let modop = ModularOpsU64::new(rlwe_q.to_u64().unwrap());
1740    //     let nttop = NttBackendU64::new(rlwe_q.to_u64().unwrap(), ring_size);
1741    //     decrypt_rlwe(&rlwe_ct, ck.sk_rlwe.values(), &mut m_out, &nttop, &modop);
1742
1743    //     println!("RLWE post PBS message: {:?}", m_out[0]);
1744    // });
1745
1746    // sample extract
1747    sample_extract(lwe_in, &trivial_rlwe_test_poly, pbs_info.modop_rlweq(), 0);
1748}
1749
1750fn mod_switch_odd(v: f64, from_q: f64, to_q: f64) -> usize {
1751    let odd_v = (((v * to_q) / (from_q)).floor()).to_usize().unwrap();
1752    //TODO(Jay): check correctness of this
1753    odd_v + ((odd_v & 1) ^ 1)
1754}
1755
1756// TODO(Jay): Add tests for sample extract
1757fn sample_extract<M: Matrix + MatrixMut, ModOp: ArithmeticOps<Element = M::MatElement>>(
1758    lwe_out: &mut M::R,
1759    rlwe_in: &M,
1760    mod_op: &ModOp,
1761    index: usize,
1762) where
1763    <M as Matrix>::R: RowMut,
1764    M::MatElement: Copy,
1765{
1766    let ring_size = rlwe_in.dimension().1;
1767
1768    // index..=0
1769    let to = &mut lwe_out.as_mut()[1..];
1770    let from = rlwe_in.get_row_slice(0);
1771    for i in 0..index + 1 {
1772        to[i] = from[index - i];
1773    }
1774
1775    // -(N..index)
1776    for i in index + 1..ring_size {
1777        to[i] = mod_op.neg(&from[ring_size + index - i]);
1778    }
1779
1780    // set b
1781    lwe_out.as_mut()[0] = *rlwe_in.get(1, index);
1782}
1783
1784/// TODO(Jay): Write tests for monomial mul
1785fn monomial_mul<El, ModOp: ArithmeticOps<Element = El>>(
1786    p_in: &[El],
1787    p_out: &mut [El],
1788    mon_exp: usize,
1789    mon_sign: bool,
1790    ring_size: usize,
1791    mod_op: &ModOp,
1792) where
1793    El: Copy,
1794{
1795    debug_assert!(p_in.as_ref().len() == ring_size);
1796    debug_assert!(p_in.as_ref().len() == p_out.as_ref().len());
1797    debug_assert!(mon_exp < ring_size);
1798
1799    p_in.as_ref().iter().enumerate().for_each(|(index, v)| {
1800        let mut to_index = index + mon_exp;
1801        let mut to_sign = mon_sign;
1802        if to_index >= ring_size {
1803            to_index = to_index - ring_size;
1804            to_sign = !to_sign;
1805        }
1806
1807        if !to_sign {
1808            p_out.as_mut()[to_index] = mod_op.neg(v);
1809        } else {
1810            p_out.as_mut()[to_index] = *v;
1811        }
1812    });
1813}
1814
1815thread_local! {
1816    static PBS_TRACER: RefCell<PBSTracer<Vec<Vec<u64>>>> =
1817RefCell::new(PBSTracer::default()); }
1818
1819#[derive(Default)]
1820struct PBSTracer<M>
1821where
1822    M: Matrix + Default,
1823{
1824    pub(crate) ct_rlwe_q_mod: M::R,
1825    pub(crate) ct_lwe_q_mod: M::R,
1826    pub(crate) ct_lwe_q_mod_after_ksk: M::R,
1827    pub(crate) ct_br_q_mod: Vec<u64>,
1828}
1829
1830impl PBSTracer<Vec<Vec<u64>>> {
1831    fn trace(&self, parameters: &BoolParameters<u64>, sk_lwe: &[i32], sk_rlwe: &[i32]) {
1832        assert!(parameters.rlwe_n().0 == sk_rlwe.len());
1833        assert!(parameters.lwe_n().0 == sk_lwe.len());
1834
1835        let modop_rlweq = ModularOpsU64::new(*parameters.rlwe_q());
1836        // noise after mod down Q -> Q_ks
1837        let m_back0 = decrypt_lwe(&self.ct_rlwe_q_mod, sk_rlwe, &modop_rlweq);
1838
1839        let modop_lweq = ModularOpsU64::<CiphertextModulus<u64>>::new(*parameters.lwe_q());
1840        // noise after mod down Q -> Q_ks
1841        let m_back1 = decrypt_lwe(&self.ct_lwe_q_mod, sk_rlwe, &modop_lweq);
1842        // noise after key switch from RLWE -> LWE
1843        let m_back2 = decrypt_lwe(&self.ct_lwe_q_mod_after_ksk, sk_lwe, &modop_lweq);
1844
1845        // noise after mod down odd from Q_ks -> q
1846        let modop_br_q = ModularOpsU64::<u64>::new(*parameters.br_q() as u64);
1847        let m_back3 = decrypt_lwe(&self.ct_br_q_mod, sk_lwe, &modop_br_q);
1848
1849        println!(
1850            "
1851            M initial mod Q: {m_back0},
1852            M after mod down Q -> Q_ks: {m_back1},
1853            M after key switch from RLWE -> LWE: {m_back2},
1854            M after mod dwon Q_ks -> q: {m_back3}
1855        "
1856        );
1857    }
1858}
1859
1860impl WithLocal for PBSTracer<Vec<Vec<u64>>> {
1861    fn with_local<F, R>(func: F) -> R
1862    where
1863        F: Fn(&Self) -> R,
1864    {
1865        PBS_TRACER.with_borrow(|t| func(t))
1866    }
1867
1868    fn with_local_mut<F, R>(func: F) -> R
1869    where
1870        F: Fn(&mut Self) -> R,
1871    {
1872        PBS_TRACER.with_borrow_mut(|t| func(t))
1873    }
1874}
1875
1876#[cfg(test)]
1877mod tests {
1878    use std::iter::Sum;
1879
1880    use rand::{thread_rng, Rng};
1881    use rand_distr::Uniform;
1882
1883    use crate::{
1884        backend::ModularOpsU64,
1885        bool,
1886        ntt::NttBackendU64,
1887        random::DEFAULT_RNG,
1888        rgsw::{
1889            self, measure_noise, public_key_encrypt_rlwe, secret_key_encrypt_rlwe,
1890            tests::{_measure_noise_rgsw, _sk_encrypt_rlwe},
1891            RgswCiphertext, RgswCiphertextEvaluationDomain, SeededRgswCiphertext,
1892            SeededRlweCiphertext,
1893        },
1894        utils::{negacyclic_mul, Stats},
1895    };
1896
1897    use super::*;
1898
1899    #[test]
1900    fn bool_encrypt_decrypt_works() {
1901        let bool_evaluator = BoolEvaluator::<
1902            Vec<Vec<u64>>,
1903            NttBackendU64,
1904            ModularOpsU64<CiphertextModulus<u64>>,
1905            ModularOpsU64<CiphertextModulus<u64>>,
1906        >::new(SP_BOOL_PARAMS);
1907        let client_key = bool_evaluator.client_key();
1908
1909        let mut m = true;
1910        for _ in 0..1000 {
1911            let lwe_ct = bool_evaluator.sk_encrypt(m, &client_key);
1912            let m_back = bool_evaluator.sk_decrypt(&lwe_ct, &client_key);
1913            assert_eq!(m, m_back);
1914            m = !m;
1915        }
1916    }
1917
1918    #[test]
1919    fn bool_nand() {
1920        DefaultSecureRng::with_local_mut(|r| {
1921            let rng = DefaultSecureRng::new_seeded([19u8; 32]);
1922            *r = rng;
1923        });
1924
1925        let mut bool_evaluator = BoolEvaluator::<
1926            Vec<Vec<u64>>,
1927            NttBackendU64,
1928            ModularOpsU64<CiphertextModulus<u64>>,
1929            ModularOpsU64<CiphertextModulus<u64>>,
1930        >::new(SP_BOOL_PARAMS);
1931
1932        // println!("{:?}", bool_evaluator.nand_test_vec);
1933        let client_key = bool_evaluator.client_key();
1934        let seeded_server_key = bool_evaluator.server_key(&client_key);
1935        let server_key_eval_domain =
1936            ServerKeyEvaluationDomain::<_, DefaultSecureRng, NttBackendU64>::from(
1937                &seeded_server_key,
1938            );
1939
1940        let mut m0 = false;
1941        let mut m1 = true;
1942        let mut ct0 = bool_evaluator.sk_encrypt(m0, &client_key);
1943        let mut ct1 = bool_evaluator.sk_encrypt(m1, &client_key);
1944
1945        for _ in 0..1000 {
1946            let ct_back = bool_evaluator.nand(&ct0, &ct1, &server_key_eval_domain);
1947
1948            let m_out = !(m0 && m1);
1949
1950            // Trace and measure PBS noise
1951            {
1952                let noise0 = {
1953                    let ideal = if m0 {
1954                        bool_evaluator.pbs_info.parameters.rlwe_q().true_el()
1955                    } else {
1956                        bool_evaluator.pbs_info.parameters.rlwe_q().false_el()
1957                    };
1958                    let n = measure_noise_lwe(
1959                        &ct0,
1960                        client_key.sk_rlwe.values(),
1961                        &bool_evaluator.pbs_info.rlwe_modop,
1962                        &ideal,
1963                    );
1964                    let v = decrypt_lwe(
1965                        &ct0,
1966                        client_key.sk_rlwe.values(),
1967                        &bool_evaluator.pbs_info.rlwe_modop,
1968                    );
1969                    (n, v)
1970                };
1971                let noise1 = {
1972                    let ideal = if m1 {
1973                        bool_evaluator.pbs_info.parameters.rlwe_q().true_el()
1974                    } else {
1975                        bool_evaluator.pbs_info.parameters.rlwe_q().false_el()
1976                    };
1977                    let n = measure_noise_lwe(
1978                        &ct1,
1979                        client_key.sk_rlwe.values(),
1980                        &bool_evaluator.pbs_info.rlwe_modop,
1981                        &ideal,
1982                    );
1983                    let v = decrypt_lwe(
1984                        &ct1,
1985                        client_key.sk_rlwe.values(),
1986                        &bool_evaluator.pbs_info.rlwe_modop,
1987                    );
1988                    (n, v)
1989                };
1990
1991                // // Trace PBS
1992                // PBSTracer::with_local(|t| {
1993                //     t.trace(
1994                //         &SP_BOOL_PARAMS,
1995                //         &client_key.sk_lwe.values(),
1996                //         client_key.sk_rlwe.values(),
1997                //     )
1998                // });
1999
2000                // Calculate noise in ciphertext post PBS
2001                let noise_out = {
2002                    let ideal = if m_out {
2003                        bool_evaluator.pbs_info.parameters.rlwe_q().true_el()
2004                    } else {
2005                        bool_evaluator.pbs_info.parameters.rlwe_q().false_el()
2006                    };
2007                    let n = measure_noise_lwe(
2008                        &ct_back,
2009                        client_key.sk_rlwe.values(),
2010                        &bool_evaluator.pbs_info.rlwe_modop,
2011                        &ideal,
2012                    );
2013                    let v = decrypt_lwe(
2014                        &ct_back,
2015                        client_key.sk_rlwe.values(),
2016                        &bool_evaluator.pbs_info.rlwe_modop,
2017                    );
2018                    (n, v)
2019                };
2020                dbg!(m0, m1, m_out);
2021                println!(
2022                    "ct0 (noise, message): {:?}  \n ct1 (noise, message): {:?} \n PBS (noise, message): {:?}", noise0, noise1, noise_out
2023                );
2024            }
2025            let m_back = bool_evaluator.sk_decrypt(&ct_back, &client_key);
2026            assert!(m_out == m_back, "Expected {m_out}, got {m_back}");
2027            println!("----------");
2028
2029            m1 = m0;
2030            m0 = m_out;
2031
2032            ct1 = ct0;
2033            ct0 = ct_back;
2034        }
2035    }
2036
2037    #[test]
2038    fn multi_party_encryption_decryption() {
2039        let bool_evaluator = BoolEvaluator::<
2040            Vec<Vec<u64>>,
2041            NttBackendU64,
2042            ModularOpsU64<CiphertextModulus<u64>>,
2043            ModularOpsU64<CiphertextModulus<u64>>,
2044        >::new(MP_BOOL_PARAMS);
2045
2046        let no_of_parties = 500;
2047        let parties = (0..no_of_parties)
2048            .map(|_| bool_evaluator.client_key())
2049            .collect_vec();
2050
2051        let mut ideal_rlwe_sk = vec![0i32; bool_evaluator.pbs_info.rlwe_n()];
2052        parties.iter().for_each(|k| {
2053            izip!(ideal_rlwe_sk.iter_mut(), k.sk_rlwe.values()).for_each(|(ideal_i, s_i)| {
2054                *ideal_i = *ideal_i + s_i;
2055            });
2056        });
2057
2058        println!("{:?}", &ideal_rlwe_sk);
2059
2060        let mut m = true;
2061        for i in 0..100 {
2062            let pk_cr_seed = [0u8; 32];
2063
2064            let public_key_share = parties
2065                .iter()
2066                .map(|k| bool_evaluator.multi_party_public_key_share(pk_cr_seed, k))
2067                .collect_vec();
2068
2069            let collective_pk = PublicKey::<
2070                Vec<Vec<u64>>,
2071                DefaultSecureRng,
2072                ModularOpsU64<CiphertextModulus<u64>>,
2073            >::from(public_key_share.as_slice());
2074            let lwe_ct = bool_evaluator.pk_encrypt(&collective_pk.key, m);
2075
2076            let decryption_shares = parties
2077                .iter()
2078                .map(|k| bool_evaluator.multi_party_decryption_share(&lwe_ct, k))
2079                .collect_vec();
2080
2081            let m_back = bool_evaluator.multi_party_decrypt(&decryption_shares, &lwe_ct);
2082
2083            {
2084                let ideal_m = if m {
2085                    bool_evaluator.pbs_info.parameters.rlwe_q().true_el()
2086                } else {
2087                    bool_evaluator.pbs_info.parameters.rlwe_q().false_el()
2088                };
2089                let noise = measure_noise_lwe(
2090                    &lwe_ct,
2091                    &ideal_rlwe_sk,
2092                    &bool_evaluator.pbs_info.rlwe_modop,
2093                    &ideal_m,
2094                );
2095                println!("Noise: {noise}");
2096            }
2097
2098            assert_eq!(m_back, m);
2099            m = !m;
2100        }
2101    }
2102
2103    fn _collecitve_public_key_gen(rlwe_q: u64, parties_rlwe_sk: &[RlweSecret]) -> Vec<Vec<u64>> {
2104        let ring_size = parties_rlwe_sk[0].values.len();
2105        assert!(ring_size.is_power_of_two());
2106        let mut rng = DefaultSecureRng::new();
2107        let nttop = NttBackendU64::new(&rlwe_q, ring_size);
2108        let modop = ModularOpsU64::new(rlwe_q);
2109
2110        // Generate Pk shares
2111        let pk_seed = [0u8; 32];
2112        let pk_shares = parties_rlwe_sk.iter().map(|sk| {
2113            let mut p_rng = DefaultSecureRng::new_seeded(pk_seed);
2114            let mut share_out = vec![0u64; ring_size];
2115            public_key_share(
2116                &mut share_out,
2117                sk.values(),
2118                &modop,
2119                &nttop,
2120                &mut p_rng,
2121                &mut rng,
2122            );
2123            share_out
2124        });
2125
2126        let mut pk_part_b = vec![0u64; ring_size];
2127        pk_shares.for_each(|share| modop.elwise_add_mut(&mut pk_part_b, &share));
2128        let mut pk_part_a = vec![0u64; ring_size];
2129        let mut p_rng = DefaultSecureRng::new_seeded(pk_seed);
2130        RandomFillUniformInModulus::random_fill(&mut p_rng, &rlwe_q, pk_part_a.as_mut_slice());
2131
2132        vec![pk_part_a, pk_part_b]
2133    }
2134
2135    fn _multi_party_all_keygen(
2136        bool_evaluator: &BoolEvaluator<
2137            Vec<Vec<u64>>,
2138            NttBackendU64,
2139            ModularOpsU64<CiphertextModulus<u64>>,
2140            ModularOpsU64<CiphertextModulus<u64>>,
2141        >,
2142        no_of_parties: usize,
2143    ) -> (
2144        Vec<ClientKey>,
2145        PublicKey<Vec<Vec<u64>>, DefaultSecureRng, ModularOpsU64<CiphertextModulus<u64>>>,
2146        Vec<
2147            CommonReferenceSeededMultiPartyServerKeyShare<
2148                Vec<Vec<u64>>,
2149                BoolParameters<u64>,
2150                [u8; 32],
2151            >,
2152        >,
2153        SeededMultiPartyServerKey<Vec<Vec<u64>>, [u8; 32], BoolParameters<u64>>,
2154        ServerKeyEvaluationDomain<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64>,
2155        ClientKey,
2156    ) {
2157        let parties = (0..no_of_parties)
2158            .map(|_| bool_evaluator.client_key())
2159            .collect_vec();
2160
2161        let mut rng = DefaultSecureRng::new();
2162
2163        // Collective public key
2164        let mut pk_cr_seed = [0u8; 32];
2165        rng.fill_bytes(&mut pk_cr_seed);
2166        let public_key_share = parties
2167            .iter()
2168            .map(|k| bool_evaluator.multi_party_public_key_share(pk_cr_seed, k))
2169            .collect_vec();
2170        let collective_pk =
2171            PublicKey::<Vec<Vec<u64>>, DefaultSecureRng, _>::from(public_key_share.as_slice());
2172
2173        // Server key
2174        let mut pbs_cr_seed = [0u8; 32];
2175        rng.fill_bytes(&mut pbs_cr_seed);
2176        let server_key_shares = parties
2177            .iter()
2178            .map(|k| {
2179                bool_evaluator.multi_party_server_key_share(pbs_cr_seed, &collective_pk.key, k)
2180            })
2181            .collect_vec();
2182        let seeded_server_key =
2183            bool_evaluator.aggregate_multi_party_server_key_shares(&server_key_shares);
2184        let server_key_eval = ServerKeyEvaluationDomain::<_, DefaultSecureRng, NttBackendU64>::from(
2185            &seeded_server_key,
2186        );
2187
2188        // construct ideal rlwe sk for meauring noise
2189        let ideal_client_key = {
2190            let mut ideal_rlwe_sk = vec![0i32; bool_evaluator.pbs_info.rlwe_n()];
2191            parties.iter().for_each(|k| {
2192                izip!(ideal_rlwe_sk.iter_mut(), k.sk_rlwe.values()).for_each(|(ideal_i, s_i)| {
2193                    *ideal_i = *ideal_i + s_i;
2194                });
2195            });
2196            let mut ideal_lwe_sk = vec![0i32; bool_evaluator.pbs_info.lwe_n()];
2197            parties.iter().for_each(|k| {
2198                izip!(ideal_lwe_sk.iter_mut(), k.sk_lwe.values()).for_each(|(ideal_i, s_i)| {
2199                    *ideal_i = *ideal_i + s_i;
2200                });
2201            });
2202
2203            ClientKey {
2204                sk_lwe: LweSecret {
2205                    values: ideal_lwe_sk,
2206                },
2207                sk_rlwe: RlweSecret {
2208                    values: ideal_rlwe_sk,
2209                },
2210            }
2211        };
2212
2213        (
2214            parties,
2215            collective_pk,
2216            server_key_shares,
2217            seeded_server_key,
2218            server_key_eval,
2219            ideal_client_key,
2220        )
2221    }
2222
2223    #[test]
2224    fn multi_party_nand() {
2225        let mut bool_evaluator = BoolEvaluator::<
2226            Vec<Vec<u64>>,
2227            NttBackendU64,
2228            ModularOpsU64<CiphertextModulus<u64>>,
2229            ModularOpsU64<CiphertextModulus<u64>>,
2230        >::new(MP_BOOL_PARAMS);
2231
2232        let (parties, collective_pk, _, _, server_key_eval, ideal_client_key) =
2233            _multi_party_all_keygen(&bool_evaluator, 2);
2234
2235        let mut m0 = true;
2236        let mut m1 = false;
2237
2238        let mut lwe0 = bool_evaluator.pk_encrypt(&collective_pk.key, m0);
2239        let mut lwe1 = bool_evaluator.pk_encrypt(&collective_pk.key, m1);
2240
2241        for _ in 0..2000 {
2242            let lwe_out = bool_evaluator.nand(&lwe0, &lwe1, &server_key_eval);
2243
2244            let m_expected = !(m0 & m1);
2245
2246            // measure noise
2247            {
2248                let noise0 = {
2249                    let ideal = if m0 {
2250                        bool_evaluator.pbs_info.rlwe_q().true_el()
2251                    } else {
2252                        bool_evaluator.pbs_info.rlwe_q().false_el()
2253                    };
2254                    let n = measure_noise_lwe(
2255                        &lwe0,
2256                        ideal_client_key.sk_rlwe.values(),
2257                        &bool_evaluator.pbs_info.rlwe_modop,
2258                        &ideal,
2259                    );
2260                    let v = decrypt_lwe(
2261                        &lwe0,
2262                        ideal_client_key.sk_rlwe.values(),
2263                        &bool_evaluator.pbs_info.rlwe_modop,
2264                    );
2265                    (n, v)
2266                };
2267                let noise1 = {
2268                    let ideal = if m1 {
2269                        bool_evaluator.pbs_info.rlwe_q().true_el()
2270                    } else {
2271                        bool_evaluator.pbs_info.rlwe_q().false_el()
2272                    };
2273                    let n = measure_noise_lwe(
2274                        &lwe1,
2275                        ideal_client_key.sk_rlwe.values(),
2276                        &bool_evaluator.pbs_info.rlwe_modop,
2277                        &ideal,
2278                    );
2279                    let v = decrypt_lwe(
2280                        &lwe1,
2281                        ideal_client_key.sk_rlwe.values(),
2282                        &bool_evaluator.pbs_info.rlwe_modop,
2283                    );
2284                    (n, v)
2285                };
2286
2287                // // Trace PBS
2288                // PBSTracer::with_local(|t| {
2289                //     t.trace(
2290                //         &MP_BOOL_PARAMS,
2291                //         &ideal_client_key.sk_lwe.values(),
2292                //         &ideal_client_key.sk_rlwe.values(),
2293                //     )
2294                // });
2295
2296                let noise_out = {
2297                    let ideal_m = if m_expected {
2298                        bool_evaluator.pbs_info.rlwe_q().true_el()
2299                    } else {
2300                        bool_evaluator.pbs_info.rlwe_q().false_el()
2301                    };
2302                    let n = measure_noise_lwe(
2303                        &lwe_out,
2304                        ideal_client_key.sk_rlwe.values(),
2305                        &bool_evaluator.pbs_info.rlwe_modop,
2306                        &ideal_m,
2307                    );
2308                    let v = decrypt_lwe(
2309                        &lwe_out,
2310                        ideal_client_key.sk_rlwe.values(),
2311                        &bool_evaluator.pbs_info.rlwe_modop,
2312                    );
2313                    (n, v)
2314                };
2315                dbg!(m0, m1, m_expected);
2316                println!(
2317                    "ct0 (noise, message): {:?}  \n ct1 (noise, message): {:?} \n PBS (noise, message): {:?}", noise0, noise1, noise_out
2318                );
2319            }
2320
2321            // multi-party decrypt
2322            let decryption_shares = parties
2323                .iter()
2324                .map(|k| bool_evaluator.multi_party_decryption_share(&lwe_out, k))
2325                .collect_vec();
2326            let m_back = bool_evaluator.multi_party_decrypt(&decryption_shares, &lwe_out);
2327
2328            // let m_back = bool_evaluator.sk_decrypt(&lwe_out, &ideal_client_key);
2329
2330            assert!(m_expected == m_back, "Expected {m_expected}, got {m_back}");
2331            m1 = m0;
2332            m0 = m_expected;
2333
2334            lwe1 = lwe0;
2335            lwe0 = lwe_out;
2336        }
2337    }
2338
2339    #[test]
2340    fn tester() {
2341        // pub(super) const TEST_MP_BOOL_PARAMS: BoolParameters<u64> =
2342        // BoolParameters::<u64> {     rlwe_q: 1152921504606830593,
2343        //     rlwe_logq: 60,
2344        //     lwe_q: 1 << 20,
2345        //     lwe_logq: 20,
2346        //     br_q: 1 << 11,
2347        //     rlwe_n: 1 << 11,
2348        //     lwe_n: 500,
2349        //     d_rgsw: 4,
2350        //     logb_rgsw: 12,
2351        //     d_lwe: 5,
2352        //     logb_lwe: 4,
2353        //     g: 5,
2354        //     w: 1,
2355        // };
2356
2357        let bool_evaluator = BoolEvaluator::<
2358            Vec<Vec<u64>>,
2359            NttBackendU64,
2360            ModularOpsU64<CiphertextModulus<u64>>,
2361            ModularOpsU64<CiphertextModulus<u64>>,
2362        >::new(MP_BOOL_PARAMS);
2363
2364        // let (_, collective_pk, _, _, server_key_eval, ideal_client_key) =
2365        //     _multi_party_all_keygen(&bool_evaluator, 20);
2366        let no_of_parties = 2;
2367        let lwe_q = bool_evaluator.pbs_info.parameters.lwe_q();
2368        let rlwe_q = bool_evaluator.pbs_info.parameters.rlwe_q();
2369        let lwe_n = bool_evaluator.pbs_info.parameters.lwe_n().0;
2370        let rlwe_n = bool_evaluator.pbs_info.parameters.rlwe_n().0;
2371        let lwe_modop = &bool_evaluator.pbs_info.lwe_modop;
2372        let rlwe_nttop = &bool_evaluator.pbs_info.rlwe_nttop;
2373        let rlwe_modop = &bool_evaluator.pbs_info.rlwe_modop;
2374
2375        let rlwe_rgsw_decomposer = &bool_evaluator.pbs_info.rlwe_rgsw_decomposer;
2376        let rlwe_rgsw_gadget_a = rlwe_rgsw_decomposer.0.gadget_vector();
2377        let rlwe_rgsw_gadget_b = rlwe_rgsw_decomposer.1.gadget_vector();
2378
2379        // let rgsw_rgsw_decomposer = &bool_evaluator
2380        //     .pbs_info
2381        //     .parameters
2382        //     .rgsw_rgsw_decomposer::<DefaultDecomposer<u64>>();
2383        // let rgsw_rgsw_gagdet_a = rgsw_rgsw_decomposer.a().gadget_vector();
2384        // let rgsw_rgsw_gagdet_b = rgsw_rgsw_decomposer.b().gadget_vector();
2385
2386        let parties = (0..no_of_parties)
2387            .map(|_| bool_evaluator.client_key())
2388            .collect_vec();
2389
2390        let ideal_client_key = {
2391            let mut ideal_rlwe_sk = vec![0i32; bool_evaluator.pbs_info.rlwe_n()];
2392            parties.iter().for_each(|k| {
2393                izip!(ideal_rlwe_sk.iter_mut(), k.sk_rlwe.values()).for_each(|(ideal_i, s_i)| {
2394                    *ideal_i = *ideal_i + s_i;
2395                });
2396            });
2397            let mut ideal_lwe_sk = vec![0i32; bool_evaluator.pbs_info.lwe_n()];
2398            parties.iter().for_each(|k| {
2399                izip!(ideal_lwe_sk.iter_mut(), k.sk_lwe.values()).for_each(|(ideal_i, s_i)| {
2400                    *ideal_i = *ideal_i + s_i;
2401                });
2402            });
2403
2404            ClientKey {
2405                sk_lwe: LweSecret {
2406                    values: ideal_lwe_sk,
2407                },
2408                sk_rlwe: RlweSecret {
2409                    values: ideal_rlwe_sk,
2410                },
2411            }
2412        };
2413
2414        // check noise in freshly encrypted RLWE ciphertext (ie var_fresh)
2415        if false {
2416            let mut rng = DefaultSecureRng::new();
2417            let mut check = Stats { samples: vec![] };
2418            for _ in 0..10 {
2419                // generate a new collective public key
2420                let mut pk_cr_seed = [0u8; 32];
2421                rng.fill_bytes(&mut pk_cr_seed);
2422                let public_key_share = parties
2423                    .iter()
2424                    .map(|k| bool_evaluator.multi_party_public_key_share(pk_cr_seed, k))
2425                    .collect_vec();
2426                let collective_pk = PublicKey::<
2427                    Vec<Vec<u64>>,
2428                    DefaultSecureRng,
2429                    ModularOpsU64<CiphertextModulus<u64>>,
2430                >::from(public_key_share.as_slice());
2431
2432                let mut m = vec![0u64; rlwe_n];
2433                RandomFillUniformInModulus::random_fill(&mut rng, rlwe_q, m.as_mut_slice());
2434                let mut rlwe_ct = vec![vec![0u64; rlwe_n]; 2];
2435                public_key_encrypt_rlwe::<_, _, _, _, i32, _>(
2436                    &mut rlwe_ct,
2437                    &collective_pk.key,
2438                    &m,
2439                    rlwe_modop,
2440                    rlwe_nttop,
2441                    &mut rng,
2442                );
2443
2444                let mut m_back = vec![0u64; rlwe_n];
2445                decrypt_rlwe(
2446                    &rlwe_ct,
2447                    ideal_client_key.sk_rlwe.values(),
2448                    &mut m_back,
2449                    rlwe_nttop,
2450                    rlwe_modop,
2451                );
2452
2453                rlwe_modop.elwise_sub_mut(m_back.as_mut_slice(), m.as_slice());
2454
2455                check.add_more(Vec::<i64>::try_convert_from(&m_back, rlwe_q).as_slice());
2456            }
2457
2458            println!("Public key Std: {}", check.std_dev().abs().log2());
2459        }
2460
2461        if true {
2462            // Generate server key shares
2463            let mut rng = DefaultSecureRng::new();
2464            let mut pk_cr_seed = [0u8; 32];
2465            rng.fill_bytes(&mut pk_cr_seed);
2466            let public_key_share = parties
2467                .iter()
2468                .map(|k| bool_evaluator.multi_party_public_key_share(pk_cr_seed, k))
2469                .collect_vec();
2470            let collective_pk = PublicKey::<
2471                Vec<Vec<u64>>,
2472                DefaultSecureRng,
2473                ModularOpsU64<CiphertextModulus<u64>>,
2474            >::from(public_key_share.as_slice());
2475
2476            let pbs_cr_seed = [0u8; 32];
2477            rng.fill_bytes(&mut pk_cr_seed);
2478            let server_key_shares = parties
2479                .iter()
2480                .map(|k| {
2481                    bool_evaluator.multi_party_server_key_share(pbs_cr_seed, &collective_pk.key, k)
2482                })
2483                .collect_vec();
2484
2485            let seeded_server_key =
2486                bool_evaluator.aggregate_multi_party_server_key_shares(&server_key_shares);
2487
2488            // Check noise in RGSW ciphertexts of ideal LWE secret elements
2489            if true {
2490                let mut check = Stats { samples: vec![] };
2491                izip!(
2492                    ideal_client_key.sk_lwe.values.iter(),
2493                    seeded_server_key.rgsw_cts.iter()
2494                )
2495                .for_each(|(s_i, rgsw_ct_i)| {
2496                    // X^{s[i]}
2497                    let mut m_si = vec![0u64; rlwe_n];
2498                    let s_i = *s_i * (bool_evaluator.pbs_info.embedding_factor as i32);
2499                    if s_i < 0 {
2500                        m_si[rlwe_n - (s_i.abs() as usize)] = rlwe_q.neg_one();
2501                    } else {
2502                        m_si[s_i as usize] = 1;
2503                    }
2504
2505                    // RLWE(-sm)
2506                    let mut neg_s_eval =
2507                        Vec::<u64>::try_convert_from(ideal_client_key.sk_rlwe.values(), rlwe_q);
2508                    rlwe_modop.elwise_neg_mut(&mut neg_s_eval);
2509                    rlwe_nttop.forward(&mut neg_s_eval);
2510                    for j in 0..rlwe_rgsw_decomposer.a().decomposition_count() {
2511                        // -s[X]*X^{s_lwe[i]}*B_j
2512                        let mut m_ideal = m_si.clone();
2513                        rlwe_nttop.forward(m_ideal.as_mut_slice());
2514                        rlwe_modop.elwise_mul_mut(m_ideal.as_mut_slice(), neg_s_eval.as_slice());
2515                        rlwe_nttop.backward(m_ideal.as_mut_slice());
2516                        rlwe_modop
2517                            .elwise_scalar_mul_mut(m_ideal.as_mut_slice(), &rlwe_rgsw_gadget_a[j]);
2518
2519                        // RLWE(-s*X^{s_lwe[i]}*B_j)
2520                        let mut rlwe_ct = vec![vec![0u64; rlwe_n]; 2];
2521                        rlwe_ct[0].copy_from_slice(&rgsw_ct_i[j]);
2522                        rlwe_ct[1].copy_from_slice(
2523                            &rgsw_ct_i[j + rlwe_rgsw_decomposer.a().decomposition_count()],
2524                        );
2525
2526                        let mut m_back = vec![0u64; rlwe_n];
2527                        decrypt_rlwe(
2528                            &rlwe_ct,
2529                            ideal_client_key.sk_rlwe.values(),
2530                            &mut m_back,
2531                            rlwe_nttop,
2532                            rlwe_modop,
2533                        );
2534
2535                        // diff
2536                        rlwe_modop.elwise_sub_mut(&mut m_back, &m_ideal);
2537                        check.add_more(&Vec::<i64>::try_convert_from(&m_back, rlwe_q));
2538                    }
2539
2540                    // RLWE'(m)
2541                    for j in 0..rlwe_rgsw_decomposer.b().decomposition_count() {
2542                        // X^{s_lwe[i]}*B_j
2543                        let mut m_ideal = m_si.clone();
2544                        rlwe_modop
2545                            .elwise_scalar_mul_mut(m_ideal.as_mut_slice(), &rlwe_rgsw_gadget_b[j]);
2546
2547                        // RLWE(X^{s_lwe[i]}*B_j)
2548                        let mut rlwe_ct = vec![vec![0u64; rlwe_n]; 2];
2549                        rlwe_ct[0].copy_from_slice(
2550                            &rgsw_ct_i[j + (2 * rlwe_rgsw_decomposer.a().decomposition_count())],
2551                        );
2552                        rlwe_ct[1].copy_from_slice(
2553                            &rgsw_ct_i[j
2554                                + (2 * rlwe_rgsw_decomposer.a().decomposition_count()
2555                                    + rlwe_rgsw_decomposer.b().decomposition_count())],
2556                        );
2557
2558                        let mut m_back = vec![0u64; rlwe_n];
2559                        decrypt_rlwe(
2560                            &rlwe_ct,
2561                            ideal_client_key.sk_rlwe.values(),
2562                            &mut m_back,
2563                            rlwe_nttop,
2564                            rlwe_modop,
2565                        );
2566
2567                        // diff
2568                        rlwe_modop.elwise_sub_mut(&mut m_back, &m_ideal);
2569                        check.add_more(&Vec::<i64>::try_convert_from(&m_back, rlwe_q));
2570                    }
2571                });
2572                println!(
2573                    "RGSW Std: {} {} ;; max={}",
2574                    check.mean(),
2575                    check.std_dev().abs().log2(),
2576                    check.samples.iter().max().unwrap()
2577                );
2578            }
2579
2580            // check noise in RLWE x RGSW(X^{s_i}) where RGSW is accunulated RGSW ciphertext
2581            if false {
2582                let mut check = Stats { samples: vec![] };
2583                // server key in Evaluation domain
2584                let server_key_eval_domain =
2585                    ServerKeyEvaluationDomain::<_, DefaultSecureRng, NttBackendU64>::from(
2586                        &seeded_server_key,
2587                    );
2588                izip!(
2589                    ideal_client_key.sk_lwe.values(),
2590                    seeded_server_key.rgsw_cts.iter()
2591                )
2592                .for_each(|(s_i, rgsw_ct_i)| {
2593                    let mut rgsw_ct_i = rgsw_ct_i.clone();
2594                    rgsw_ct_i
2595                        .iter_mut()
2596                        .for_each(|ri| rlwe_nttop.forward(ri.as_mut()));
2597
2598                    let mut m = vec![0u64; rlwe_n];
2599                    RandomFillUniformInModulus::random_fill(&mut rng, rlwe_q, m.as_mut_slice());
2600                    let mut rlwe_ct = vec![vec![0u64; rlwe_n]; 2];
2601                    public_key_encrypt_rlwe::<_, _, _, _, i32, _>(
2602                        &mut rlwe_ct,
2603                        &collective_pk.key,
2604                        &m,
2605                        rlwe_modop,
2606                        rlwe_nttop,
2607                        &mut rng,
2608                    );
2609
2610                    // RLWE(m*X^{s[i]}) = RLWE(m) x RGSW(X^{s[i]})
2611                    let mut rlwe_after = RlweCiphertext::<_, DefaultSecureRng>::new_trivial(vec![
2612                        vec![0u64; rlwe_n],
2613                        m.clone(),
2614                    ]);
2615                    // let mut rlwe_after =
2616                    //     RlweCiphertext::<_, DefaultSecureRng>::from_raw(rlwe_ct.clone(), false);
2617                    let mut scratch = vec![
2618                        vec![0u64; rlwe_n];
2619                        std::cmp::max(
2620                            rlwe_rgsw_decomposer.0.decomposition_count(),
2621                            rlwe_rgsw_decomposer.1.decomposition_count()
2622                        ) + 2
2623                    ];
2624                    rlwe_by_rgsw(
2625                        &mut rlwe_after,
2626                        &rgsw_ct_i,
2627                        &mut scratch,
2628                        rlwe_rgsw_decomposer,
2629                        rlwe_nttop,
2630                        rlwe_modop,
2631                    );
2632
2633                    // m1 = X^{s[i]}
2634                    let mut m1 = vec![0u64; rlwe_n];
2635                    let s_i = *s_i * (bool_evaluator.pbs_info.embedding_factor as i32);
2636                    if s_i < 0 {
2637                        m1[rlwe_n - (s_i.abs() as usize)] = rlwe_q.neg_one()
2638                    } else {
2639                        m1[s_i as usize] = 1;
2640                    }
2641
2642                    // (m+e) * m1
2643                    let mut m_plus_e_times_m1 = m.clone();
2644                    // decrypt_rlwe(
2645                    //     &rlwe_ct,
2646                    //     ideal_client_key.sk_rlwe.values(),
2647                    //     &mut m_plus_e_times_m1,
2648                    //     rlwe_nttop,
2649                    //     rlwe_modop,
2650                    // );
2651                    rlwe_nttop.forward(m_plus_e_times_m1.as_mut_slice());
2652                    rlwe_nttop.forward(m1.as_mut_slice());
2653                    rlwe_modop.elwise_mul_mut(m_plus_e_times_m1.as_mut_slice(), m1.as_slice());
2654                    rlwe_nttop.backward(m_plus_e_times_m1.as_mut_slice());
2655
2656                    // Resulting RLWE ciphertext will equal: (m0m1 + em1) + e_{rlsw x rgsw}.
2657                    // Hence, resulting rlwe ciphertext will have error em1 + e_{rlwe x rgsw}.
2658                    // Here we're only concerned with e_{rlwe x rgsw}, that is noise caused due to
2659                    // RLWExRGSW. Also note, in practice m1 is a monomial, for ex, X^{s_{i}}, for
2660                    // some i and var(em1) = var(e).
2661                    let mut m_plus_e_times_m1_more_e = vec![0u64; rlwe_n];
2662                    decrypt_rlwe(
2663                        &rlwe_after,
2664                        ideal_client_key.sk_rlwe.values(),
2665                        &mut m_plus_e_times_m1_more_e,
2666                        rlwe_nttop,
2667                        rlwe_modop,
2668                    );
2669
2670                    // diff
2671                    rlwe_modop.elwise_sub_mut(
2672                        m_plus_e_times_m1_more_e.as_mut_slice(),
2673                        m_plus_e_times_m1.as_slice(),
2674                    );
2675
2676                    let noise = measure_noise(
2677                        &rlwe_after,
2678                        &m_plus_e_times_m1,
2679                        rlwe_nttop,
2680                        rlwe_modop,
2681                        ideal_client_key.sk_rlwe.values(),
2682                    );
2683                    print!("NOISE: {}", noise);
2684
2685                    check.add_more(&Vec::<i64>::try_convert_from(
2686                        &m_plus_e_times_m1_more_e,
2687                        rlwe_q,
2688                    ));
2689                });
2690                println!(
2691                    "RLWE x RGSW, where RGSW has noise var_brk, std: {} {}",
2692                    check.std_dev(),
2693                    check.std_dev().abs().log2()
2694                )
2695            }
2696        }
2697
2698        // Check noise in fresh RGSW ciphertexts, ie X^{s_j[i]}, must equal noise in
2699        // fresh RLWE ciphertext
2700        if true {}
2701        // test LWE ksk from RLWE -> LWE
2702        // if false {
2703        //     let logp = 2;
2704        //     let mut rng = DefaultSecureRng::new();
2705
2706        //     let m = 1;
2707        //     let encoded_m = m << (lwe_logq - logp);
2708
2709        //     // Encrypt
2710        //     let mut lwe_ct = vec![0u64; rlwe_n + 1];
2711        //     encrypt_lwe(
2712        //         &mut lwe_ct,
2713        //         &encoded_m,
2714        //         ideal_client_key.sk_rlwe.values(),
2715        //         lwe_modop,
2716        //         &mut rng,
2717        //     );
2718
2719        //     // key switch
2720        //     let lwe_decomposer = &bool_evaluator.decomposer_lwe;
2721        //     let mut lwe_out = vec![0u64; lwe_n + 1];
2722        //     lwe_key_switch(
2723        //         &mut lwe_out,
2724        //         &lwe_ct,
2725        //         &server_key_eval.lwe_ksk,
2726        //         lwe_modop,
2727        //         lwe_decomposer,
2728        //     );
2729
2730        //     let encoded_m_back = decrypt_lwe(&lwe_out,
2731        // ideal_client_key.sk_lwe.values(), lwe_modop);     let m_back
2732        // =         ((encoded_m_back as f64 * (1 << logp) as f64) /
2733        // (lwe_q as f64)).round() as u64;     dbg!(m_back, m);
2734
2735        //     let noise = measure_noise_lwe(
2736        //         &lwe_out,
2737        //         ideal_client_key.sk_lwe.values(),
2738        //         lwe_modop,
2739        //         &encoded_m,
2740        //     );
2741
2742        //     println!("Noise: {noise}");
2743        // }
2744
2745        // Measure noise in RGSW ciphertexts of ideal LWE secrets
2746        // if true {
2747        //     let gadget_vec = gadget_vector(
2748        //         bool_evaluator.parameters.rlwe_logq,
2749        //         bool_evaluator.parameters.logb_rgsw,
2750        //         bool_evaluator.parameters.d_rgsw,
2751        //     );
2752
2753        //     for i in 0..20 {
2754        //         // measure noise in RGSW(s[i])
2755        //         let si =
2756        //             ideal_client_key.sk_lwe.values[i] *
2757        // (bool_evaluator.embedding_factor as i32);         let mut
2758        // si_poly = vec![0u64; rlwe_n];         if si < 0 {
2759        //             si_poly[rlwe_n - (si.abs() as usize)] = rlwe_q - 1;
2760        //         } else {
2761        //             si_poly[(si.abs() as usize)] = 1;
2762        //         }
2763
2764        //         let mut rgsw_si = server_key_eval.rgsw_cts[i].clone();
2765        //         rgsw_si
2766        //             .iter_mut()
2767        //             .for_each(|ri| rlwe_nttop.backward(ri.as_mut()));
2768
2769        //         println!("####### Noise in RGSW(X^s_{i}) #######");
2770        //         _measure_noise_rgsw(
2771        //             &rgsw_si,
2772        //             &si_poly,
2773        //             ideal_client_key.sk_rlwe.values(),
2774        //             &gadget_vec,
2775        //             rlwe_q,
2776        //         );
2777        //         println!("####### ##################### #######");
2778        //     }
2779        // }
2780
2781        // // measure noise grwoth in RLWExRGSW
2782        // if true {
2783        //     let mut rng = DefaultSecureRng::new();
2784        //     let mut carry_m = vec![0u64; rlwe_n];
2785        //     RandomUniformDist1::random_fill(&mut rng, &rlwe_q,
2786        // carry_m.as_mut_slice());
2787
2788        //     // RGSW(carrym)
2789        //     let trivial_rlwect = vec![vec![0u64; rlwe_n], carry_m.clone()];
2790        //     let mut rlwe_ct = RlweCiphertext::<_,
2791        // DefaultSecureRng>::from_raw(trivial_rlwect, true);
2792
2793        //     let mut scratch_matrix_dplus2_ring = vec![vec![0u64; rlwe_n];
2794        // d_rgsw + 2];     let mul_mod =
2795        //         |v0: &u64, v1: &u64| (((*v0 as u128 * *v1 as u128) % (rlwe_q as u128)) as u64);
2796
2797        //     for i in 0..bool_evaluator.parameters.lwe_n {
2798        //         rlwe_by_rgsw(
2799        //             &mut rlwe_ct,
2800        //             server_key_eval.rgsw_ct_lwe_si(i),
2801        //             &mut scratch_matrix_dplus2_ring,
2802        //             rlwe_decomposer,
2803        //             rlwe_nttop,
2804        //             rlwe_modop,
2805        //         );
2806
2807        //         // carry_m[X] * s_i[X]
2808        //         let si =
2809        //             ideal_client_key.sk_lwe.values[i] *
2810        // (bool_evaluator.embedding_factor as i32);         let mut
2811        // si_poly = vec![0u64; rlwe_n];         if si < 0 {
2812        //             si_poly[rlwe_n - (si.abs() as usize)] = rlwe_q - 1;
2813        //         } else {
2814        //             si_poly[(si.abs() as usize)] = 1;
2815        //         }
2816        //         carry_m = negacyclic_mul(&carry_m, &si_poly, mul_mod,
2817        // rlwe_q);
2818
2819        //         let noise = measure_noise(
2820        //             &rlwe_ct,
2821        //             &carry_m,
2822        //             rlwe_nttop,
2823        //             rlwe_modop,
2824        //             ideal_client_key.sk_rlwe.values(),
2825        //         );
2826        //         println!("Noise RLWE(carry_m) accumulating {i}^th secret
2827        // monomial: {noise}");     }
2828        // }
2829
2830        // // Check galois keys
2831        // if false {
2832        //     let g = bool_evaluator.g() as isize;
2833        //     let mut rng = DefaultSecureRng::new();
2834        //     let mut scratch_matrix_dplus2_ring = vec![vec![0u64; rlwe_n];
2835        // d_rgsw + 2];     for i in [g, -g] {
2836        //         let mut m = vec![0u64; rlwe_n];
2837        //         RandomUniformDist1::random_fill(&mut rng, &rlwe_q,
2838        // m.as_mut_slice());         let mut rlwe_ct = {
2839        //             let mut data = vec![vec![0u64; rlwe_n]; 2];
2840        //             public_key_encrypt_rlwe(
2841        //                 &mut data,
2842        //                 &collective_pk.key,
2843        //                 &m,
2844        //                 rlwe_modop,
2845        //                 rlwe_nttop,
2846        //                 &mut rng,
2847        //             );
2848        //             RlweCiphertext::<_, DefaultSecureRng>::from_raw(data,
2849        // false)         };
2850
2851        //         let auto_key = server_key_eval.galois_key_for_auto(i);
2852        //         let (auto_map_index, auto_map_sign) =
2853        // generate_auto_map(rlwe_n, i);         galois_auto(
2854        //             &mut rlwe_ct,
2855        //             auto_key,
2856        //             &mut scratch_matrix_dplus2_ring,
2857        //             &auto_map_index,
2858        //             &auto_map_sign,
2859        //             rlwe_modop,
2860        //             rlwe_nttop,
2861        //             rlwe_decomposer,
2862        //         );
2863
2864        //         // send m(X) -> m(X^i)
2865        //         let mut m_k = vec![0u64; rlwe_n];
2866        //         izip!(m.iter(), auto_map_index.iter(),
2867        // auto_map_sign.iter()).for_each(             |(mi, to_index, to_sign)|
2868        // {                 if !to_sign {
2869        //                     m_k[*to_index] = rlwe_q - *mi;
2870        //                 } else {
2871        //                     m_k[*to_index] = *mi;
2872        //                 }
2873        //             },
2874        //         );
2875
2876        //         // measure noise
2877        //         let noise = measure_noise(
2878        //             &rlwe_ct,
2879        //             &m_k,
2880        //             rlwe_nttop,
2881        //             rlwe_modop,
2882        //             ideal_client_key.sk_rlwe.values(),
2883        //         );
2884
2885        //         println!("Noise after auto k={i}: {noise}");
2886        //     }
2887        // }
2888    }
2889}