arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/bool/keys.rs
src/bool/keys.rsRaw
   1use std::{collections::HashMap, hash::Hash, marker::PhantomData};
   2
   3use crate::{
   4    backend::{ModInit, VectorOps},
   5    lwe::LweSecret,
   6    pbs::WithShoupRepr,
   7    random::{NewWithSeed, RandomFillUniformInModulus},
   8    rgsw::RlweSecret,
   9    utils::{ToShoup, WithLocal},
  10    Decryptor, Encryptor, Matrix, MatrixEntity, MatrixMut, MultiPartyDecryptor, RowEntity, RowMut,
  11};
  12
  13use super::{parameters, BoolEvaluator, BoolParameters, CiphertextModulus};
  14
  15trait SinglePartyClientKey {
  16    type Element;
  17    fn sk_rlwe(&self) -> &[Self::Element];
  18    fn sk_lwe(&self) -> &[Self::Element];
  19}
  20
  21trait InteractiveMultiPartyClientKey {
  22    type Element;
  23    fn sk_rlwe(&self) -> &[Self::Element];
  24    fn sk_lwe(&self) -> &[Self::Element];
  25}
  26
  27trait NonInteractiveMultiPartyClientKey {
  28    type Element;
  29    fn sk_rlwe(&self) -> &[Self::Element];
  30    fn sk_u_rlwe(&self) -> &[Self::Element];
  31    fn sk_lwe(&self) -> &[Self::Element];
  32}
  33
  34/// Client key with RLWE and LWE secrets
  35#[derive(Clone)]
  36pub struct ClientKey {
  37    sk_rlwe: RlweSecret,
  38    sk_lwe: LweSecret,
  39}
  40
  41/// Client key with RLWE and LWE secrets
  42#[derive(Clone)]
  43pub struct ThrowMeAwayKey {
  44    sk_rlwe: RlweSecret,
  45    sk_u_rlwe: RlweSecret,
  46    sk_lwe: LweSecret,
  47}
  48
  49mod impl_ck {
  50    use super::*;
  51
  52    // Client key
  53    impl ClientKey {
  54        pub(in super::super) fn new(sk_rlwe: RlweSecret, sk_lwe: LweSecret) -> Self {
  55            Self { sk_rlwe, sk_lwe }
  56        }
  57
  58        pub(in super::super) fn sk_rlwe(&self) -> &RlweSecret {
  59            &self.sk_rlwe
  60        }
  61
  62        pub(in super::super) fn sk_lwe(&self) -> &LweSecret {
  63            &self.sk_lwe
  64        }
  65    }
  66
  67    // Client key
  68    impl ThrowMeAwayKey {
  69        pub(in super::super) fn new(
  70            sk_rlwe: RlweSecret,
  71            sk_u_rlwe: RlweSecret,
  72            sk_lwe: LweSecret,
  73        ) -> Self {
  74            Self {
  75                sk_rlwe,
  76                sk_u_rlwe,
  77                sk_lwe,
  78            }
  79        }
  80
  81        pub(in super::super) fn sk_rlwe(&self) -> &RlweSecret {
  82            &self.sk_rlwe
  83        }
  84
  85        pub(in super::super) fn sk_u_rlwe(&self) -> &RlweSecret {
  86            &self.sk_u_rlwe
  87        }
  88
  89        pub(in super::super) fn sk_lwe(&self) -> &LweSecret {
  90            &self.sk_lwe
  91        }
  92    }
  93
  94    impl Encryptor<bool, Vec<u64>> for ClientKey {
  95        fn encrypt(&self, m: &bool) -> Vec<u64> {
  96            BoolEvaluator::with_local(|e| e.sk_encrypt(*m, self))
  97        }
  98    }
  99
 100    impl Decryptor<bool, Vec<u64>> for ClientKey {
 101        fn decrypt(&self, c: &Vec<u64>) -> bool {
 102            BoolEvaluator::with_local(|e| e.sk_decrypt(c, self))
 103        }
 104    }
 105
 106    impl MultiPartyDecryptor<bool, Vec<u64>> for ClientKey {
 107        type DecryptionShare = u64;
 108
 109        fn gen_decryption_share(&self, c: &Vec<u64>) -> Self::DecryptionShare {
 110            BoolEvaluator::with_local(|e| e.multi_party_decryption_share(c, &self))
 111        }
 112
 113        fn aggregate_decryption_shares(
 114            &self,
 115            c: &Vec<u64>,
 116            shares: &[Self::DecryptionShare],
 117        ) -> bool {
 118            BoolEvaluator::with_local(|e| e.multi_party_decrypt(shares, c))
 119        }
 120    }
 121}
 122
 123/// Public key
 124pub struct PublicKey<M, Rng, ModOp> {
 125    key: M,
 126    _phantom: PhantomData<(Rng, ModOp)>,
 127}
 128
 129pub(super) mod impl_pk {
 130    use super::*;
 131
 132    impl<M, R, Mo> PublicKey<M, R, Mo> {
 133        pub(in super::super) fn key(&self) -> &M {
 134            &self.key
 135        }
 136    }
 137
 138    impl<Rng, ModOp> Encryptor<bool, Vec<u64>> for PublicKey<Vec<Vec<u64>>, Rng, ModOp> {
 139        fn encrypt(&self, m: &bool) -> Vec<u64> {
 140            BoolEvaluator::with_local(|e| e.pk_encrypt(&self.key, *m))
 141        }
 142    }
 143
 144    impl<Rng, ModOp> Encryptor<[bool], Vec<Vec<u64>>> for PublicKey<Vec<Vec<u64>>, Rng, ModOp> {
 145        fn encrypt(&self, m: &[bool]) -> Vec<Vec<u64>> {
 146            BoolEvaluator::with_local(|e| e.pk_encrypt_batched(&self.key, m))
 147        }
 148    }
 149
 150    impl<
 151            M: MatrixMut + MatrixEntity,
 152            Rng: NewWithSeed
 153                + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
 154            ModOp,
 155        > From<SeededPublicKey<M::R, Rng::Seed, BoolParameters<M::MatElement>, ModOp>>
 156        for PublicKey<M, Rng, ModOp>
 157    where
 158        <M as Matrix>::R: RowMut,
 159        M::MatElement: Copy,
 160    {
 161        fn from(
 162            value: SeededPublicKey<M::R, Rng::Seed, BoolParameters<M::MatElement>, ModOp>,
 163        ) -> Self {
 164            let mut prng = Rng::new_with_seed(value.seed);
 165
 166            let mut key = M::zeros(2, value.parameters.rlwe_n().0);
 167            // sample A
 168            RandomFillUniformInModulus::random_fill(
 169                &mut prng,
 170                value.parameters.rlwe_q(),
 171                key.get_row_mut(0),
 172            );
 173            // Copy over B
 174            key.get_row_mut(1).copy_from_slice(value.part_b.as_ref());
 175
 176            PublicKey {
 177                key,
 178                _phantom: PhantomData,
 179            }
 180        }
 181    }
 182
 183    impl<
 184            M: MatrixMut + MatrixEntity,
 185            Rng: NewWithSeed
 186                + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
 187            ModOp: VectorOps<Element = M::MatElement> + ModInit<M = CiphertextModulus<M::MatElement>>,
 188        >
 189        From<
 190            &[CommonReferenceSeededCollectivePublicKeyShare<
 191                M::R,
 192                Rng::Seed,
 193                BoolParameters<M::MatElement>,
 194            >],
 195        > for PublicKey<M, Rng, ModOp>
 196    where
 197        <M as Matrix>::R: RowMut,
 198        Rng::Seed: Copy + PartialEq,
 199        M::MatElement: PartialEq + Copy,
 200    {
 201        fn from(
 202            value: &[CommonReferenceSeededCollectivePublicKeyShare<
 203                M::R,
 204                Rng::Seed,
 205                BoolParameters<M::MatElement>,
 206            >],
 207        ) -> Self {
 208            assert!(value.len() > 0);
 209
 210            let parameters = &value[0].parameters;
 211            let mut key = M::zeros(2, parameters.rlwe_n().0);
 212
 213            // sample A
 214            let seed = value[0].cr_seed;
 215            let mut main_rng = Rng::new_with_seed(seed);
 216            RandomFillUniformInModulus::random_fill(
 217                &mut main_rng,
 218                parameters.rlwe_q(),
 219                key.get_row_mut(0),
 220            );
 221
 222            // Sum all Bs
 223            let rlweq_modop = ModOp::new(parameters.rlwe_q().clone());
 224            value.iter().for_each(|share_i| {
 225                assert!(share_i.cr_seed == seed);
 226                assert!(&share_i.parameters == parameters);
 227
 228                rlweq_modop.elwise_add_mut(key.get_row_mut(1), share_i.share.as_ref());
 229            });
 230
 231            PublicKey {
 232                key,
 233                _phantom: PhantomData,
 234            }
 235        }
 236    }
 237}
 238
 239/// Seeded public key
 240struct SeededPublicKey<Ro, S, P, ModOp> {
 241    part_b: Ro,
 242    seed: S,
 243    parameters: P,
 244    _phantom: PhantomData<ModOp>,
 245}
 246
 247mod impl_seeded_pk {
 248    use super::*;
 249
 250    impl<R, S, ModOp>
 251        From<&[CommonReferenceSeededCollectivePublicKeyShare<R, S, BoolParameters<R::Element>>]>
 252        for SeededPublicKey<R, S, BoolParameters<R::Element>, ModOp>
 253    where
 254        ModOp: VectorOps<Element = R::Element> + ModInit<M = CiphertextModulus<R::Element>>,
 255        S: PartialEq + Clone,
 256        R: RowMut + RowEntity + Clone,
 257        R::Element: Clone + PartialEq,
 258    {
 259        fn from(
 260            value: &[CommonReferenceSeededCollectivePublicKeyShare<
 261                R,
 262                S,
 263                BoolParameters<R::Element>,
 264            >],
 265        ) -> Self {
 266            assert!(value.len() > 0);
 267
 268            let parameters = &value[0].parameters;
 269            let cr_seed = value[0].cr_seed.clone();
 270
 271            // Sum all Bs
 272            let rlweq_modop = ModOp::new(parameters.rlwe_q().clone());
 273            let mut part_b = value[0].share.clone();
 274            value.iter().skip(1).for_each(|share_i| {
 275                assert!(&share_i.cr_seed == &cr_seed);
 276                assert!(&share_i.parameters == parameters);
 277
 278                rlweq_modop.elwise_add_mut(part_b.as_mut(), share_i.share.as_ref());
 279            });
 280
 281            Self {
 282                part_b,
 283                seed: cr_seed,
 284                parameters: parameters.clone(),
 285                _phantom: PhantomData,
 286            }
 287        }
 288    }
 289}
 290
 291/// CRS seeded collective public key share
 292pub struct CommonReferenceSeededCollectivePublicKeyShare<Ro, S, P> {
 293    share: Ro,
 294    cr_seed: S,
 295    parameters: P,
 296}
 297impl<Ro, S, P> CommonReferenceSeededCollectivePublicKeyShare<Ro, S, P> {
 298    pub(super) fn new(share: Ro, cr_seed: S, parameters: P) -> Self {
 299        CommonReferenceSeededCollectivePublicKeyShare {
 300            share,
 301            cr_seed,
 302            parameters,
 303        }
 304    }
 305}
 306
 307/// CRS seeded Multi-party server key share
 308pub struct CommonReferenceSeededMultiPartyServerKeyShare<M: Matrix, P, S> {
 309    rgsw_cts: Vec<M>,
 310    /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
 311    /// to -g is at 0
 312    auto_keys: HashMap<usize, M>,
 313    lwe_ksk: M::R,
 314    /// Common reference seed
 315    cr_seed: S,
 316    parameters: P,
 317}
 318
 319impl<M: Matrix, P, S> CommonReferenceSeededMultiPartyServerKeyShare<M, P, S> {
 320    pub(super) fn new(
 321        rgsw_cts: Vec<M>,
 322        auto_keys: HashMap<usize, M>,
 323        lwe_ksk: M::R,
 324        cr_seed: S,
 325        parameters: P,
 326    ) -> Self {
 327        CommonReferenceSeededMultiPartyServerKeyShare {
 328            rgsw_cts,
 329            auto_keys,
 330            lwe_ksk,
 331            cr_seed,
 332            parameters,
 333        }
 334    }
 335
 336    pub(super) fn cr_seed(&self) -> &S {
 337        &self.cr_seed
 338    }
 339
 340    pub(super) fn parameters(&self) -> &P {
 341        &self.parameters
 342    }
 343
 344    pub(super) fn auto_keys(&self) -> &HashMap<usize, M> {
 345        &self.auto_keys
 346    }
 347
 348    pub(super) fn rgsw_cts(&self) -> &[M] {
 349        &self.rgsw_cts
 350    }
 351
 352    pub(super) fn lwe_ksk(&self) -> &M::R {
 353        &self.lwe_ksk
 354    }
 355}
 356
 357/// CRS seeded MultiParty server key
 358pub struct SeededMultiPartyServerKey<M: Matrix, S, P> {
 359    rgsw_cts: Vec<M>,
 360    /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
 361    /// to -g is at 0
 362    auto_keys: HashMap<usize, M>,
 363    lwe_ksk: M::R,
 364    cr_seed: S,
 365    parameters: P,
 366}
 367
 368impl<M: Matrix, S, P> SeededMultiPartyServerKey<M, S, P> {
 369    pub(super) fn new(
 370        rgsw_cts: Vec<M>,
 371        auto_keys: HashMap<usize, M>,
 372        lwe_ksk: M::R,
 373        cr_seed: S,
 374        parameters: P,
 375    ) -> Self {
 376        SeededMultiPartyServerKey {
 377            rgsw_cts,
 378            auto_keys,
 379            lwe_ksk,
 380            cr_seed,
 381            parameters,
 382        }
 383    }
 384
 385    pub(super) fn rgsw_cts(&self) -> &[M] {
 386        &self.rgsw_cts
 387    }
 388}
 389
 390/// Seeded single party server key
 391pub struct SeededSinglePartyServerKey<M: Matrix, P, S> {
 392    /// Rgsw cts of LWE secret elements
 393    pub(crate) rgsw_cts: Vec<M>,
 394    /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
 395    /// to -g is at 0
 396    pub(crate) auto_keys: HashMap<usize, M>,
 397    /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
 398    pub(crate) lwe_ksk: M::R,
 399    /// Parameters
 400    pub(crate) parameters: P,
 401    /// Main seed
 402    pub(crate) seed: S,
 403}
 404impl<M: Matrix, S> SeededSinglePartyServerKey<M, BoolParameters<M::MatElement>, S> {
 405    pub(super) fn from_raw(
 406        auto_keys: HashMap<usize, M>,
 407        rgsw_cts: Vec<M>,
 408        lwe_ksk: M::R,
 409        parameters: BoolParameters<M::MatElement>,
 410        seed: S,
 411    ) -> Self {
 412        // sanity checks
 413        auto_keys.iter().for_each(|v| {
 414            assert!(
 415                v.1.dimension()
 416                    == (
 417                        parameters.auto_decomposition_count().0,
 418                        parameters.rlwe_n().0
 419                    )
 420            )
 421        });
 422
 423        let (part_a_d, part_b_d) = parameters.rlwe_rgsw_decomposition_count();
 424        rgsw_cts.iter().for_each(|v| {
 425            assert!(v.dimension() == (part_a_d.0 * 2 + part_b_d.0, parameters.rlwe_n().0))
 426        });
 427        assert!(
 428            lwe_ksk.as_ref().len()
 429                == (parameters.lwe_decomposition_count().0 * parameters.rlwe_n().0)
 430        );
 431
 432        SeededSinglePartyServerKey {
 433            rgsw_cts,
 434            auto_keys,
 435            lwe_ksk,
 436            parameters,
 437            seed,
 438        }
 439    }
 440}
 441
 442/// Server key in evaluation domain
 443pub(crate) struct ServerKeyEvaluationDomain<M, P, R, N> {
 444    /// Rgsw cts of LWE secret elements
 445    rgsw_cts: Vec<M>,
 446    /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
 447    /// to -g is at 0
 448    galois_keys: HashMap<usize, M>,
 449    /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
 450    lwe_ksk: M,
 451    parameters: P,
 452    _phanton: PhantomData<(R, N)>,
 453}
 454
 455pub(super) mod impl_server_key_eval_domain {
 456    use itertools::{izip, Itertools};
 457
 458    use crate::{
 459        backend::Modulus,
 460        bool::{NonInteractiveMultiPartyCrs, SeededNonInteractiveMultiPartyServerKey},
 461        ntt::{Ntt, NttInit},
 462        pbs::PbsKey,
 463    };
 464
 465    use super::*;
 466
 467    impl<M, Mod, R, N> ServerKeyEvaluationDomain<M, Mod, R, N> {
 468        pub(in super::super) fn rgsw_cts(&self) -> &[M] {
 469            &self.rgsw_cts
 470        }
 471    }
 472
 473    impl<
 474            M: MatrixMut + MatrixEntity,
 475            R: RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>
 476                + NewWithSeed,
 477            N: NttInit<CiphertextModulus<M::MatElement>> + Ntt<Element = M::MatElement>,
 478        > From<&SeededSinglePartyServerKey<M, BoolParameters<M::MatElement>, R::Seed>>
 479        for ServerKeyEvaluationDomain<M, BoolParameters<M::MatElement>, R, N>
 480    where
 481        <M as Matrix>::R: RowMut,
 482        M::MatElement: Copy,
 483        R::Seed: Clone,
 484    {
 485        fn from(
 486            value: &SeededSinglePartyServerKey<M, BoolParameters<M::MatElement>, R::Seed>,
 487        ) -> Self {
 488            let mut main_prng = R::new_with_seed(value.seed.clone());
 489            let parameters = &value.parameters;
 490            let g = parameters.g() as isize;
 491            let ring_size = value.parameters.rlwe_n().0;
 492            let lwe_n = value.parameters.lwe_n().0;
 493            let rlwe_q = value.parameters.rlwe_q();
 494            let lwq_q = value.parameters.lwe_q();
 495
 496            let nttop = N::new(rlwe_q, ring_size);
 497
 498            // galois keys
 499            let mut auto_keys = HashMap::new();
 500            let auto_decomp_count = parameters.auto_decomposition_count().0;
 501            let auto_element_dlogs = parameters.auto_element_dlogs();
 502            for i in auto_element_dlogs.into_iter() {
 503                let seeded_auto_key = value.auto_keys.get(&i).unwrap();
 504                assert!(seeded_auto_key.dimension() == (auto_decomp_count, ring_size));
 505
 506                let mut data = M::zeros(auto_decomp_count * 2, ring_size);
 507
 508                // sample RLWE'_A(-s(X^k))
 509                data.iter_rows_mut().take(auto_decomp_count).for_each(|ri| {
 510                    RandomFillUniformInModulus::random_fill(&mut main_prng, &rlwe_q, ri.as_mut())
 511                });
 512
 513                // copy over RLWE'B_(-s(X^k))
 514                izip!(
 515                    data.iter_rows_mut().skip(auto_decomp_count),
 516                    seeded_auto_key.iter_rows()
 517                )
 518                .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
 519
 520                // Send to Evaluation domain
 521                data.iter_rows_mut()
 522                    .for_each(|ri| nttop.forward(ri.as_mut()));
 523
 524                auto_keys.insert(i, data);
 525            }
 526
 527            // RGSW ciphertexts
 528            let (rlrg_a_decomp, rlrg_b_decomp) = parameters.rlwe_rgsw_decomposition_count();
 529            let rgsw_cts = value
 530                .rgsw_cts
 531                .iter()
 532                .map(|seeded_rgsw_si| {
 533                    assert!(
 534                        seeded_rgsw_si.dimension()
 535                            == (rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0, ring_size)
 536                    );
 537
 538                    let mut data = M::zeros(rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0 * 2, ring_size);
 539
 540                    // copy over RLWE'(-sm)
 541                    izip!(
 542                        data.iter_rows_mut().take(rlrg_a_decomp.0 * 2),
 543                        seeded_rgsw_si.iter_rows().take(rlrg_a_decomp.0 * 2)
 544                    )
 545                    .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
 546
 547                    // sample RLWE'_A(m)
 548                    data.iter_rows_mut()
 549                        .skip(rlrg_a_decomp.0 * 2)
 550                        .take(rlrg_b_decomp.0)
 551                        .for_each(|ri| {
 552                            RandomFillUniformInModulus::random_fill(
 553                                &mut main_prng,
 554                                &rlwe_q,
 555                                ri.as_mut(),
 556                            )
 557                        });
 558
 559                    // copy over RLWE'_B(m)
 560                    izip!(
 561                        data.iter_rows_mut()
 562                            .skip(rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0),
 563                        seeded_rgsw_si.iter_rows().skip(rlrg_a_decomp.0 * 2)
 564                    )
 565                    .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
 566
 567                    // send polynomials to evaluation domain
 568                    data.iter_rows_mut()
 569                        .for_each(|ri| nttop.forward(ri.as_mut()));
 570
 571                    data
 572                })
 573                .collect_vec();
 574
 575            // LWE ksk
 576            let lwe_ksk = {
 577                let d = parameters.lwe_decomposition_count().0;
 578                assert!(value.lwe_ksk.as_ref().len() == d * ring_size);
 579
 580                let mut data = M::zeros(d * ring_size, lwe_n + 1);
 581                izip!(data.iter_rows_mut(), value.lwe_ksk.as_ref().iter()).for_each(
 582                    |(lwe_i, bi)| {
 583                        RandomFillUniformInModulus::random_fill(
 584                            &mut main_prng,
 585                            &lwq_q,
 586                            &mut lwe_i.as_mut()[1..],
 587                        );
 588                        lwe_i.as_mut()[0] = *bi;
 589                    },
 590                );
 591
 592                data
 593            };
 594
 595            ServerKeyEvaluationDomain {
 596                rgsw_cts,
 597                galois_keys: auto_keys,
 598                lwe_ksk,
 599                parameters: parameters.clone(),
 600                _phanton: PhantomData,
 601            }
 602        }
 603    }
 604
 605    impl<
 606            M: MatrixMut + MatrixEntity,
 607            Rng: NewWithSeed,
 608            N: NttInit<CiphertextModulus<M::MatElement>> + Ntt<Element = M::MatElement>,
 609        > From<&SeededMultiPartyServerKey<M, Rng::Seed, BoolParameters<M::MatElement>>>
 610        for ServerKeyEvaluationDomain<M, BoolParameters<M::MatElement>, Rng, N>
 611    where
 612        <M as Matrix>::R: RowMut,
 613        Rng::Seed: Copy,
 614        Rng: RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
 615        M::MatElement: Copy,
 616    {
 617        fn from(
 618            value: &SeededMultiPartyServerKey<M, Rng::Seed, BoolParameters<M::MatElement>>,
 619        ) -> Self {
 620            let g = value.parameters.g() as isize;
 621            let rlwe_n = value.parameters.rlwe_n().0;
 622            let lwe_n = value.parameters.lwe_n().0;
 623            let rlwe_q = value.parameters.rlwe_q();
 624            let lwe_q = value.parameters.lwe_q();
 625
 626            let mut main_prng = Rng::new_with_seed(value.cr_seed);
 627
 628            let rlwe_nttop = N::new(rlwe_q, rlwe_n);
 629
 630            // auto keys
 631            let mut auto_keys = HashMap::new();
 632            let auto_d_count = value.parameters.auto_decomposition_count().0;
 633            let auto_element_dlogs = value.parameters.auto_element_dlogs();
 634            for i in auto_element_dlogs.into_iter() {
 635                let mut key = M::zeros(auto_d_count * 2, rlwe_n);
 636
 637                // sample a
 638                key.iter_rows_mut().take(auto_d_count).for_each(|ri| {
 639                    RandomFillUniformInModulus::random_fill(&mut main_prng, &rlwe_q, ri.as_mut())
 640                });
 641
 642                let key_part_b = value.auto_keys.get(&i).unwrap();
 643                assert!(key_part_b.dimension() == (auto_d_count, rlwe_n));
 644                izip!(
 645                    key.iter_rows_mut().skip(auto_d_count),
 646                    key_part_b.iter_rows()
 647                )
 648                .for_each(|(to_ri, from_ri)| {
 649                    to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 650                });
 651
 652                // send to evaluation domain
 653                key.iter_rows_mut()
 654                    .for_each(|ri| rlwe_nttop.forward(ri.as_mut()));
 655
 656                auto_keys.insert(i, key);
 657            }
 658
 659            // rgsw cts
 660            let (rlrg_d_a, rlrg_d_b) = value.parameters.rlwe_rgsw_decomposition_count();
 661            let rgsw_ct_out = rlrg_d_a.0 * 2 + rlrg_d_b.0 * 2;
 662            let rgsw_cts = value
 663                .rgsw_cts
 664                .iter()
 665                .map(|ct_i_in| {
 666                    assert!(ct_i_in.dimension() == (rgsw_ct_out, rlwe_n));
 667                    let mut eval_ct_i_out = M::zeros(rgsw_ct_out, rlwe_n);
 668
 669                    izip!(eval_ct_i_out.iter_rows_mut(), ct_i_in.iter_rows()).for_each(
 670                        |(to_ri, from_ri)| {
 671                            to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 672                            rlwe_nttop.forward(to_ri.as_mut());
 673                        },
 674                    );
 675
 676                    eval_ct_i_out
 677                })
 678                .collect_vec();
 679
 680            // lwe ksk
 681            let d_lwe = value.parameters.lwe_decomposition_count().0;
 682            let mut lwe_ksk = M::zeros(rlwe_n * d_lwe, lwe_n + 1);
 683            izip!(lwe_ksk.iter_rows_mut(), value.lwe_ksk.as_ref().iter()).for_each(
 684                |(lwe_i, bi)| {
 685                    RandomFillUniformInModulus::random_fill(
 686                        &mut main_prng,
 687                        &lwe_q,
 688                        &mut lwe_i.as_mut()[1..],
 689                    );
 690                    lwe_i.as_mut()[0] = *bi;
 691                },
 692            );
 693
 694            ServerKeyEvaluationDomain {
 695                rgsw_cts,
 696                galois_keys: auto_keys,
 697                lwe_ksk,
 698                parameters: value.parameters.clone(),
 699                _phanton: PhantomData,
 700            }
 701        }
 702    }
 703
 704    impl<M: Matrix, P, R, N> PbsKey for ServerKeyEvaluationDomain<M, P, R, N> {
 705        type AutoKey = M;
 706        type LweKskKey = M;
 707        type RgswCt = M;
 708
 709        fn galois_key_for_auto(&self, k: usize) -> &Self::AutoKey {
 710            self.galois_keys.get(&k).unwrap()
 711        }
 712        fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::RgswCt {
 713            &self.rgsw_cts[si]
 714        }
 715
 716        fn lwe_ksk(&self) -> &Self::LweKskKey {
 717            &self.lwe_ksk
 718        }
 719    }
 720}
 721
 722pub(crate) struct NonInteractiveServerKeyEvaluationDomain<M, P, R, N> {
 723    /// RGSW ciphertexts ideal lwe secret key elements under ideal rlwe secret
 724    rgsw_cts: Vec<M>,
 725    /// Automorphism keys under ideal rlwe secret
 726    auto_keys: HashMap<usize, M>,
 727    /// LWE key switching key from Q -> Q_{ks}
 728    lwe_ksk: M,
 729    /// Key switching key from user j to ideal secret key s. User j's ksk is at
 730    /// j'th element
 731    ui_to_s_ksks: Vec<M>,
 732    parameters: P,
 733    _phanton: PhantomData<(R, N)>,
 734}
 735
 736pub(super) mod impl_non_interactive_server_key_eval_domain {
 737    use itertools::{izip, Itertools};
 738
 739    use crate::{bool::NonInteractiveMultiPartyCrs, random::RandomFill, Ntt, NttInit};
 740
 741    use super::*;
 742
 743    impl<M, Rng, N>
 744        From<
 745            SeededNonInteractiveMultiPartyServerKey<
 746                M,
 747                NonInteractiveMultiPartyCrs<Rng::Seed>,
 748                BoolParameters<M::MatElement>,
 749            >,
 750        > for NonInteractiveServerKeyEvaluationDomain<M, BoolParameters<M::MatElement>, Rng, N>
 751    where
 752        M: MatrixMut + MatrixEntity + Clone,
 753        Rng: NewWithSeed
 754            + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>
 755            + RandomFill<<Rng as NewWithSeed>::Seed>,
 756        N: Ntt<Element = M::MatElement> + NttInit<CiphertextModulus<M::MatElement>>,
 757        M::R: RowMut,
 758        M::MatElement: Copy,
 759        Rng::Seed: Clone + Copy + Default,
 760    {
 761        fn from(
 762            value: SeededNonInteractiveMultiPartyServerKey<
 763                M,
 764                NonInteractiveMultiPartyCrs<Rng::Seed>,
 765                BoolParameters<M::MatElement>,
 766            >,
 767        ) -> Self {
 768            let rlwe_nttop = N::new(value.parameters.rlwe_q(), value.parameters.rlwe_n().0);
 769            let ring_size = value.parameters.rlwe_n().0;
 770
 771            // RGSW cts
 772            // copy over rgsw cts and send to evaluation domain
 773            let mut rgsw_cts = value.rgsw_cts.clone();
 774            rgsw_cts.iter_mut().for_each(|c| {
 775                c.iter_rows_mut()
 776                    .for_each(|ri| rlwe_nttop.forward(ri.as_mut()))
 777            });
 778
 779            // Auto keys
 780            // populate pseudo random part of auto keys. Then send auto keys to
 781            // evaluation domain
 782            let mut auto_keys = HashMap::new();
 783            let auto_seed = value.cr_seed.auto_keys_cts_seed::<Rng>();
 784            let mut auto_prng = Rng::new_with_seed(auto_seed);
 785            let auto_element_dlogs = value.parameters.auto_element_dlogs();
 786            let d_auto = value.parameters.auto_decomposition_count().0;
 787            auto_element_dlogs.iter().for_each(|el| {
 788                let auto_part_b = value
 789                    .auto_keys
 790                    .get(el)
 791                    .expect(&format!("Auto key for element g^{el} not found"));
 792
 793                assert!(auto_part_b.dimension() == (d_auto, ring_size));
 794
 795                let mut auto_ct = M::zeros(d_auto, ring_size);
 796
 797                // sample part A
 798                auto_ct.iter_rows_mut().take(d_auto).for_each(|ri| {
 799                    RandomFillUniformInModulus::random_fill(
 800                        &mut auto_prng,
 801                        value.parameters.rlwe_q(),
 802                        ri.as_mut(),
 803                    )
 804                });
 805
 806                // Copy over part B
 807                izip!(
 808                    auto_ct.iter_rows_mut().skip(d_auto),
 809                    auto_part_b.iter_rows()
 810                )
 811                .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
 812
 813                // send to evaluation domain
 814                auto_ct
 815                    .iter_rows_mut()
 816                    .for_each(|r| rlwe_nttop.forward(r.as_mut()));
 817
 818                auto_keys.insert(*el, auto_ct);
 819            });
 820
 821            // LWE ksk
 822            // populate pseudo random part of lwe ciphertexts in ksk and copy over part b
 823            // elements
 824            let lwe_ksk_seed = value.cr_seed.lwe_ksk_cts_seed::<Rng>();
 825            let mut lwe_ksk_prng = Rng::new_with_seed(lwe_ksk_seed);
 826            let mut lwe_ksk = M::zeros(
 827                value.parameters.lwe_decomposition_count().0 * ring_size,
 828                value.parameters.lwe_n().0 + 1,
 829            );
 830            lwe_ksk.iter_rows_mut().for_each(|ri| {
 831                // first element is resereved for part b. Only sample a_is in the rest
 832                RandomFillUniformInModulus::random_fill(
 833                    &mut lwe_ksk_prng,
 834                    value.parameters.lwe_q(),
 835                    &mut ri.as_mut()[1..],
 836                )
 837            });
 838            // copy over part bs
 839            izip!(value.lwe_ksk.as_ref().iter(), lwe_ksk.iter_rows_mut()).for_each(
 840                |(b_el, lwe_ct)| {
 841                    lwe_ct.as_mut()[0] = *b_el;
 842                },
 843            );
 844
 845            // u_i to s ksk
 846            let d_uitos = value
 847                .parameters
 848                .non_interactive_ui_to_s_key_switch_decomposition_count()
 849                .0;
 850            let total_users = *value.ui_to_s_ksks_key_order.iter().max().unwrap();
 851            let ui_to_s_ksks = (0..total_users)
 852                .map(|user_index| {
 853                    let user_i_seed = value.cr_seed.ui_to_s_ks_seed_for_user_i::<Rng>(user_index);
 854                    let mut prng = Rng::new_with_seed(user_i_seed);
 855
 856                    let mut ksk_ct = M::zeros(d_uitos * 2, ring_size);
 857
 858                    ksk_ct.iter_rows_mut().take(d_uitos).for_each(|r| {
 859                        RandomFillUniformInModulus::random_fill(
 860                            &mut prng,
 861                            value.parameters.rlwe_q(),
 862                            r.as_mut(),
 863                        );
 864                    });
 865
 866                    let incoming_ksk_partb_ref =
 867                        &value.ui_to_s_ksks[value.ui_to_s_ksks_key_order[user_index]];
 868                    assert!(ksk_ct.dimension() == (d_uitos, ring_size));
 869                    izip!(
 870                        ksk_ct.iter_rows_mut().skip(d_uitos),
 871                        incoming_ksk_partb_ref.iter_rows()
 872                    )
 873                    .for_each(|(to_ri, from_ri)| {
 874                        to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 875                    });
 876
 877                    ksk_ct
 878                        .iter_rows_mut()
 879                        .for_each(|r| rlwe_nttop.forward(r.as_mut()));
 880                    ksk_ct
 881                })
 882                .collect_vec();
 883
 884            NonInteractiveServerKeyEvaluationDomain {
 885                rgsw_cts,
 886                auto_keys,
 887                lwe_ksk,
 888                ui_to_s_ksks,
 889                parameters: value.parameters.clone(),
 890                _phanton: PhantomData,
 891            }
 892        }
 893    }
 894}
 895
 896pub struct SeededNonInteractiveMultiPartyServerKey<M: Matrix, S, P> {
 897    /// u_i to s key switching keys in random order
 898    ui_to_s_ksks: Vec<M>,
 899    /// Defines order for u_i to s key switchin keys by storing the index of
 900    /// user j's ksk in `ui_to_s_ksks` at index `j`. Find user j's u_i to s ksk
 901    /// at `ui_to_s_ksks[ui_to_s_ksks_key_order[j]]`
 902    ui_to_s_ksks_key_order: Vec<usize>,
 903    /// RGSW ciphertets
 904    rgsw_cts: Vec<M>,
 905    auto_keys: HashMap<usize, M>,
 906    lwe_ksk: M::R,
 907    cr_seed: S,
 908    parameters: P,
 909}
 910
 911impl<M: Matrix, S, P> SeededNonInteractiveMultiPartyServerKey<M, S, P> {
 912    pub(super) fn new(
 913        ui_to_s_ksks: Vec<M>,
 914        ui_to_s_ksks_key_order: Vec<usize>,
 915        rgsw_cts: Vec<M>,
 916        auto_keys: HashMap<usize, M>,
 917        lwe_ksk: M::R,
 918        cr_seed: S,
 919        parameters: P,
 920    ) -> Self {
 921        Self {
 922            ui_to_s_ksks,
 923            ui_to_s_ksks_key_order,
 924            rgsw_cts,
 925            auto_keys,
 926            lwe_ksk,
 927            cr_seed,
 928            parameters,
 929        }
 930    }
 931}
 932
 933pub(crate) struct ShoupNonInteractiveServerKeyEvaluationDomain<M> {
 934    /// RGSW ciphertexts ideal lwe secret key elements under ideal rlwe secret
 935    rgsw_cts: Vec<NormalAndShoup<M>>,
 936    /// Automorphism keys under ideal rlwe secret
 937    auto_keys: HashMap<usize, NormalAndShoup<M>>,
 938    /// LWE key switching key from Q -> Q_{ks}
 939    lwe_ksk: M,
 940    /// Key switching key from user j to ideal secret key s. User j's ksk is at
 941    /// j'th element
 942    ui_to_s_ksks: Vec<NormalAndShoup<M>>,
 943}
 944
 945mod impl_shoup_non_interactive_server_key_eval_domain {
 946    use itertools::Itertools;
 947    use num_traits::{FromPrimitive, PrimInt, ToPrimitive};
 948
 949    use super::*;
 950    use crate::{backend::Modulus, pbs::PbsKey};
 951
 952    impl<M: Matrix + ToShoup<Modulus = M::MatElement>, R, N>
 953        From<NonInteractiveServerKeyEvaluationDomain<M, BoolParameters<M::MatElement>, R, N>>
 954        for ShoupNonInteractiveServerKeyEvaluationDomain<M>
 955    where
 956        M::MatElement: FromPrimitive + ToPrimitive + PrimInt,
 957    {
 958        fn from(
 959            value: NonInteractiveServerKeyEvaluationDomain<M, BoolParameters<M::MatElement>, R, N>,
 960        ) -> Self {
 961            let rlwe_q = value.parameters.rlwe_q().q().unwrap();
 962
 963            let rgsw_cts = value
 964                .rgsw_cts
 965                .into_iter()
 966                .map(|m| NormalAndShoup::new_with_modulus(m, rlwe_q))
 967                .collect_vec();
 968
 969            let mut auto_keys = HashMap::new();
 970            value.auto_keys.into_iter().for_each(|(k, v)| {
 971                auto_keys.insert(k, NormalAndShoup::new_with_modulus(v, rlwe_q));
 972            });
 973
 974            let ui_to_s_ksks = value
 975                .ui_to_s_ksks
 976                .into_iter()
 977                .map(|m| NormalAndShoup::new_with_modulus(m, rlwe_q))
 978                .collect_vec();
 979
 980            Self {
 981                rgsw_cts,
 982                auto_keys,
 983                lwe_ksk: value.lwe_ksk,
 984                ui_to_s_ksks,
 985            }
 986        }
 987    }
 988
 989    impl<M: Matrix> PbsKey for ShoupNonInteractiveServerKeyEvaluationDomain<M> {
 990        type AutoKey = NormalAndShoup<M>;
 991        type LweKskKey = M;
 992        type RgswCt = NormalAndShoup<M>;
 993
 994        fn galois_key_for_auto(&self, k: usize) -> &Self::AutoKey {
 995            self.auto_keys.get(&k).unwrap()
 996        }
 997        fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::RgswCt {
 998            &self.rgsw_cts[si]
 999        }
1000
1001        fn lwe_ksk(&self) -> &Self::LweKskKey {
1002            &self.lwe_ksk
1003        }
1004    }
1005}
1006
1007/// Server key in evaluation domain with Shoup representations
1008pub(crate) struct ShoupServerKeyEvaluationDomain<M> {
1009    /// Rgsw cts of LWE secret elements
1010    rgsw_cts: Vec<NormalAndShoup<M>>,
1011    /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
1012    /// to -g is at 0
1013    galois_keys: HashMap<usize, NormalAndShoup<M>>,
1014    /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
1015    lwe_ksk: M,
1016}
1017
1018mod shoup_server_key_eval_domain {
1019    use itertools::{izip, Itertools};
1020    use num_traits::{FromPrimitive, PrimInt};
1021
1022    use crate::{backend::Modulus, pbs::PbsKey};
1023
1024    use super::*;
1025
1026    impl<M: MatrixMut + MatrixEntity + ToShoup<Modulus = M::MatElement>, R, N>
1027        From<ServerKeyEvaluationDomain<M, BoolParameters<M::MatElement>, R, N>>
1028        for ShoupServerKeyEvaluationDomain<M>
1029    where
1030        <M as Matrix>::R: RowMut,
1031        M::MatElement: PrimInt + FromPrimitive,
1032    {
1033        fn from(value: ServerKeyEvaluationDomain<M, BoolParameters<M::MatElement>, R, N>) -> Self {
1034            let q = value.parameters.rlwe_q().q().unwrap();
1035            // Rgsw ciphertexts
1036            let rgsw_cts = value
1037                .rgsw_cts
1038                .into_iter()
1039                .map(|ct| NormalAndShoup::new_with_modulus(ct, q))
1040                .collect_vec();
1041
1042            let mut auto_keys = HashMap::new();
1043            value.galois_keys.into_iter().for_each(|(index, key)| {
1044                auto_keys.insert(index, NormalAndShoup::new_with_modulus(key, q));
1045            });
1046
1047            Self {
1048                rgsw_cts,
1049                galois_keys: auto_keys,
1050                lwe_ksk: value.lwe_ksk,
1051            }
1052        }
1053    }
1054
1055    impl<M: Matrix> PbsKey for ShoupServerKeyEvaluationDomain<M> {
1056        type AutoKey = NormalAndShoup<M>;
1057        type LweKskKey = M;
1058        type RgswCt = NormalAndShoup<M>;
1059
1060        fn galois_key_for_auto(&self, k: usize) -> &Self::AutoKey {
1061            self.galois_keys.get(&k).unwrap()
1062        }
1063        fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::RgswCt {
1064            &self.rgsw_cts[si]
1065        }
1066
1067        fn lwe_ksk(&self) -> &Self::LweKskKey {
1068            &self.lwe_ksk
1069        }
1070    }
1071}
1072
1073/// Stores normal and shoup representation of Matrix elements (Normal, Shoup)
1074pub(crate) struct NormalAndShoup<M>(M, M);
1075
1076impl<M: ToShoup> NormalAndShoup<M> {
1077    fn new_with_modulus(value: M, modulus: <M as ToShoup>::Modulus) -> Self {
1078        let value_shoup = M::to_shoup(&value, modulus);
1079        NormalAndShoup(value, value_shoup)
1080    }
1081}
1082
1083impl<M> AsRef<M> for NormalAndShoup<M> {
1084    fn as_ref(&self) -> &M {
1085        &self.0
1086    }
1087}
1088
1089impl<M> WithShoupRepr for NormalAndShoup<M> {
1090    type M = M;
1091    fn shoup_repr(&self) -> &Self::M {
1092        &self.1
1093    }
1094}