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