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