arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/rgsw.rs
src/rgsw.rsRaw
   1use std::{
   2    clone,
   3    fmt::Debug,
   4    marker::PhantomData,
   5    ops::{Div, Neg, Sub},
   6};
   7
   8use itertools::{izip, Itertools};
   9use num_traits::{PrimInt, ToPrimitive, Zero};
  10
  11use crate::{
  12    backend::{ArithmeticOps, VectorOps},
  13    decomposer::{self, Decomposer},
  14    ntt::{self, Ntt, NttInit},
  15    random::{DefaultSecureRng, NewWithSeed, RandomGaussianDist, RandomUniformDist},
  16    utils::{fill_random_ternary_secret_with_hamming_weight, TryConvertFrom, WithLocal},
  17    Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret,
  18};
  19
  20pub struct SeededAutoKey<M, S>
  21where
  22    M: Matrix,
  23{
  24    data: M,
  25    seed: S,
  26    modulus: M::MatElement,
  27}
  28
  29impl<M: Matrix + MatrixEntity, S> SeededAutoKey<M, S> {
  30    fn from_raw(data: M, seed: S, modulus: M::MatElement) -> Self {
  31        assert!(data.dimension().0 % 3 == 0);
  32
  33        SeededAutoKey {
  34            data,
  35            seed,
  36            modulus,
  37        }
  38    }
  39
  40    fn empty(ring_size: usize, d_rgsw: usize, seed: S, modulus: M::MatElement) -> Self {
  41        SeededAutoKey {
  42            data: M::zeros(d_rgsw, ring_size),
  43            seed,
  44            modulus: modulus,
  45        }
  46    }
  47}
  48
  49pub struct AutoKeyEvaluationDomain<M, R, N> {
  50    data: M,
  51    _phantom: PhantomData<(R, N)>,
  52}
  53
  54impl<
  55        M: MatrixMut + MatrixEntity,
  56        R: RandomUniformDist<[M::MatElement], Parameters = M::MatElement> + NewWithSeed,
  57        N: NttInit<Element = M::MatElement> + Ntt<Element = M::MatElement>,
  58    > From<&SeededAutoKey<M, R::Seed>> for AutoKeyEvaluationDomain<M, R, N>
  59where
  60    <M as Matrix>::R: RowMut,
  61    M::MatElement: Copy,
  62    R::Seed: Clone,
  63{
  64    fn from(value: &SeededAutoKey<M, R::Seed>) -> Self {
  65        let (d, ring_size) = value.data.dimension();
  66        let mut data = M::zeros(2 * d, ring_size);
  67
  68        // sample RLWE'_A(-s(X^k))
  69        let mut p_rng = R::new_with_seed(value.seed.clone());
  70        data.iter_rows_mut().take(d).for_each(|r| {
  71            RandomUniformDist::random_fill(&mut p_rng, &value.modulus, r.as_mut());
  72        });
  73
  74        // copy over RLWE'_B(-s(X^k))
  75        izip!(data.iter_rows_mut().skip(d), value.data.iter_rows()).for_each(|(to_r, from_r)| {
  76            to_r.as_mut().copy_from_slice(from_r.as_ref());
  77        });
  78
  79        // send RLWE'(-s(X^k)) polynomials to evaluation domain
  80        let ntt_op = N::new(value.modulus, ring_size);
  81        data.iter_rows_mut()
  82            .for_each(|r| ntt_op.forward(r.as_mut()));
  83
  84        AutoKeyEvaluationDomain {
  85            data,
  86            _phantom: PhantomData,
  87        }
  88    }
  89}
  90
  91pub struct RgswCiphertext<M: Matrix> {
  92    data: M,
  93    modulus: M::MatElement,
  94}
  95
  96pub struct SeededRgswCiphertext<M, S>
  97where
  98    M: Matrix,
  99{
 100    pub(crate) data: M,
 101    seed: S,
 102    modulus: M::MatElement,
 103}
 104
 105impl<M: Matrix + MatrixEntity, S> SeededRgswCiphertext<M, S> {
 106    fn from_raw(data: M, seed: S, modulus: M::MatElement) -> SeededRgswCiphertext<M, S> {
 107        assert!(data.dimension().0 % 3 == 0);
 108
 109        SeededRgswCiphertext {
 110            data,
 111            seed,
 112            modulus,
 113        }
 114    }
 115
 116    pub(crate) fn empty(
 117        ring_size: usize,
 118        d_rgsw: usize,
 119        seed: S,
 120        modulus: M::MatElement,
 121    ) -> SeededRgswCiphertext<M, S> {
 122        SeededRgswCiphertext {
 123            data: M::zeros(d_rgsw * 3, ring_size),
 124            seed,
 125            modulus: modulus,
 126        }
 127    }
 128}
 129
 130impl<M: Debug + Matrix, S: Debug> Debug for SeededRgswCiphertext<M, S>
 131where
 132    M::MatElement: Debug,
 133{
 134    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 135        f.debug_struct("SeededRgswCiphertext")
 136            .field("data", &self.data)
 137            .field("seed", &self.seed)
 138            .field("modulus", &self.modulus)
 139            .finish()
 140    }
 141}
 142
 143pub struct RgswCiphertextEvaluationDomain<M, R, N> {
 144    pub(crate) data: M,
 145    _phantom: PhantomData<(R, N)>,
 146}
 147
 148impl<
 149        M: MatrixMut + MatrixEntity,
 150        R: NewWithSeed + RandomUniformDist<[M::MatElement], Parameters = M::MatElement>,
 151        N: NttInit<Element = M::MatElement> + Ntt<Element = M::MatElement> + Debug,
 152    > From<&SeededRgswCiphertext<M, R::Seed>> for RgswCiphertextEvaluationDomain<M, R, N>
 153where
 154    <M as Matrix>::R: RowMut,
 155    M::MatElement: Copy,
 156    R::Seed: Clone,
 157    M: Debug,
 158{
 159    fn from(value: &SeededRgswCiphertext<M, R::Seed>) -> Self {
 160        let d = value.data.dimension().0.div(3);
 161
 162        let mut data = M::zeros(4 * d, value.data.dimension().1);
 163
 164        // copy RLWE'(-sm)
 165        izip!(data.iter_rows_mut().take(2 * d), value.data.iter_rows()).for_each(
 166            |(to_ri, from_ri)| {
 167                to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 168            },
 169        );
 170
 171        // sample A polynomials of RLWE'(m) - RLWE'A(m)
 172        // TODO(Jay): Do we want to be generic over RandomGenerator used here? I think
 173        // not.
 174        let mut p_rng = R::new_with_seed(value.seed.clone());
 175        izip!(data.iter_rows_mut().skip(2 * d).take(d))
 176            .for_each(|ri| p_rng.random_fill(&value.modulus, ri.as_mut()));
 177
 178        // RLWE'_B(m)
 179        izip!(
 180            data.iter_rows_mut().skip(3 * d),
 181            value.data.iter_rows().skip(2 * d)
 182        )
 183        .for_each(|(to_ri, from_ri)| {
 184            to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 185        });
 186
 187        // Send polynomials to evaluation domain
 188        let ring_size = data.dimension().1;
 189        let nttop = N::new(value.modulus, ring_size);
 190        data.iter_rows_mut()
 191            .for_each(|ri| nttop.forward(ri.as_mut()));
 192
 193        Self {
 194            data: data,
 195            _phantom: PhantomData,
 196        }
 197    }
 198}
 199
 200impl<
 201        M: MatrixMut + MatrixEntity,
 202        R,
 203        N: NttInit<Element = M::MatElement> + Ntt<Element = M::MatElement>,
 204    > From<&RgswCiphertext<M>> for RgswCiphertextEvaluationDomain<M, R, N>
 205where
 206    <M as Matrix>::R: RowMut,
 207    M::MatElement: Copy,
 208    M: Debug,
 209{
 210    fn from(value: &RgswCiphertext<M>) -> Self {
 211        assert!(value.data.dimension().0 % 4 == 0);
 212        let d = value.data.dimension().0.div(4);
 213
 214        let mut data = M::zeros(4 * d, value.data.dimension().1);
 215
 216        // copy RLWE'(-sm)
 217        izip!(data.iter_rows_mut().take(2 * d), value.data.iter_rows()).for_each(
 218            |(to_ri, from_ri)| {
 219                to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 220            },
 221        );
 222
 223        // copy RLWE'(m)
 224        izip!(
 225            data.iter_rows_mut().skip(2 * d),
 226            value.data.iter_rows().skip(2 * d)
 227        )
 228        .for_each(|(to_ri, from_ri)| {
 229            to_ri.as_mut().copy_from_slice(from_ri.as_ref());
 230        });
 231
 232        // Send polynomials to evaluation domain
 233        let ring_size = data.dimension().1;
 234        let nttop = N::new(value.modulus, ring_size);
 235        data.iter_rows_mut()
 236            .for_each(|ri| nttop.forward(ri.as_mut()));
 237
 238        Self {
 239            data: data,
 240            _phantom: PhantomData,
 241        }
 242    }
 243}
 244
 245impl<M: Debug, R, N> Debug for RgswCiphertextEvaluationDomain<M, R, N> {
 246    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 247        f.debug_struct("RgswCiphertextEvaluationDomain")
 248            .field("data", &self.data)
 249            .field("_phantom", &self._phantom)
 250            .finish()
 251    }
 252}
 253
 254impl<M: Matrix, R, N> Matrix for RgswCiphertextEvaluationDomain<M, R, N> {
 255    type MatElement = M::MatElement;
 256    type R = M::R;
 257
 258    fn dimension(&self) -> (usize, usize) {
 259        self.data.dimension()
 260    }
 261}
 262
 263impl<M: Matrix, R, N> AsRef<[M::R]> for RgswCiphertextEvaluationDomain<M, R, N> {
 264    fn as_ref(&self) -> &[M::R] {
 265        self.data.as_ref()
 266    }
 267}
 268
 269pub struct SeededRlweCiphertext<R, S>
 270where
 271    R: Row,
 272{
 273    pub(crate) data: R,
 274    pub(crate) seed: S,
 275    pub(crate) modulus: R::Element,
 276}
 277
 278impl<R: RowEntity, S> SeededRlweCiphertext<R, S> {
 279    pub(crate) fn empty(ring_size: usize, seed: S, modulus: R::Element) -> Self {
 280        SeededRlweCiphertext {
 281            data: R::zeros(ring_size),
 282            seed,
 283            modulus,
 284        }
 285    }
 286}
 287
 288pub struct RlweCiphertext<M, Rng> {
 289    pub(crate) data: M,
 290    pub(crate) is_trivial: bool,
 291    _phatom: PhantomData<Rng>,
 292}
 293
 294impl<M, Rng> RlweCiphertext<M, Rng> {
 295    pub(crate) fn from_raw(data: M, is_trivial: bool) -> Self {
 296        RlweCiphertext {
 297            data,
 298            is_trivial,
 299            _phatom: PhantomData,
 300        }
 301    }
 302}
 303
 304impl<M: Matrix, Rng> Matrix for RlweCiphertext<M, Rng> {
 305    type MatElement = M::MatElement;
 306    type R = M::R;
 307
 308    fn dimension(&self) -> (usize, usize) {
 309        self.data.dimension()
 310    }
 311}
 312
 313impl<M: MatrixMut, Rng> MatrixMut for RlweCiphertext<M, Rng> where <M as Matrix>::R: RowMut {}
 314
 315impl<M: Matrix, Rng> AsRef<[<M as Matrix>::R]> for RlweCiphertext<M, Rng> {
 316    fn as_ref(&self) -> &[<M as Matrix>::R] {
 317        self.data.as_ref()
 318    }
 319}
 320
 321impl<M: MatrixMut, Rng> AsMut<[<M as Matrix>::R]> for RlweCiphertext<M, Rng>
 322where
 323    <M as Matrix>::R: RowMut,
 324{
 325    fn as_mut(&mut self) -> &mut [<M as Matrix>::R] {
 326        self.data.as_mut()
 327    }
 328}
 329
 330impl<M, Rng> IsTrivial for RlweCiphertext<M, Rng> {
 331    fn is_trivial(&self) -> bool {
 332        self.is_trivial
 333    }
 334    fn set_not_trivial(&mut self) {
 335        self.is_trivial = false;
 336    }
 337}
 338
 339impl<R: Row, M: MatrixEntity<R = R, MatElement = R::Element> + MatrixMut, Rng: NewWithSeed>
 340    From<&SeededRlweCiphertext<R, Rng::Seed>> for RlweCiphertext<M, Rng>
 341where
 342    Rng::Seed: Clone,
 343    Rng: RandomUniformDist<[M::MatElement], Parameters = M::MatElement>,
 344    <M as Matrix>::R: RowMut,
 345    R::Element: Copy,
 346{
 347    fn from(value: &SeededRlweCiphertext<R, Rng::Seed>) -> Self {
 348        let mut data = M::zeros(2, value.data.as_ref().len());
 349
 350        // sample a
 351        let mut p_rng = Rng::new_with_seed(value.seed.clone());
 352        RandomUniformDist::random_fill(&mut p_rng, &value.modulus, data.get_row_mut(0));
 353
 354        data.get_row_mut(1).copy_from_slice(value.data.as_ref());
 355
 356        RlweCiphertext {
 357            data,
 358            is_trivial: false,
 359            _phatom: PhantomData,
 360        }
 361    }
 362}
 363
 364pub trait IsTrivial {
 365    fn is_trivial(&self) -> bool;
 366    fn set_not_trivial(&mut self);
 367}
 368
 369pub struct SeededRlwePublicKey<Ro: Row, S> {
 370    data: Ro,
 371    seed: S,
 372    modulus: Ro::Element,
 373}
 374
 375impl<Ro: RowEntity, S> SeededRlwePublicKey<Ro, S> {
 376    pub(crate) fn empty(ring_size: usize, seed: S, modulus: Ro::Element) -> Self {
 377        Self {
 378            data: Ro::zeros(ring_size),
 379            seed,
 380            modulus,
 381        }
 382    }
 383}
 384
 385pub struct RlwePublicKey<M, R> {
 386    data: M,
 387    _phantom: PhantomData<R>,
 388}
 389
 390impl<
 391        M: MatrixMut + MatrixEntity,
 392        Rng: NewWithSeed + RandomUniformDist<[M::MatElement], Parameters = M::MatElement>,
 393    > From<&SeededRlwePublicKey<M::R, Rng::Seed>> for RlwePublicKey<M, Rng>
 394where
 395    <M as Matrix>::R: RowMut,
 396    M::MatElement: Copy,
 397    Rng::Seed: Copy,
 398{
 399    fn from(value: &SeededRlwePublicKey<M::R, Rng::Seed>) -> Self {
 400        let mut data = M::zeros(2, value.data.as_ref().len());
 401
 402        // sample a
 403        let mut p_rng = Rng::new_with_seed(value.seed);
 404        RandomUniformDist::random_fill(&mut p_rng, &value.modulus, data.get_row_mut(0));
 405
 406        // copy over b
 407        data.get_row_mut(1).copy_from_slice(value.data.as_ref());
 408
 409        Self {
 410            data,
 411            _phantom: PhantomData,
 412        }
 413    }
 414}
 415
 416#[derive(Clone)]
 417pub struct RlweSecret {
 418    pub(crate) values: Vec<i32>,
 419}
 420
 421impl Secret for RlweSecret {
 422    type Element = i32;
 423    fn values(&self) -> &[Self::Element] {
 424        &self.values
 425    }
 426}
 427
 428impl RlweSecret {
 429    pub fn random(hw: usize, n: usize) -> RlweSecret {
 430        DefaultSecureRng::with_local_mut(|rng| {
 431            let mut out = vec![0i32; n];
 432            fill_random_ternary_secret_with_hamming_weight(&mut out, hw, rng);
 433
 434            RlweSecret { values: out }
 435        })
 436    }
 437}
 438
 439pub(crate) fn generate_auto_map(ring_size: usize, k: isize) -> (Vec<usize>, Vec<bool>) {
 440    assert!(k & 1 == 1, "Auto {k} must be odd");
 441
 442    let k = if k < 0 {
 443        // k is -ve, return k%(2*N)
 444        (2 * ring_size) - (k.abs() as usize % (2 * ring_size))
 445    } else {
 446        k as usize
 447    };
 448    let (auto_map_index, auto_sign_index): (Vec<usize>, Vec<bool>) = (0..ring_size)
 449        .into_iter()
 450        .map(|i| {
 451            let mut to_index = (i * k) % (2 * ring_size);
 452            let mut sign = true;
 453
 454            // wrap around. false implies negative
 455            if to_index >= ring_size {
 456                to_index = to_index - ring_size;
 457                sign = false;
 458            }
 459
 460            (to_index, sign)
 461        })
 462        .unzip();
 463    (auto_map_index, auto_sign_index)
 464}
 465
 466/// Generates RLWE Key switching key to key switch ciphertext RLWE_{from_s}(m)
 467/// to RLWE_{to_s}(m).
 468///
 469/// Key switching equals
 470///     \sum decompose(c_1)_i * RLWE_{to_s}(\beta^i -from_s)
 471/// Hence, key switchin key equals RLWE'(-from_s) = RLWE(-from_s), RLWE(beta^1
 472/// -from_s), ..., RLWE(beta^{d-1} -from_s).
 473///
 474/// - ksk_out: Output Key switching key. Key switching key stores only part B
 475///   polynomials of ksk RLWE ciphertexts (i.e. RLWE'_B(-from_s)) in coefficient
 476///   domain
 477/// - neg_from_s: Negative of secret polynomial to key switch from
 478/// - to_s: secret polynomial to key switch to.
 479pub(crate) fn rlwe_ksk_gen<
 480    Mmut: MatrixMut + MatrixEntity,
 481    ModOp: ArithmeticOps<Element = Mmut::MatElement> + VectorOps<Element = Mmut::MatElement>,
 482    NttOp: Ntt<Element = Mmut::MatElement>,
 483    R: RandomGaussianDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
 484    PR: RandomUniformDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
 485>(
 486    ksk_out: &mut Mmut,
 487    neg_from_s: Mmut::R,
 488    mut to_s: Mmut::R,
 489    gadget_vector: &[Mmut::MatElement],
 490    mod_op: &ModOp,
 491    ntt_op: &NttOp,
 492    p_rng: &mut PR,
 493    rng: &mut R,
 494) where
 495    <Mmut as Matrix>::R: RowMut,
 496{
 497    let ring_size = neg_from_s.as_ref().len();
 498    let d = gadget_vector.len();
 499    assert!(ksk_out.dimension() == (d, ring_size));
 500
 501    let q = ArithmeticOps::modulus(mod_op);
 502
 503    ntt_op.forward(to_s.as_mut());
 504
 505    // RLWE'_{to_s}(-from_s)
 506    let mut part_a = {
 507        let mut a = Mmut::zeros(d, ring_size);
 508        a.iter_rows_mut()
 509            .for_each(|ai| RandomUniformDist::random_fill(p_rng, &q, ai.as_mut()));
 510        a
 511    };
 512    izip!(
 513        part_a.iter_rows_mut(),
 514        ksk_out.iter_rows_mut(),
 515        gadget_vector.iter(),
 516    )
 517    .for_each(|(ai, bi, beta_i)| {
 518        // si * ai
 519        ntt_op.forward(ai.as_mut());
 520        mod_op.elwise_mul_mut(ai.as_mut(), to_s.as_ref());
 521        ntt_op.backward(ai.as_mut());
 522
 523        // ei + to_s*ai
 524        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
 525        mod_op.elwise_add_mut(bi.as_mut(), ai.as_ref());
 526
 527        // beta_i * -from_s
 528        // use ai as scratch space
 529        mod_op.elwise_scalar_mul(ai.as_mut(), neg_from_s.as_ref(), beta_i);
 530
 531        // bi = ei + to_s*ai + beta_i*-from_s
 532        mod_op.elwise_add_mut(bi.as_mut(), ai.as_ref());
 533    });
 534}
 535
 536pub(crate) fn galois_key_gen<
 537    Mmut: MatrixMut + MatrixEntity,
 538    ModOp: ArithmeticOps<Element = Mmut::MatElement> + VectorOps<Element = Mmut::MatElement>,
 539    NttOp: Ntt<Element = Mmut::MatElement>,
 540    S,
 541    R: RandomGaussianDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
 542    PR: RandomUniformDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
 543>(
 544    ksk_out: &mut Mmut,
 545    s: &[S],
 546    auto_k: isize,
 547    gadget_vector: &[Mmut::MatElement],
 548    mod_op: &ModOp,
 549    ntt_op: &NttOp,
 550    p_rng: &mut PR,
 551    rng: &mut R,
 552) where
 553    <Mmut as Matrix>::R: RowMut,
 554    Mmut::R: TryConvertFrom<[S], Parameters = Mmut::MatElement> + RowEntity,
 555    Mmut::MatElement: Copy + Sub<Output = Mmut::MatElement>,
 556{
 557    let ring_size = s.len();
 558    let (auto_map_index, auto_map_sign) = generate_auto_map(ring_size, auto_k);
 559
 560    let q = ArithmeticOps::modulus(mod_op);
 561
 562    // s(X) -> -s(X^k)
 563    let s = Mmut::R::try_convert_from(s, &q);
 564    let mut neg_s_auto = Mmut::R::zeros(s.as_ref().len());
 565    izip!(s.as_ref(), auto_map_index.iter(), auto_map_sign.iter()).for_each(
 566        |(el, to_index, sign)| {
 567            // if sign is +ve (true), then negate because we need -s(X) (i.e. do the
 568            // opposite than the usual case)
 569            if *sign {
 570                neg_s_auto.as_mut()[*to_index] = q - *el;
 571            } else {
 572                neg_s_auto.as_mut()[*to_index] = *el;
 573            }
 574        },
 575    );
 576
 577    // Ksk from -s(X^k) to s(X)
 578    rlwe_ksk_gen(
 579        ksk_out,
 580        neg_s_auto,
 581        s,
 582        gadget_vector,
 583        mod_op,
 584        ntt_op,
 585        p_rng,
 586        rng,
 587    );
 588}
 589
 590pub(crate) fn routine<R: RowMut, ModOp: VectorOps<Element = R::Element>>(
 591    write_to_row: &mut [R::Element],
 592    matrix_a: &[R],
 593    matrix_b: &[R],
 594    mod_op: &ModOp,
 595) {
 596    izip!(matrix_a.iter(), matrix_b.iter()).for_each(|(a, b)| {
 597        mod_op.elwise_fma_mut(write_to_row, a.as_ref(), b.as_ref());
 598    });
 599}
 600
 601/// Decomposes ring polynomial r(X) into d polynomials using decomposer into
 602/// output matrix decomp_r
 603///
 604/// Note that decomposition of r(X) requires decomposition of each of
 605/// coefficients.
 606///
 607/// - decomp_r: must have dimensions d x ring_size. i^th decomposed polynomial
 608///   will be stored at i^th row.
 609pub(crate) fn decompose_r<R: RowMut, D: Decomposer<Element = R::Element>>(
 610    r: &[R::Element],
 611    decomp_r: &mut [R],
 612    decomposer: &D,
 613) where
 614    R::Element: Copy,
 615{
 616    let ring_size = r.len();
 617    let d = decomposer.d();
 618
 619    for ri in 0..ring_size {
 620        let el_decomposed = decomposer.decompose(&r[ri]);
 621        for j in 0..d {
 622            decomp_r[j].as_mut()[ri] = el_decomposed[j];
 623        }
 624    }
 625}
 626
 627/// Sends RLWE_{s}(X) -> RLWE_{s}(X^k) where k is some galois element
 628pub(crate) fn galois_auto<
 629    MT: Matrix + IsTrivial + MatrixMut,
 630    Mmut: MatrixMut<MatElement = MT::MatElement>,
 631    ModOp: ArithmeticOps<Element = MT::MatElement> + VectorOps<Element = MT::MatElement>,
 632    NttOp: Ntt<Element = MT::MatElement>,
 633    D: Decomposer<Element = MT::MatElement>,
 634>(
 635    rlwe_in: &mut MT,
 636    ksk: &Mmut,
 637    scratch_matrix_dplus2_ring: &mut Mmut,
 638    auto_map_index: &[usize],
 639    auto_map_sign: &[bool],
 640    mod_op: &ModOp,
 641    ntt_op: &NttOp,
 642    decomposer: &D,
 643) where
 644    <Mmut as Matrix>::R: RowMut,
 645    <MT as Matrix>::R: RowMut,
 646    MT::MatElement: Copy + Zero,
 647{
 648    let d = decomposer.d();
 649
 650    let (scratch_matrix_d_ring, tmp_rlwe_out) = scratch_matrix_dplus2_ring.split_at_row_mut(d);
 651
 652    // send b(X) -> b(X^k)
 653    izip!(
 654        rlwe_in.get_row(1),
 655        auto_map_index.iter(),
 656        auto_map_sign.iter()
 657    )
 658    .for_each(|(el_in, to_index, sign)| {
 659        if !*sign {
 660            tmp_rlwe_out[1].as_mut()[*to_index] = mod_op.neg(el_in);
 661        } else {
 662            tmp_rlwe_out[1].as_mut()[*to_index] = *el_in;
 663            // scratch_matrix_dplus2_ring.set(d + 1, *to_index, *el_in);
 664        }
 665    });
 666
 667    if !rlwe_in.is_trivial() {
 668        // send a(X) -> a(X^k) and decompose a(X^k)
 669        izip!(
 670            rlwe_in.get_row(0),
 671            auto_map_index.iter(),
 672            auto_map_sign.iter()
 673        )
 674        .for_each(|(el_in, to_index, sign)| {
 675            let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
 676
 677            let el_out_decomposed = decomposer.decompose(&el_out);
 678            for j in 0..d {
 679                scratch_matrix_d_ring[j].as_mut()[*to_index] = el_out_decomposed[j];
 680            }
 681        });
 682
 683        // transform decomposed a(X^k) to evaluation domain
 684        scratch_matrix_d_ring.iter_mut().for_each(|r| {
 685            ntt_op.forward(r.as_mut());
 686        });
 687
 688        // RLWE(m^k) = a', b'; RLWE(m) = a, b
 689        // key switch: (a * RLWE'(s(X^k)))
 690        let (ksk_a, ksk_b) = ksk.split_at_row(d);
 691        tmp_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
 692        // a' = decomp<a> * RLWE'_A(s(X^k))
 693        routine(
 694            tmp_rlwe_out[0].as_mut(),
 695            scratch_matrix_d_ring,
 696            ksk_a,
 697            mod_op,
 698        );
 699        // send b(X^k) to evaluation domain
 700        ntt_op.forward(tmp_rlwe_out[1].as_mut());
 701        // b' = b(X^k)
 702        // b' += decomp<a(X^k)> * RLWE'_B(s(X^k))
 703        routine(
 704            tmp_rlwe_out[1].as_mut(),
 705            scratch_matrix_d_ring,
 706            ksk_b,
 707            mod_op,
 708        );
 709
 710        // transform RLWE(m^k) to coefficient domain
 711        tmp_rlwe_out
 712            .iter_mut()
 713            .for_each(|r| ntt_op.backward(r.as_mut()));
 714
 715        rlwe_in
 716            .get_row_mut(0)
 717            .copy_from_slice(tmp_rlwe_out[0].as_ref());
 718    }
 719
 720    rlwe_in
 721        .get_row_mut(1)
 722        .copy_from_slice(tmp_rlwe_out[1].as_ref());
 723}
 724
 725/// Returns RLWE(m0m1) = RLWE(m0) x RGSW(m1). Mutates rlwe_in inplace to equal
 726/// RLWE(m0m1)
 727///
 728/// - rlwe_in: is RLWE(m0) with polynomials in coefficient domain
 729/// - rgsw_in: is RGSW(m1) with polynomials in evaluation domain
 730/// - scratch_matrix_d_ring: is a matrix of dimension (d_rgsw, ring_size) used
 731///   as scratch space to store decomposed Ring elements temporarily
 732pub(crate) fn rlwe_by_rgsw<
 733    Mmut: MatrixMut,
 734    MT: Matrix<MatElement = Mmut::MatElement> + MatrixMut<MatElement = Mmut::MatElement> + IsTrivial,
 735    D: Decomposer<Element = Mmut::MatElement>,
 736    ModOp: VectorOps<Element = Mmut::MatElement>,
 737    NttOp: Ntt<Element = Mmut::MatElement>,
 738>(
 739    rlwe_in: &mut MT,
 740    rgsw_in: &Mmut,
 741    scratch_matrix_dplus2_ring: &mut Mmut,
 742    decomposer: &D,
 743    ntt_op: &NttOp,
 744    mod_op: &ModOp,
 745) where
 746    Mmut::MatElement: Copy + Zero,
 747    <Mmut as Matrix>::R: RowMut,
 748    <MT as Matrix>::R: RowMut,
 749{
 750    let d_rgsw = decomposer.d();
 751    assert!(scratch_matrix_dplus2_ring.dimension() == (d_rgsw + 2, rlwe_in.dimension().1));
 752    assert!(rgsw_in.dimension() == (d_rgsw * 4, rlwe_in.dimension().1));
 753
 754    // decomposed RLWE x RGSW
 755    let (rlwe_dash_nsm, rlwe_dash_m) = rgsw_in.split_at_row(d_rgsw * 2);
 756    let (scratch_matrix_d_ring, scratch_rlwe_out) =
 757        scratch_matrix_dplus2_ring.split_at_row_mut(d_rgsw);
 758    scratch_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
 759    scratch_rlwe_out[1].as_mut().fill(Mmut::MatElement::zero());
 760    // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
 761    if !rlwe_in.is_trivial() {
 762        // a_in = 0 when RLWE_in is trivial RLWE ciphertext
 763        // decomp<a_in>
 764        decompose_r(rlwe_in.get_row_slice(0), scratch_matrix_d_ring, decomposer);
 765        scratch_matrix_d_ring
 766            .iter_mut()
 767            .for_each(|r| ntt_op.forward(r.as_mut()));
 768        // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
 769        routine(
 770            scratch_rlwe_out[0].as_mut(),
 771            scratch_matrix_d_ring.as_ref(),
 772            &rlwe_dash_nsm[..d_rgsw],
 773            mod_op,
 774        );
 775        // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
 776        routine(
 777            scratch_rlwe_out[1].as_mut(),
 778            scratch_matrix_d_ring.as_ref(),
 779            &rlwe_dash_nsm[d_rgsw..],
 780            mod_op,
 781        );
 782    }
 783    // decomp<b_in>
 784    decompose_r(rlwe_in.get_row_slice(1), scratch_matrix_d_ring, decomposer);
 785    scratch_matrix_d_ring
 786        .iter_mut()
 787        .for_each(|r| ntt_op.forward(r.as_mut()));
 788    // a_out += decomp<b_in> \cdot RLWE_A'(m)
 789    routine(
 790        scratch_rlwe_out[0].as_mut(),
 791        scratch_matrix_d_ring.as_ref(),
 792        &rlwe_dash_m[..d_rgsw],
 793        mod_op,
 794    );
 795    // b_out += decomp<b_in> \cdot RLWE_B'(m)
 796    routine(
 797        scratch_rlwe_out[1].as_mut(),
 798        scratch_matrix_d_ring.as_ref(),
 799        &rlwe_dash_m[d_rgsw..],
 800        mod_op,
 801    );
 802
 803    // transform rlwe_out to coefficient domain
 804    scratch_rlwe_out
 805        .iter_mut()
 806        .for_each(|r| ntt_op.backward(r.as_mut()));
 807
 808    rlwe_in
 809        .get_row_mut(0)
 810        .copy_from_slice(scratch_rlwe_out[0].as_mut());
 811    rlwe_in
 812        .get_row_mut(1)
 813        .copy_from_slice(scratch_rlwe_out[1].as_mut());
 814    rlwe_in.set_not_trivial();
 815}
 816
 817/// Inplace mutates rlwe_0_eval_domain to equal RGSW(m0m1) = RGSW(m0)xRGSW(m1)
 818/// in evaluation domain
 819///
 820/// - rgsw_0_eval_domain: RGSW(m0) in evaluation domain
 821/// - rgsw_1: RGSW(m1)
 822/// - scratch_matrix_d_plus_rgsw_by_ring: scratch space matrix of size
 823///   (d+(d*4))xring_size, where d equals d_rgsw
 824pub(crate) fn rgsw_by_rgsw_inplace<
 825    Mmut: MatrixMut,
 826    D: Decomposer<Element = Mmut::MatElement>,
 827    ModOp: VectorOps<Element = Mmut::MatElement>,
 828    NttOp: Ntt<Element = Mmut::MatElement>,
 829>(
 830    rgsw_0_eval_domain: &mut Mmut,
 831    rgsw_1: &Mmut,
 832    decomposer: &D,
 833    scratch_matrix_d_plus_rgsw_by_ring: &mut Mmut,
 834    ntt_op: &NttOp,
 835    mod_op: &ModOp,
 836) where
 837    <Mmut as Matrix>::R: RowMut,
 838    Mmut::MatElement: Copy + Zero,
 839{
 840    let d_rgsw = decomposer.d();
 841    assert!(rgsw_0_eval_domain.dimension().0 == 4 * d_rgsw);
 842    let ring_size = rgsw_0_eval_domain.dimension().1;
 843    assert!(rgsw_1.dimension() == (4 * d_rgsw, ring_size));
 844    assert!(scratch_matrix_d_plus_rgsw_by_ring.dimension() == (d_rgsw + (d_rgsw * 4), ring_size));
 845
 846    let (decomp_r_space, rgsw_space) = scratch_matrix_d_plus_rgsw_by_ring.split_at_row_mut(d_rgsw);
 847
 848    // zero rgsw_space
 849    rgsw_space
 850        .iter_mut()
 851        .for_each(|ri| ri.as_mut().fill(Mmut::MatElement::zero()));
 852    let (rlwe_dash_space_nsm, rlwe_dash_space_m) = rgsw_space.split_at_mut(d_rgsw * 2);
 853    let (rlwe_dash_space_nsm_parta, rlwe_dash_space_nsm_partb) =
 854        rlwe_dash_space_nsm.split_at_mut(d_rgsw);
 855    let (rlwe_dash_space_m_parta, rlwe_dash_space_m_partb) = rlwe_dash_space_m.split_at_mut(d_rgsw);
 856
 857    let (rgsw0_nsm, rgsw0_m) = rgsw_0_eval_domain.split_at_row(d_rgsw * 2);
 858    let (rgsw1_nsm, rgsw1_m) = rgsw_1.split_at_row(d_rgsw * 2);
 859
 860    // RGSW x RGSW
 861    izip!(
 862        rgsw1_nsm
 863            .iter()
 864            .take(d_rgsw)
 865            .chain(rgsw1_m.iter().take(d_rgsw)),
 866        rgsw1_nsm
 867            .iter()
 868            .skip(d_rgsw)
 869            .chain(rgsw1_m.iter().skip(d_rgsw)),
 870        rlwe_dash_space_nsm_parta
 871            .iter_mut()
 872            .chain(rlwe_dash_space_m_parta.iter_mut()),
 873        rlwe_dash_space_nsm_partb
 874            .iter_mut()
 875            .chain(rlwe_dash_space_m_partb.iter_mut()),
 876    )
 877    .for_each(|(rlwe_a, rlwe_b, rlwe_out_a, rlwe_out_b)| {
 878        // Part A
 879        decompose_r(rlwe_a.as_ref(), decomp_r_space.as_mut(), decomposer);
 880        decomp_r_space
 881            .iter_mut()
 882            .for_each(|ri| ntt_op.forward(ri.as_mut()));
 883        routine(
 884            rlwe_out_a.as_mut(),
 885            decomp_r_space,
 886            &rgsw0_nsm[..d_rgsw],
 887            mod_op,
 888        );
 889        routine(
 890            rlwe_out_b.as_mut(),
 891            decomp_r_space,
 892            &rgsw0_nsm[d_rgsw..],
 893            mod_op,
 894        );
 895
 896        // Part B
 897        decompose_r(rlwe_b.as_ref(), decomp_r_space.as_mut(), decomposer);
 898        decomp_r_space
 899            .iter_mut()
 900            .for_each(|ri| ntt_op.forward(ri.as_mut()));
 901        routine(
 902            rlwe_out_a.as_mut(),
 903            decomp_r_space,
 904            &rgsw0_m[..d_rgsw],
 905            mod_op,
 906        );
 907        routine(
 908            rlwe_out_b.as_mut(),
 909            decomp_r_space,
 910            &rgsw0_m[d_rgsw..],
 911            mod_op,
 912        );
 913    });
 914
 915    // copy over RGSW(m0m1) into RGSW(m0)
 916    izip!(rgsw_0_eval_domain.iter_rows_mut(), rgsw_space.iter())
 917        .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()))
 918}
 919
 920/// Encrypts message m as a RGSW ciphertext.
 921///
 922/// - m_eval: is `m` is evaluation domain
 923/// - out_rgsw: RGSW(m) is stored as single matrix of dimension (d_rgsw * 3,
 924///   ring_size). The matrix has the following structure [RLWE'_A(-sm) ||
 925///   RLWE'_B(-sm) || RLWE'_B(m)]^T and RLWE'_A(m) is generated via seed (where
 926///   p_rng is assumed to be seeded with seed)
 927pub(crate) fn secret_key_encrypt_rgsw<
 928    Mmut: MatrixMut + MatrixEntity,
 929    S,
 930    R: RandomGaussianDist<[Mmut::MatElement], Parameters = Mmut::MatElement>
 931        + RandomUniformDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
 932    PR: RandomUniformDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
 933    ModOp: VectorOps<Element = Mmut::MatElement>,
 934    NttOp: Ntt<Element = Mmut::MatElement>,
 935>(
 936    out_rgsw: &mut Mmut,
 937    m: &[Mmut::MatElement],
 938    gadget_vector: &[Mmut::MatElement],
 939    s: &[S],
 940    mod_op: &ModOp,
 941    ntt_op: &NttOp,
 942    p_rng: &mut PR,
 943    rng: &mut R,
 944) where
 945    <Mmut as Matrix>::R: RowMut + RowEntity + TryConvertFrom<[S], Parameters = Mmut::MatElement>,
 946    Mmut::MatElement: Copy,
 947{
 948    let d = gadget_vector.len();
 949    let q = mod_op.modulus();
 950    let ring_size = s.len();
 951    assert!(out_rgsw.dimension() == (d * 3, ring_size));
 952    assert!(m.as_ref().len() == ring_size);
 953
 954    // RLWE(-sm), RLWE(m)
 955    let (rlwe_dash_nsm, b_rlwe_dash_m) = out_rgsw.split_at_row_mut(d * 2);
 956
 957    let mut s_eval = Mmut::R::try_convert_from(s, &q);
 958    ntt_op.forward(s_eval.as_mut());
 959
 960    let mut scratch_space = Mmut::R::zeros(ring_size);
 961
 962    // RLWE'(-sm)
 963    let (a_rlwe_dash_nsm, b_rlwe_dash_nsm) = rlwe_dash_nsm.split_at_mut(d);
 964    izip!(
 965        a_rlwe_dash_nsm.iter_mut(),
 966        b_rlwe_dash_nsm.iter_mut(),
 967        gadget_vector.iter()
 968    )
 969    .for_each(|(ai, bi, beta_i)| {
 970        // Sample a_i
 971        RandomUniformDist::random_fill(rng, &q, ai.as_mut());
 972
 973        // a_i * s
 974        scratch_space.as_mut().copy_from_slice(ai.as_ref());
 975        ntt_op.forward(scratch_space.as_mut());
 976        mod_op.elwise_mul_mut(scratch_space.as_mut(), s_eval.as_ref());
 977        ntt_op.backward(scratch_space.as_mut());
 978
 979        // b_i = e_i + a_i * s
 980        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
 981        mod_op.elwise_add_mut(bi.as_mut(), scratch_space.as_ref());
 982
 983        // a_i + \beta_i * m
 984        mod_op.elwise_scalar_mul(scratch_space.as_mut(), m.as_ref(), beta_i);
 985        mod_op.elwise_add_mut(ai.as_mut(), scratch_space.as_ref());
 986    });
 987
 988    // RLWE(m)
 989    let mut a_rlwe_dash_m = {
 990        // polynomials of part A of RLWE'(m) are sampled from seed
 991        let mut a = Mmut::zeros(d, ring_size);
 992        a.iter_rows_mut()
 993            .for_each(|ai| RandomUniformDist::random_fill(p_rng, &q, ai.as_mut()));
 994        a
 995    };
 996
 997    izip!(
 998        a_rlwe_dash_m.iter_rows_mut(),
 999        b_rlwe_dash_m.iter_mut(),
1000        gadget_vector.iter()
1001    )
1002    .for_each(|(ai, bi, beta_i)| {
1003        // ai * s
1004        ntt_op.forward(ai.as_mut());
1005        mod_op.elwise_mul_mut(ai.as_mut(), s_eval.as_ref());
1006        ntt_op.backward(ai.as_mut());
1007
1008        // beta_i * m
1009        mod_op.elwise_scalar_mul(scratch_space.as_mut(), m.as_ref(), beta_i);
1010
1011        // Sample e_i
1012        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
1013        // e_i + beta_i * m + ai*s
1014        mod_op.elwise_add_mut(bi.as_mut(), scratch_space.as_ref());
1015        mod_op.elwise_add_mut(bi.as_mut(), ai.as_ref());
1016    });
1017}
1018
1019pub(crate) fn public_key_encrypt_rgsw<
1020    Mmut: MatrixMut + MatrixEntity,
1021    M: Matrix<MatElement = Mmut::MatElement>,
1022    R: RandomGaussianDist<[Mmut::MatElement], Parameters = Mmut::MatElement>
1023        + RandomUniformDist<[u8], Parameters = u8>
1024        + RandomUniformDist<usize, Parameters = usize>,
1025    ModOp: VectorOps<Element = Mmut::MatElement>,
1026    NttOp: Ntt<Element = Mmut::MatElement>,
1027>(
1028    out_rgsw: &mut Mmut,
1029    m: &[M::MatElement],
1030    public_key: &M,
1031    gadget_vector: &[Mmut::MatElement],
1032    mod_op: &ModOp,
1033    ntt_op: &NttOp,
1034    rng: &mut R,
1035) where
1036    <Mmut as Matrix>::R: RowMut + RowEntity + TryConvertFrom<[i32], Parameters = Mmut::MatElement>,
1037    Mmut::MatElement: Copy,
1038{
1039    let ring_size = public_key.dimension().1;
1040    let d = gadget_vector.len();
1041    assert!(public_key.dimension().0 == 2);
1042    assert!(out_rgsw.dimension() == (d * 4, ring_size));
1043
1044    let mut pk_eval = Mmut::zeros(2, ring_size);
1045    izip!(pk_eval.iter_rows_mut(), public_key.iter_rows()).for_each(|(to_i, from_i)| {
1046        to_i.as_mut().copy_from_slice(from_i.as_ref());
1047        ntt_op.forward(to_i.as_mut());
1048    });
1049    let p0 = pk_eval.get_row_slice(0);
1050    let p1 = pk_eval.get_row_slice(1);
1051
1052    let q = mod_op.modulus();
1053
1054    // RGSW(m) = RLWE'(-sm), RLWE(m)
1055    let (rlwe_dash_nsm, rlwe_dash_m) = out_rgsw.split_at_row_mut(2 * d);
1056
1057    // RLWE(-sm)
1058    let (rlwe_dash_nsm_parta, rlwe_dash_nsm_partb) = rlwe_dash_nsm.split_at_mut(d);
1059    izip!(
1060        rlwe_dash_nsm_parta.iter_mut(),
1061        rlwe_dash_nsm_partb.iter_mut(),
1062        gadget_vector.iter()
1063    )
1064    .for_each(|(ai, bi, beta_i)| {
1065        // sample ephemeral secret u_i
1066        let mut u = vec![0i32; ring_size];
1067        fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
1068        let mut u_eval = Mmut::R::try_convert_from(u.as_ref(), &q);
1069        ntt_op.forward(u_eval.as_mut());
1070
1071        let mut u_eval_copy = Mmut::R::zeros(ring_size);
1072        u_eval_copy.as_mut().copy_from_slice(u_eval.as_ref());
1073
1074        // p0 * u
1075        mod_op.elwise_mul_mut(u_eval.as_mut(), p0.as_ref());
1076        // p1 * u
1077        mod_op.elwise_mul_mut(u_eval_copy.as_mut(), p1.as_ref());
1078        ntt_op.backward(u_eval.as_mut());
1079        ntt_op.backward(u_eval_copy.as_mut());
1080
1081        // sample error
1082        RandomGaussianDist::random_fill(rng, &q, ai.as_mut());
1083        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
1084
1085        // a = p0*u+e0
1086        mod_op.elwise_add_mut(ai.as_mut(), u_eval.as_ref());
1087        // b = p1*u+e1
1088        mod_op.elwise_add_mut(bi.as_mut(), u_eval_copy.as_ref());
1089
1090        // a = p0*u + e0 + \beta*m
1091        // use u_eval as scratch
1092        mod_op.elwise_scalar_mul(u_eval.as_mut(), m.as_ref(), beta_i);
1093        mod_op.elwise_add_mut(ai.as_mut(), u_eval.as_ref());
1094    });
1095
1096    // RLWE(m)
1097    let (rlwe_dash_m_parta, rlwe_dash_m_partb) = rlwe_dash_m.split_at_mut(d);
1098    izip!(
1099        rlwe_dash_m_parta.iter_mut(),
1100        rlwe_dash_m_partb.iter_mut(),
1101        gadget_vector.iter()
1102    )
1103    .for_each(|(ai, bi, beta_i)| {
1104        // sample ephemeral secret u_i
1105        let mut u = vec![0i32; ring_size];
1106        fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
1107        let mut u_eval = Mmut::R::try_convert_from(u.as_ref(), &q);
1108        ntt_op.forward(u_eval.as_mut());
1109
1110        let mut u_eval_copy = Mmut::R::zeros(ring_size);
1111        u_eval_copy.as_mut().copy_from_slice(u_eval.as_ref());
1112
1113        // p0 * u
1114        mod_op.elwise_mul_mut(u_eval.as_mut(), p0.as_ref());
1115        // p1 * u
1116        mod_op.elwise_mul_mut(u_eval_copy.as_mut(), p1.as_ref());
1117        ntt_op.backward(u_eval.as_mut());
1118        ntt_op.backward(u_eval_copy.as_mut());
1119
1120        // sample error
1121        RandomGaussianDist::random_fill(rng, &q, ai.as_mut());
1122        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
1123
1124        // a = p0*u+e0
1125        mod_op.elwise_add_mut(ai.as_mut(), u_eval.as_ref());
1126        // b = p1*u+e1
1127        mod_op.elwise_add_mut(bi.as_mut(), u_eval_copy.as_ref());
1128
1129        // b = p1*u + e0 + \beta*m
1130        // use u_eval as scratch
1131        mod_op.elwise_scalar_mul(u_eval.as_mut(), m.as_ref(), beta_i);
1132        mod_op.elwise_add_mut(bi.as_mut(), u_eval.as_ref());
1133    });
1134}
1135
1136/// Encrypt polynomial m(X) as RLWE ciphertext.
1137///
1138/// - rlwe_out: returned RLWE ciphertext RLWE(m) in coefficient domain. RLWE
1139///   ciphertext is a matirx with first row consiting polynomial `a` and the
1140///   second rows consting polynomial `b`
1141pub(crate) fn secret_key_encrypt_rlwe<
1142    Ro: Row + RowMut + RowEntity,
1143    ModOp: VectorOps<Element = Ro::Element>,
1144    NttOp: Ntt<Element = Ro::Element>,
1145    S,
1146    R: RandomGaussianDist<[Ro::Element], Parameters = Ro::Element>,
1147    PR: RandomUniformDist<[Ro::Element], Parameters = Ro::Element>,
1148>(
1149    m: &Ro,
1150    b_rlwe_out: &mut Ro,
1151    s: &[S],
1152    mod_op: &ModOp,
1153    ntt_op: &NttOp,
1154    p_rng: &mut PR,
1155    rng: &mut R,
1156) where
1157    Ro: TryConvertFrom<[S], Parameters = Ro::Element> + Debug,
1158{
1159    let ring_size = s.len();
1160    assert!(m.as_ref().len() == ring_size);
1161    assert!(b_rlwe_out.as_ref().len() == ring_size);
1162
1163    let q = mod_op.modulus();
1164
1165    // sample a
1166    let mut a = {
1167        let mut a = Ro::zeros(ring_size);
1168        RandomUniformDist::random_fill(p_rng, &q, a.as_mut());
1169        a
1170    };
1171
1172    // s * a
1173    let mut sa = Ro::try_convert_from(s, &q);
1174    ntt_op.forward(sa.as_mut());
1175    ntt_op.forward(a.as_mut());
1176    mod_op.elwise_mul_mut(sa.as_mut(), a.as_ref());
1177    ntt_op.backward(sa.as_mut());
1178
1179    // sample e
1180    RandomGaussianDist::random_fill(rng, &q, b_rlwe_out.as_mut());
1181    mod_op.elwise_add_mut(b_rlwe_out.as_mut(), m.as_ref());
1182    mod_op.elwise_add_mut(b_rlwe_out.as_mut(), sa.as_ref());
1183}
1184
1185/// Generates RLWE public key
1186pub(crate) fn gen_rlwe_public_key<
1187    Ro: RowMut + RowEntity,
1188    S,
1189    ModOp: VectorOps<Element = Ro::Element>,
1190    NttOp: Ntt<Element = Ro::Element>,
1191    PRng: RandomUniformDist<[Ro::Element], Parameters = Ro::Element>,
1192    Rng: RandomGaussianDist<[Ro::Element], Parameters = Ro::Element>,
1193>(
1194    part_b_out: &mut Ro,
1195    s: &[S],
1196    ntt_op: &NttOp,
1197    mod_op: &ModOp,
1198    p_rng: &mut PRng,
1199    rng: &mut Rng,
1200) where
1201    Ro: TryConvertFrom<[S], Parameters = Ro::Element>,
1202{
1203    let ring_size = s.len();
1204    assert!(part_b_out.as_ref().len() == ring_size);
1205
1206    let q = mod_op.modulus();
1207
1208    // sample a
1209    let mut a = {
1210        let mut tmp = Ro::zeros(ring_size);
1211        RandomUniformDist::random_fill(p_rng, &q, tmp.as_mut());
1212        tmp
1213    };
1214    ntt_op.forward(a.as_mut());
1215
1216    // s*a
1217    let mut sa = Ro::try_convert_from(s, &q);
1218    ntt_op.forward(sa.as_mut());
1219    mod_op.elwise_mul_mut(sa.as_mut(), a.as_ref());
1220    ntt_op.backward(sa.as_mut());
1221
1222    // s*a + e
1223    RandomGaussianDist::random_fill(rng, &q, part_b_out.as_mut());
1224    mod_op.elwise_add_mut(part_b_out.as_mut(), sa.as_ref());
1225}
1226
1227/// Decrypts degree 1 RLWE ciphertext RLWE(m) and returns m
1228///
1229/// - rlwe_ct: input degree 1 ciphertext RLWE(m).
1230pub(crate) fn decrypt_rlwe<
1231    R: RowMut,
1232    M: Matrix<MatElement = R::Element>,
1233    ModOp: VectorOps<Element = R::Element>,
1234    NttOp: Ntt<Element = R::Element>,
1235    S,
1236>(
1237    rlwe_ct: &M,
1238    s: &[S],
1239    m_out: &mut R,
1240    ntt_op: &NttOp,
1241    mod_op: &ModOp,
1242) where
1243    R: TryConvertFrom<[S], Parameters = R::Element>,
1244    R::Element: Copy,
1245{
1246    let ring_size = s.len();
1247    assert!(rlwe_ct.dimension() == (2, ring_size));
1248    assert!(m_out.as_ref().len() == ring_size);
1249
1250    // transform a to evluation form
1251    m_out.as_mut().copy_from_slice(rlwe_ct.get_row_slice(0));
1252    ntt_op.forward(m_out.as_mut());
1253
1254    // -s*a
1255    let mut s = R::try_convert_from(&s, &mod_op.modulus());
1256    ntt_op.forward(s.as_mut());
1257    mod_op.elwise_mul_mut(m_out.as_mut(), s.as_ref());
1258    mod_op.elwise_neg_mut(m_out.as_mut());
1259    ntt_op.backward(m_out.as_mut());
1260
1261    // m+e = b - s*a
1262    mod_op.elwise_add_mut(m_out.as_mut(), rlwe_ct.get_row_slice(1));
1263}
1264
1265// Measures noise in degree 1 RLWE ciphertext against encoded ideal message
1266// encoded_m
1267pub(crate) fn measure_noise<
1268    Mmut: MatrixMut + Matrix,
1269    ModOp: VectorOps<Element = Mmut::MatElement>,
1270    NttOp: Ntt<Element = Mmut::MatElement>,
1271    S,
1272>(
1273    rlwe_ct: &Mmut,
1274    encoded_m_ideal: &Mmut::R,
1275    ntt_op: &NttOp,
1276    mod_op: &ModOp,
1277    s: &[S],
1278) -> f64
1279where
1280    <Mmut as Matrix>::R: RowMut,
1281    Mmut::R: RowEntity + TryConvertFrom<[S], Parameters = Mmut::MatElement>,
1282    Mmut::MatElement: PrimInt + ToPrimitive + Debug,
1283{
1284    let ring_size = s.len();
1285    assert!(rlwe_ct.dimension() == (2, ring_size));
1286    assert!(encoded_m_ideal.as_ref().len() == ring_size);
1287
1288    // -(s * a)
1289    let q = VectorOps::modulus(mod_op);
1290    let mut s = Mmut::R::try_convert_from(s, &q);
1291    ntt_op.forward(s.as_mut());
1292    let mut a = Mmut::R::zeros(ring_size);
1293    a.as_mut().copy_from_slice(rlwe_ct.get_row_slice(0));
1294    ntt_op.forward(a.as_mut());
1295    mod_op.elwise_mul_mut(s.as_mut(), a.as_ref());
1296    mod_op.elwise_neg_mut(s.as_mut());
1297    ntt_op.backward(s.as_mut());
1298
1299    // m+e = b - s*a
1300    let mut m_plus_e = s;
1301    mod_op.elwise_add_mut(m_plus_e.as_mut(), rlwe_ct.get_row_slice(1));
1302
1303    // difference
1304    mod_op.elwise_sub_mut(m_plus_e.as_mut(), encoded_m_ideal.as_ref());
1305
1306    let mut max_diff_bits = f64::MIN;
1307    m_plus_e.as_ref().iter().for_each(|v| {
1308        let mut v = *v;
1309        println!("{:?}", v);
1310        if v >= (q >> 1) {
1311            // v is -ve
1312            v = q - v;
1313        }
1314
1315        let bits = (v.to_f64().unwrap()).log2();
1316
1317        if max_diff_bits < bits {
1318            max_diff_bits = bits;
1319        }
1320    });
1321
1322    return max_diff_bits;
1323}
1324
1325#[cfg(test)]
1326pub(crate) mod tests {
1327    use std::{ops::Mul, vec};
1328
1329    use itertools::{izip, Itertools};
1330    use rand::{thread_rng, Rng};
1331
1332    use crate::{
1333        backend::{ModInit, ModularOpsU64, VectorOps},
1334        decomposer::{gadget_vector, DefaultDecomposer},
1335        ntt::{self, Ntt, NttBackendU64, NttInit},
1336        random::{DefaultSecureRng, NewWithSeed, RandomUniformDist},
1337        rgsw::{
1338            gen_rlwe_public_key, measure_noise, public_key_encrypt_rgsw, AutoKeyEvaluationDomain,
1339            RgswCiphertext, RgswCiphertextEvaluationDomain, RlweCiphertext, RlwePublicKey,
1340            SeededAutoKey, SeededRgswCiphertext, SeededRlweCiphertext, SeededRlwePublicKey,
1341        },
1342        utils::{generate_prime, negacyclic_mul, TryConvertFrom},
1343        Matrix, Secret,
1344    };
1345
1346    use super::{
1347        decrypt_rlwe, galois_auto, galois_key_gen, generate_auto_map, rgsw_by_rgsw_inplace,
1348        rlwe_by_rgsw, secret_key_encrypt_rgsw, secret_key_encrypt_rlwe, RlweSecret,
1349    };
1350
1351    #[test]
1352    fn rlwe_encrypt_decryption() {
1353        let logq = 50;
1354        let logp = 2;
1355        let ring_size = 1 << 4;
1356        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
1357        let p = 1u64 << logp;
1358
1359        let mut rng = DefaultSecureRng::new();
1360
1361        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1362
1363        // sample m0
1364        let mut m0 = vec![0u64; ring_size as usize];
1365        RandomUniformDist::<[u64]>::random_fill(&mut rng, &(1u64 << logp), m0.as_mut_slice());
1366
1367        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1368        let mod_op = ModularOpsU64::new(q);
1369
1370        // encrypt m0
1371        let mut rlwe_seed = [0u8; 32];
1372        rng.fill_bytes(&mut rlwe_seed);
1373        let mut seeded_rlwe_in_ct =
1374            SeededRlweCiphertext::<_, [u8; 32]>::empty(ring_size as usize, rlwe_seed, q);
1375        let mut p_rng = DefaultSecureRng::new_with_seed(rlwe_seed);
1376        let encoded_m = m0
1377            .iter()
1378            .map(|v| (((*v as f64) * q as f64) / (p as f64)).round() as u64)
1379            .collect_vec();
1380        secret_key_encrypt_rlwe(
1381            &encoded_m,
1382            &mut seeded_rlwe_in_ct.data,
1383            s.values(),
1384            &mod_op,
1385            &ntt_op,
1386            &mut p_rng,
1387            &mut rng,
1388        );
1389        let rlwe_in_ct =
1390            RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_rlwe_in_ct);
1391
1392        let mut encoded_m_back = vec![0u64; ring_size as usize];
1393        decrypt_rlwe(
1394            &rlwe_in_ct,
1395            s.values(),
1396            &mut encoded_m_back,
1397            &ntt_op,
1398            &mod_op,
1399        );
1400        let m_back = encoded_m_back
1401            .iter()
1402            .map(|v| (((*v as f64 * p as f64) / q as f64).round() as u64) % p)
1403            .collect_vec();
1404        assert_eq!(m0, m_back);
1405    }
1406
1407    #[test]
1408    fn rlwe_by_rgsw_works() {
1409        let logq = 50;
1410        let logp = 2;
1411        let ring_size = 1 << 9;
1412        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
1413        let p = 1u64 << logp;
1414        let d_rgsw = 10;
1415        let logb = 5;
1416
1417        let mut rng = DefaultSecureRng::new_seeded([0u8; 32]);
1418
1419        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1420
1421        let mut m0 = vec![0u64; ring_size as usize];
1422        RandomUniformDist::<[u64]>::random_fill(&mut rng, &(1u64 << logp), m0.as_mut_slice());
1423        let mut m1 = vec![0u64; ring_size as usize];
1424        m1[thread_rng().gen_range(0..ring_size) as usize] = 1;
1425
1426        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1427        let mod_op = ModularOpsU64::new(q);
1428        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
1429
1430        // Encrypt m1 as RGSW(m1)
1431        let rgsw_ct = {
1432            //TODO(Jay): Figure out better way to test secret key and public key variant of
1433            // RGSW ciphertext encryption within the same test
1434
1435            if true {
1436                // Encryption m1 as RGSW(m1) using secret key
1437                _sk_encrypt_rgsw(&m1, s.values(), &gadget_vector, &mod_op, &ntt_op)
1438            } else {
1439                // Encrypt m1 as RGSW(m1) as public key
1440
1441                // first create public key
1442                let mut pk_seed = [0u8; 32];
1443                rng.fill_bytes(&mut pk_seed);
1444                let mut pk_prng = DefaultSecureRng::new_seeded(pk_seed);
1445                let mut seeded_pk =
1446                    SeededRlwePublicKey::<Vec<u64>, _>::empty(ring_size as usize, pk_seed, q);
1447                gen_rlwe_public_key(
1448                    &mut seeded_pk.data,
1449                    s.values(),
1450                    &ntt_op,
1451                    &mod_op,
1452                    &mut pk_prng,
1453                    &mut rng,
1454                );
1455                let pk = RlwePublicKey::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_pk);
1456
1457                let rgsw_ct = _pk_encrypt_rgsw(&m1, &pk, &gadget_vector, &mod_op, &ntt_op);
1458                RgswCiphertextEvaluationDomain::<_, DefaultSecureRng, NttBackendU64>::from(
1459                    &RgswCiphertext {
1460                        data: rgsw_ct.data,
1461                        modulus: q,
1462                    },
1463                )
1464            }
1465        };
1466
1467        // Encrypt m0 as RLWE(m0)
1468        let mut rlwe_in_ct = {
1469            let mut rlwe_seed = [0u8; 32];
1470            rng.fill_bytes(&mut rlwe_seed);
1471            let mut seeded_rlwe_in_ct =
1472                SeededRlweCiphertext::<_, [u8; 32]>::empty(ring_size as usize, rlwe_seed, q);
1473            let mut p_rng = DefaultSecureRng::new_seeded(rlwe_seed);
1474            let encoded_m = m0
1475                .iter()
1476                .map(|v| (((*v as f64) * q as f64) / (p as f64)).round() as u64)
1477                .collect_vec();
1478            secret_key_encrypt_rlwe(
1479                &encoded_m,
1480                &mut seeded_rlwe_in_ct.data,
1481                s.values(),
1482                &mod_op,
1483                &ntt_op,
1484                &mut p_rng,
1485                &mut rng,
1486            );
1487
1488            RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_rlwe_in_ct)
1489        };
1490
1491        // RLWE(m0m1) = RLWE(m0) x RGSW(m1)
1492        let mut scratch_space = vec![vec![0u64; ring_size as usize]; d_rgsw + 2];
1493        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
1494        rlwe_by_rgsw(
1495            &mut rlwe_in_ct,
1496            &rgsw_ct.data,
1497            &mut scratch_space,
1498            &decomposer,
1499            &ntt_op,
1500            &mod_op,
1501        );
1502
1503        // Decrypt RLWE(m0m1)
1504        let mut encoded_m0m1_back = vec![0u64; ring_size as usize];
1505        decrypt_rlwe(
1506            &rlwe_in_ct,
1507            s.values(),
1508            &mut encoded_m0m1_back,
1509            &ntt_op,
1510            &mod_op,
1511        );
1512        let m0m1_back = encoded_m0m1_back
1513            .iter()
1514            .map(|v| (((*v as f64 * p as f64) / (q as f64)).round() as u64) % p)
1515            .collect_vec();
1516
1517        let mul_mod = |v0: &u64, v1: &u64| (v0 * v1) % p;
1518        let m0m1 = negacyclic_mul(&m0, &m1, mul_mod, p);
1519
1520        {
1521            // measure noise
1522            let encoded_m_ideal = m0m1
1523                .iter()
1524                .map(|v| (((*v as f64) * q as f64) / (p as f64)).round() as u64)
1525                .collect_vec();
1526
1527            let noise = measure_noise(&rlwe_in_ct, &encoded_m_ideal, &ntt_op, &mod_op, s.values());
1528            println!("Noise RLWE(m0m1)(= RLWE(m0)xRGSW(m1)) : {noise}");
1529        }
1530
1531        assert!(
1532            m0m1 == m0m1_back,
1533            "Expected {:?} \n Got {:?}",
1534            m0m1,
1535            m0m1_back
1536        );
1537    }
1538
1539    #[test]
1540    fn rlwe_by_rgsw_noise_growth() {
1541        let logq = 31;
1542        let ring_size = 1 << 9;
1543        let q = generate_prime(logq, ring_size * 2, 1u64 << logq).unwrap();
1544        let d_rgsw = 6;
1545        let logb = 5;
1546
1547        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1548
1549        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1550        let mod_op = ModularOpsU64::new(q);
1551        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
1552        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
1553
1554        let mul_mod = |v0: &u64, v1: &u64| ((*v0 as u128 * *v1 as u128) % (q as u128)) as u64;
1555
1556        let mut carry_m = vec![0u64; ring_size as usize];
1557        carry_m[thread_rng().gen_range(0..ring_size) as usize] = 1;
1558        let mut rlwe = vec![vec![0u64; ring_size as usize], carry_m.clone()];
1559        let mut rlwe = RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from_raw(rlwe, true);
1560
1561        let mut scratch_matrix_dplus2_ring = vec![vec![0u64; ring_size as usize]; d_rgsw + 2];
1562        for i in 0..1usize {
1563            // Encrypt monomial as RGSW
1564            let mut m = vec![0u64; ring_size as usize];
1565            m[thread_rng().gen_range(0..ring_size) as usize] = if i & 1 == 1 { 1 } else { q - 1 };
1566            let rgsw_ct = _sk_encrypt_rgsw(&m, s.values(), &gadget_vector, &mod_op, &ntt_op);
1567
1568            // RLWE(carry_m * m) = RLWE(carry_m) x RGSW(m)
1569            rlwe_by_rgsw(
1570                &mut rlwe,
1571                &rgsw_ct.data,
1572                &mut scratch_matrix_dplus2_ring,
1573                &decomposer,
1574                &ntt_op,
1575                &mod_op,
1576            );
1577
1578            carry_m = negacyclic_mul(&carry_m, &m, mul_mod, q);
1579            let noise = measure_noise(&rlwe, &carry_m, &ntt_op, &mod_op, s.values());
1580
1581            println!("Noise RLWE(carry_m) after {i}^th iteration: {noise}");
1582        }
1583    }
1584
1585    fn _pk_encrypt_rgsw(
1586        m: &[u64],
1587        public_key: &RlwePublicKey<Vec<Vec<u64>>, DefaultSecureRng>,
1588        gadget_vector: &[u64],
1589        mod_op: &ModularOpsU64,
1590        ntt_op: &NttBackendU64,
1591    ) -> RgswCiphertext<Vec<Vec<u64>>> {
1592        let (_, ring_size) = Matrix::dimension(&public_key.data);
1593        let d_rgsw = gadget_vector.len();
1594
1595        let mut rng = DefaultSecureRng::new();
1596
1597        assert!(m.len() == ring_size);
1598
1599        // public key encrypt RGSW(m1)
1600        let mut rgsw_ct = vec![vec![0u64; ring_size]; d_rgsw * 4];
1601        public_key_encrypt_rgsw(
1602            &mut rgsw_ct,
1603            m,
1604            &public_key.data,
1605            gadget_vector,
1606            mod_op,
1607            ntt_op,
1608            &mut rng,
1609        );
1610
1611        RgswCiphertext {
1612            data: rgsw_ct,
1613            modulus: mod_op.modulus(),
1614        }
1615    }
1616
1617    /// Encrypts m as RGSW ciphertext RGSW(m) using supplied secret key. Returns
1618    /// unseeded RGSW ciphertext in coefficient domain
1619    fn _sk_encrypt_rgsw(
1620        m: &[u64],
1621        s: &[i32],
1622        gadget_vector: &[u64],
1623        mod_op: &ModularOpsU64,
1624        ntt_op: &NttBackendU64,
1625    ) -> RgswCiphertextEvaluationDomain<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64> {
1626        let ring_size = s.len();
1627        assert!(m.len() == s.len());
1628
1629        let d_rgsw = gadget_vector.len();
1630        let q = mod_op.modulus();
1631
1632        let mut rng = DefaultSecureRng::new();
1633        let mut rgsw_seed = [0u8; 32];
1634        rng.fill_bytes(&mut rgsw_seed);
1635        let mut seeded_rgsw_ct = SeededRgswCiphertext::<Vec<Vec<u64>>, [u8; 32]>::empty(
1636            ring_size as usize,
1637            d_rgsw,
1638            rgsw_seed,
1639            q,
1640        );
1641        let mut p_rng = DefaultSecureRng::new_seeded(rgsw_seed);
1642        secret_key_encrypt_rgsw(
1643            &mut seeded_rgsw_ct.data,
1644            m,
1645            &gadget_vector,
1646            s,
1647            mod_op,
1648            ntt_op,
1649            &mut p_rng,
1650            &mut rng,
1651        );
1652
1653        RgswCiphertextEvaluationDomain::<_, DefaultSecureRng, NttBackendU64>::from(&seeded_rgsw_ct)
1654    }
1655
1656    /// Prints noise in RGSW ciphertext RGSW(m).
1657    ///
1658    /// - rgsw_ct: RGSW ciphertext in coefficient domain
1659    pub(crate) fn _measure_noise_rgsw(
1660        rgsw_ct: &[Vec<u64>],
1661        m: &[u64],
1662        s: &[i32],
1663        gadget_vector: &[u64],
1664        q: u64,
1665    ) {
1666        let d_rgsw = gadget_vector.len();
1667        let ring_size = s.len();
1668        assert!(Matrix::dimension(&rgsw_ct) == (d_rgsw * 2 * 2, ring_size));
1669        assert!(m.len() == ring_size);
1670
1671        let mod_op = ModularOpsU64::new(q);
1672        let ntt_op = NttBackendU64::new(q, ring_size);
1673
1674        let mul_mod = |a: &u64, b: &u64| ((*a as u128 * *b as u128) % q as u128) as u64;
1675        let s_poly = Vec::<u64>::try_convert_from(s, &q);
1676        let mut neg_s = s_poly.clone();
1677        mod_op.elwise_neg_mut(neg_s.as_mut());
1678        let neg_sm0m1 = negacyclic_mul(&neg_s, &m, mul_mod, q);
1679        for i in 0..2 {
1680            for j in 0..d_rgsw {
1681                let ideal_m = {
1682                    if i == 0 {
1683                        // RLWE(\beta^j -s * m)
1684                        let mut beta_neg_sm0m1 = vec![0u64; ring_size as usize];
1685                        mod_op.elwise_scalar_mul(
1686                            beta_neg_sm0m1.as_mut(),
1687                            &neg_sm0m1,
1688                            &gadget_vector[j],
1689                        );
1690                        beta_neg_sm0m1
1691                    } else {
1692                        // RLWE(\beta^j  m)
1693                        let mut beta_m0m1 = vec![0u64; ring_size as usize];
1694                        mod_op.elwise_scalar_mul(beta_m0m1.as_mut(), &m, &gadget_vector[j]);
1695                        beta_m0m1
1696                    }
1697                };
1698
1699                let mut rlwe = vec![vec![0u64; ring_size as usize]; 2];
1700                rlwe[0].copy_from_slice(rgsw_ct.get_row_slice((i * 2 * d_rgsw) + j));
1701                rlwe[1].copy_from_slice(rgsw_ct.get_row_slice((i * 2 * d_rgsw) + d_rgsw + j));
1702                let noise = measure_noise(&rlwe, &ideal_m, &ntt_op, &mod_op, s);
1703
1704                if i == 0 {
1705                    println!(r"Noise RLWE(\beta^{j} -sm0m1): {noise}");
1706                } else {
1707                    println!(r"Noise RLWE(\beta^{j} m0m1): {noise}");
1708                }
1709            }
1710            // m0m1
1711        }
1712    }
1713
1714    #[test]
1715    fn pk_rgsw_by_rgsw() {
1716        let logq = 60;
1717        let logp = 2;
1718        let ring_size = 1 << 11;
1719        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
1720        let p = 1u64 << logp;
1721        let d_rgsw = 15;
1722        let logb = 4;
1723
1724        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1725
1726        let mut rng = DefaultSecureRng::new();
1727        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1728        let mod_op = ModularOpsU64::new(q);
1729        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
1730        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
1731
1732        // Public Key
1733        let public_key = {
1734            let mut pk_seed = [0u8; 32];
1735            rng.fill_bytes(&mut pk_seed);
1736            let mut pk_prng = DefaultSecureRng::new_seeded(pk_seed);
1737            let mut seeded_pk =
1738                SeededRlwePublicKey::<Vec<u64>, _>::empty(ring_size as usize, pk_seed, q);
1739            gen_rlwe_public_key(
1740                &mut seeded_pk.data,
1741                s.values(),
1742                &ntt_op,
1743                &mod_op,
1744                &mut pk_prng,
1745                &mut rng,
1746            );
1747            RlwePublicKey::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_pk)
1748        };
1749
1750        let mut carry_m = vec![0u64; ring_size as usize];
1751        carry_m[thread_rng().gen_range(0..ring_size) as usize] = 1;
1752
1753        // RGSW(carry_m)
1754        let rgsw_carrym = _pk_encrypt_rgsw(&carry_m, &public_key, &gadget_vector, &mod_op, &ntt_op);
1755        let mut rgsw_carrym =
1756            RgswCiphertextEvaluationDomain::<_, DefaultSecureRng, NttBackendU64>::from(
1757                &rgsw_carrym,
1758            );
1759
1760        let mut scratch_matrix_d_plus_rgsw_by_ring =
1761            vec![vec![0u64; ring_size as usize]; d_rgsw + (d_rgsw * 4)];
1762
1763        for i in 0..10 {
1764            let mut m = vec![0u64; ring_size as usize];
1765            m[thread_rng().gen_range(0..ring_size) as usize] = q - 1;
1766            let rgsw_m = _pk_encrypt_rgsw(&m, &public_key, &gadget_vector, &mod_op, &ntt_op);
1767
1768            rgsw_by_rgsw_inplace(
1769                &mut rgsw_carrym.data,
1770                &rgsw_m.data,
1771                &decomposer,
1772                &mut scratch_matrix_d_plus_rgsw_by_ring,
1773                &ntt_op,
1774                &mod_op,
1775            );
1776
1777            // measure noise
1778            let mul_mod = |a: &u64, b: &u64| ((*a as u128 * *b as u128) % q as u128) as u64;
1779            carry_m = negacyclic_mul(&carry_m, &m, mul_mod, q);
1780            println!("########### Noise RGSW(carrym) in {i}^th loop ###########");
1781            let mut rgsw_carrym_clone = rgsw_carrym.data.clone();
1782            rgsw_carrym_clone
1783                .iter_mut()
1784                .for_each(|ri| ntt_op.backward(ri.as_mut()));
1785            _measure_noise_rgsw(&rgsw_carrym_clone, &carry_m, s.values(), &gadget_vector, q);
1786        }
1787    }
1788
1789    #[test]
1790    fn sk_rgsw_by_rgsw() {
1791        let logq = 60;
1792        let logp = 2;
1793        let ring_size = 1 << 11;
1794        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
1795        let p = 1u64 << logp;
1796        let d_rgsw = 15;
1797        let logb = 4;
1798
1799        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1800
1801        let mut rng = DefaultSecureRng::new();
1802        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1803        let mod_op = ModularOpsU64::new(q);
1804        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
1805        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
1806
1807        let mut carry_m = vec![0u64; ring_size as usize];
1808        carry_m[thread_rng().gen_range(0..ring_size) as usize] = 1;
1809
1810        // RGSW(carry_m)
1811        let mut rgsw_carrym =
1812            _sk_encrypt_rgsw(&carry_m, s.values(), &gadget_vector, &mod_op, &ntt_op);
1813
1814        let mut scratch_matrix_d_plus_rgsw_by_ring =
1815            vec![vec![0u64; ring_size as usize]; d_rgsw + (d_rgsw * 4)];
1816
1817        for i in 0..1 {
1818            let mut m = vec![0u64; ring_size as usize];
1819            m[thread_rng().gen_range(0..ring_size) as usize] = if (i & 1) == 1 { q - 1 } else { 1 };
1820            let rgsw_m = {
1821                let mut rgsw_eval =
1822                    _sk_encrypt_rgsw(&m, s.values(), &gadget_vector, &mod_op, &ntt_op).data;
1823                rgsw_eval
1824                    .iter_mut()
1825                    .for_each(|ri| ntt_op.backward(ri.as_mut()));
1826                rgsw_eval
1827            };
1828
1829            rgsw_by_rgsw_inplace(
1830                &mut rgsw_carrym.data,
1831                &rgsw_m,
1832                &decomposer,
1833                &mut scratch_matrix_d_plus_rgsw_by_ring,
1834                &ntt_op,
1835                &mod_op,
1836            );
1837
1838            // measure noise
1839            let mul_mod = |a: &u64, b: &u64| ((*a as u128 * *b as u128) % q as u128) as u64;
1840            carry_m = negacyclic_mul(&carry_m, &m, mul_mod, q);
1841            println!("########### Noise RGSW(carrym) in {i}^th loop ###########");
1842            let mut rgsw_carrym_clone = rgsw_carrym.data.clone();
1843            rgsw_carrym_clone
1844                .iter_mut()
1845                .for_each(|ri| ntt_op.backward(ri.as_mut()));
1846            _measure_noise_rgsw(&rgsw_carrym_clone, &carry_m, s.values(), &gadget_vector, q);
1847        }
1848    }
1849
1850    #[test]
1851    fn galois_auto_works() {
1852        let logq = 50;
1853        let ring_size = 1 << 4;
1854        let q = generate_prime(logq, 2 * ring_size, 1u64 << logq).unwrap();
1855        let logp = 3;
1856        let p = 1u64 << logp;
1857        let d_rgsw = 10;
1858        let logb = 5;
1859
1860        let mut rng = DefaultSecureRng::new();
1861        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1862
1863        let mut m = vec![0u64; ring_size as usize];
1864        RandomUniformDist::random_fill(&mut rng, &p, m.as_mut_slice());
1865        let encoded_m = m
1866            .iter()
1867            .map(|v| (((*v as f64 * q as f64) / (p as f64)).round() as u64))
1868            .collect_vec();
1869
1870        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1871        let mod_op = ModularOpsU64::new(q);
1872
1873        // RLWE_{s}(m)
1874        let mut seed_rlwe = [0u8; 32];
1875        rng.fill_bytes(&mut seed_rlwe);
1876        let mut seeded_rlwe_m = SeededRlweCiphertext::empty(ring_size as usize, seed_rlwe, q);
1877        let mut p_rng = DefaultSecureRng::new_seeded(seed_rlwe);
1878        secret_key_encrypt_rlwe(
1879            &encoded_m,
1880            &mut seeded_rlwe_m.data,
1881            s.values(),
1882            &mod_op,
1883            &ntt_op,
1884            &mut p_rng,
1885            &mut rng,
1886        );
1887        let mut rlwe_m = RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_rlwe_m);
1888
1889        let auto_k = -5;
1890
1891        // Generate galois key to key switch from s^k to s
1892        let mut seed_auto = [0u8; 32];
1893        rng.fill_bytes(&mut seed_auto);
1894        let mut seeded_auto_key = SeededAutoKey::empty(ring_size as usize, d_rgsw, seed_auto, q);
1895        let mut p_rng = DefaultSecureRng::new_seeded(seed_auto);
1896        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
1897        galois_key_gen(
1898            &mut seeded_auto_key.data,
1899            s.values(),
1900            auto_k,
1901            &gadget_vector,
1902            &mod_op,
1903            &ntt_op,
1904            &mut p_rng,
1905            &mut rng,
1906        );
1907        let auto_key =
1908            AutoKeyEvaluationDomain::<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64>::from(
1909                &seeded_auto_key,
1910            );
1911
1912        // Send RLWE_{s}(m) -> RLWE_{s}(m^k)
1913        let mut scratch_space = vec![vec![0u64; ring_size as usize]; d_rgsw + 2];
1914        let (auto_map_index, auto_map_sign) = generate_auto_map(ring_size as usize, auto_k);
1915        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
1916        galois_auto(
1917            &mut rlwe_m,
1918            &auto_key.data,
1919            &mut scratch_space,
1920            &auto_map_index,
1921            &auto_map_sign,
1922            &mod_op,
1923            &ntt_op,
1924            &decomposer,
1925        );
1926
1927        let rlwe_m_k = rlwe_m;
1928
1929        // Decrypt RLWE_{s}(m^k) and check
1930        let mut encoded_m_k_back = vec![0u64; ring_size as usize];
1931        decrypt_rlwe(
1932            &rlwe_m_k,
1933            s.values(),
1934            &mut encoded_m_k_back,
1935            &ntt_op,
1936            &mod_op,
1937        );
1938        let m_k_back = encoded_m_k_back
1939            .iter()
1940            .map(|v| (((*v as f64 * p as f64) / q as f64).round() as u64) % p)
1941            .collect_vec();
1942
1943        let mut m_k = vec![0u64; ring_size as usize];
1944        // Send \delta m -> \delta m^k
1945        izip!(m.iter(), auto_map_index.iter(), auto_map_sign.iter()).for_each(
1946            |(v, to_index, sign)| {
1947                if !*sign {
1948                    m_k[*to_index] = (p - *v) % p;
1949                } else {
1950                    m_k[*to_index] = *v;
1951                }
1952            },
1953        );
1954
1955        {
1956            let encoded_m_k = m_k
1957                .iter()
1958                .map(|v| ((*v as f64 * q as f64) / p as f64).round() as u64)
1959                .collect_vec();
1960
1961            let noise = measure_noise(&rlwe_m_k, &encoded_m_k, &ntt_op, &mod_op, s.values());
1962            println!("Ksk noise: {noise}");
1963        }
1964
1965        // FIXME(Jay): Galios autormophism will incur high error unless we fix in
1966        // accurate decomoposition of Decomposer when q is prime
1967        assert_eq!(m_k_back, m_k);
1968        // dbg!(m_k_back, m_k, q);
1969    }
1970}