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, Signed, 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    pub(crate) 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
 466pub(crate) fn routine<R: RowMut, ModOp: VectorOps<Element = R::Element>>(
 467    write_to_row: &mut [R::Element],
 468    matrix_a: &[R],
 469    matrix_b: &[R],
 470    mod_op: &ModOp,
 471) {
 472    izip!(matrix_a.iter(), matrix_b.iter()).for_each(|(a, b)| {
 473        mod_op.elwise_fma_mut(write_to_row, a.as_ref(), b.as_ref());
 474    });
 475}
 476
 477/// Decomposes ring polynomial r(X) into d polynomials using decomposer into
 478/// output matrix decomp_r
 479///
 480/// Note that decomposition of r(X) requires decomposition of each of
 481/// coefficients.
 482///
 483/// - decomp_r: must have dimensions d x ring_size. i^th decomposed polynomial
 484///   will be stored at i^th row.
 485pub(crate) fn decompose_r<R: RowMut, D: Decomposer<Element = R::Element>>(
 486    r: &[R::Element],
 487    decomp_r: &mut [R],
 488    decomposer: &D,
 489) where
 490    R::Element: Copy,
 491{
 492    let ring_size = r.len();
 493    let d = decomposer.d();
 494
 495    for ri in 0..ring_size {
 496        let el_decomposed = decomposer.decompose(&r[ri]);
 497        for j in 0..d {
 498            decomp_r[j].as_mut()[ri] = el_decomposed[j];
 499        }
 500    }
 501}
 502
 503/// Sends RLWE_{s}(X) -> RLWE_{s}(X^k) where k is some galois element
 504pub(crate) fn galois_auto<
 505    MT: Matrix + IsTrivial + MatrixMut,
 506    Mmut: MatrixMut<MatElement = MT::MatElement>,
 507    ModOp: ArithmeticOps<Element = MT::MatElement> + VectorOps<Element = MT::MatElement>,
 508    NttOp: Ntt<Element = MT::MatElement>,
 509    D: Decomposer<Element = MT::MatElement>,
 510>(
 511    rlwe_in: &mut MT,
 512    ksk: &Mmut,
 513    scratch_matrix_dplus2_ring: &mut Mmut,
 514    auto_map_index: &[usize],
 515    auto_map_sign: &[bool],
 516    mod_op: &ModOp,
 517    ntt_op: &NttOp,
 518    decomposer: &D,
 519) where
 520    <Mmut as Matrix>::R: RowMut,
 521    <MT as Matrix>::R: RowMut,
 522    MT::MatElement: Copy + Zero,
 523{
 524    let d = decomposer.d();
 525
 526    let (scratch_matrix_d_ring, tmp_rlwe_out) = scratch_matrix_dplus2_ring.split_at_row_mut(d);
 527
 528    // send b(X) -> b(X^k)
 529    izip!(
 530        rlwe_in.get_row(1),
 531        auto_map_index.iter(),
 532        auto_map_sign.iter()
 533    )
 534    .for_each(|(el_in, to_index, sign)| {
 535        if !*sign {
 536            tmp_rlwe_out[1].as_mut()[*to_index] = mod_op.neg(el_in);
 537        } else {
 538            tmp_rlwe_out[1].as_mut()[*to_index] = *el_in;
 539            // scratch_matrix_dplus2_ring.set(d + 1, *to_index, *el_in);
 540        }
 541    });
 542
 543    if !rlwe_in.is_trivial() {
 544        // send a(X) -> a(X^k) and decompose a(X^k)
 545        izip!(
 546            rlwe_in.get_row(0),
 547            auto_map_index.iter(),
 548            auto_map_sign.iter()
 549        )
 550        .for_each(|(el_in, to_index, sign)| {
 551            let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
 552
 553            let el_out_decomposed = decomposer.decompose(&el_out);
 554            for j in 0..d {
 555                scratch_matrix_d_ring[j].as_mut()[*to_index] = el_out_decomposed[j];
 556            }
 557        });
 558
 559        // transform decomposed a(X^k) to evaluation domain
 560        scratch_matrix_d_ring.iter_mut().for_each(|r| {
 561            ntt_op.forward(r.as_mut());
 562        });
 563
 564        // RLWE(m^k) = a', b'; RLWE(m) = a, b
 565        // key switch: (a * RLWE'(s(X^k)))
 566        let (ksk_a, ksk_b) = ksk.split_at_row(d);
 567        tmp_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
 568        // a' = decomp<a> * RLWE'_A(s(X^k))
 569        routine(
 570            tmp_rlwe_out[0].as_mut(),
 571            scratch_matrix_d_ring,
 572            ksk_a,
 573            mod_op,
 574        );
 575        // send b(X^k) to evaluation domain
 576        ntt_op.forward(tmp_rlwe_out[1].as_mut());
 577        // b' = b(X^k)
 578        // b' += decomp<a(X^k)> * RLWE'_B(s(X^k))
 579        routine(
 580            tmp_rlwe_out[1].as_mut(),
 581            scratch_matrix_d_ring,
 582            ksk_b,
 583            mod_op,
 584        );
 585
 586        // transform RLWE(m^k) to coefficient domain
 587        tmp_rlwe_out
 588            .iter_mut()
 589            .for_each(|r| ntt_op.backward(r.as_mut()));
 590
 591        rlwe_in
 592            .get_row_mut(0)
 593            .copy_from_slice(tmp_rlwe_out[0].as_ref());
 594    }
 595
 596    rlwe_in
 597        .get_row_mut(1)
 598        .copy_from_slice(tmp_rlwe_out[1].as_ref());
 599}
 600
 601/// Returns RLWE(m0m1) = RLWE(m0) x RGSW(m1). Mutates rlwe_in inplace to equal
 602/// RLWE(m0m1)
 603///
 604/// - rlwe_in: is RLWE(m0) with polynomials in coefficient domain
 605/// - rgsw_in: is RGSW(m1) with polynomials in evaluation domain
 606/// - scratch_matrix_d_ring: is a matrix of dimension (d_rgsw, ring_size) used
 607///   as scratch space to store decomposed Ring elements temporarily
 608pub(crate) fn rlwe_by_rgsw<
 609    Mmut: MatrixMut,
 610    MT: Matrix<MatElement = Mmut::MatElement> + MatrixMut<MatElement = Mmut::MatElement> + IsTrivial,
 611    D: Decomposer<Element = Mmut::MatElement>,
 612    ModOp: VectorOps<Element = Mmut::MatElement>,
 613    NttOp: Ntt<Element = Mmut::MatElement>,
 614>(
 615    rlwe_in: &mut MT,
 616    rgsw_in: &Mmut,
 617    scratch_matrix_dplus2_ring: &mut Mmut,
 618    decomposer: &D,
 619    ntt_op: &NttOp,
 620    mod_op: &ModOp,
 621) where
 622    Mmut::MatElement: Copy + Zero,
 623    <Mmut as Matrix>::R: RowMut,
 624    <MT as Matrix>::R: RowMut,
 625{
 626    let d_rgsw = decomposer.d();
 627    assert!(scratch_matrix_dplus2_ring.dimension() == (d_rgsw + 2, rlwe_in.dimension().1));
 628    assert!(rgsw_in.dimension() == (d_rgsw * 4, rlwe_in.dimension().1));
 629
 630    // decomposed RLWE x RGSW
 631    let (rlwe_dash_nsm, rlwe_dash_m) = rgsw_in.split_at_row(d_rgsw * 2);
 632    let (scratch_matrix_d_ring, scratch_rlwe_out) =
 633        scratch_matrix_dplus2_ring.split_at_row_mut(d_rgsw);
 634    scratch_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
 635    scratch_rlwe_out[1].as_mut().fill(Mmut::MatElement::zero());
 636    // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
 637    if !rlwe_in.is_trivial() {
 638        // a_in = 0 when RLWE_in is trivial RLWE ciphertext
 639        // decomp<a_in>
 640        decompose_r(rlwe_in.get_row_slice(0), scratch_matrix_d_ring, decomposer);
 641        scratch_matrix_d_ring
 642            .iter_mut()
 643            .for_each(|r| ntt_op.forward(r.as_mut()));
 644        // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
 645        routine(
 646            scratch_rlwe_out[0].as_mut(),
 647            scratch_matrix_d_ring.as_ref(),
 648            &rlwe_dash_nsm[..d_rgsw],
 649            mod_op,
 650        );
 651        // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
 652        routine(
 653            scratch_rlwe_out[1].as_mut(),
 654            scratch_matrix_d_ring.as_ref(),
 655            &rlwe_dash_nsm[d_rgsw..],
 656            mod_op,
 657        );
 658    }
 659    // decomp<b_in>
 660    decompose_r(rlwe_in.get_row_slice(1), scratch_matrix_d_ring, decomposer);
 661    scratch_matrix_d_ring
 662        .iter_mut()
 663        .for_each(|r| ntt_op.forward(r.as_mut()));
 664    // a_out += decomp<b_in> \cdot RLWE_A'(m)
 665    routine(
 666        scratch_rlwe_out[0].as_mut(),
 667        scratch_matrix_d_ring.as_ref(),
 668        &rlwe_dash_m[..d_rgsw],
 669        mod_op,
 670    );
 671    // b_out += decomp<b_in> \cdot RLWE_B'(m)
 672    routine(
 673        scratch_rlwe_out[1].as_mut(),
 674        scratch_matrix_d_ring.as_ref(),
 675        &rlwe_dash_m[d_rgsw..],
 676        mod_op,
 677    );
 678
 679    // transform rlwe_out to coefficient domain
 680    scratch_rlwe_out
 681        .iter_mut()
 682        .for_each(|r| ntt_op.backward(r.as_mut()));
 683
 684    rlwe_in
 685        .get_row_mut(0)
 686        .copy_from_slice(scratch_rlwe_out[0].as_mut());
 687    rlwe_in
 688        .get_row_mut(1)
 689        .copy_from_slice(scratch_rlwe_out[1].as_mut());
 690    rlwe_in.set_not_trivial();
 691}
 692
 693/// Inplace mutates rlwe_0 to equal RGSW(m0m1) = RGSW(m0)xRGSW(m1)
 694/// in evaluation domain
 695///
 696/// Warning -
 697/// Pass a fresh RGSW ciphertext as the second operand, i.e. as `rgsw_1`.
 698/// This is to assure minimal error growth in the resulting RGSW ciphertext.
 699/// RGSWxRGSW boils down to d_rgsw*2 RLWExRGSW multiplications. Hence, the noise
 700/// growth in resulting ciphertext depends on the norm of second RGSW
 701/// ciphertext, not the first. This is useful in cases where one is accumulating
 702/// multiple RGSW ciphertexts into 1. In which case, pass the accumulating RGSW
 703/// ciphertext as rlwe_0 (the one with higher noise) and subsequent RGSW
 704/// ciphertexts, with lower noise, to be accumulated as second
 705/// operand.
 706///
 707/// - rgsw_0: RGSW(m0)
 708/// - rgsw_1_eval: RGSW(m1) in Evaluation domain
 709/// - scratch_matrix_d_plus_rgsw_by_ring: scratch space matrix of size
 710///   (d+(d*4))xring_size, where d equals d_rgsw
 711pub(crate) fn rgsw_by_rgsw_inplace<
 712    Mmut: MatrixMut,
 713    D: Decomposer<Element = Mmut::MatElement>,
 714    ModOp: VectorOps<Element = Mmut::MatElement>,
 715    NttOp: Ntt<Element = Mmut::MatElement>,
 716>(
 717    rgsw_0: &mut Mmut,
 718    rgsw_1_eval: &Mmut,
 719    decomposer: &D,
 720    scratch_matrix_d_plus_rgsw_by_ring: &mut Mmut,
 721    ntt_op: &NttOp,
 722    mod_op: &ModOp,
 723) where
 724    <Mmut as Matrix>::R: RowMut,
 725    Mmut::MatElement: Copy + Zero,
 726{
 727    let d_rgsw = decomposer.d();
 728    assert!(rgsw_0.dimension().0 == 4 * d_rgsw);
 729    let ring_size = rgsw_0.dimension().1;
 730    assert!(rgsw_1_eval.dimension() == (4 * d_rgsw, ring_size));
 731    assert!(scratch_matrix_d_plus_rgsw_by_ring.dimension() == (d_rgsw + (d_rgsw * 4), ring_size));
 732
 733    let (decomp_r_space, rgsw_space) = scratch_matrix_d_plus_rgsw_by_ring.split_at_row_mut(d_rgsw);
 734
 735    // zero rgsw_space
 736    rgsw_space
 737        .iter_mut()
 738        .for_each(|ri| ri.as_mut().fill(Mmut::MatElement::zero()));
 739    let (rlwe_dash_space_nsm, rlwe_dash_space_m) = rgsw_space.split_at_mut(d_rgsw * 2);
 740    let (rlwe_dash_space_nsm_parta, rlwe_dash_space_nsm_partb) =
 741        rlwe_dash_space_nsm.split_at_mut(d_rgsw);
 742    let (rlwe_dash_space_m_parta, rlwe_dash_space_m_partb) = rlwe_dash_space_m.split_at_mut(d_rgsw);
 743
 744    let (rgsw0_nsm, rgsw0_m) = rgsw_0.split_at_row(d_rgsw * 2);
 745    let (rgsw1_nsm, rgsw1_m) = rgsw_1_eval.split_at_row(d_rgsw * 2);
 746
 747    // RGSW x RGSW
 748    izip!(
 749        rgsw0_nsm
 750            .iter()
 751            .take(d_rgsw)
 752            .chain(rgsw0_m.iter().take(d_rgsw)),
 753        rgsw0_nsm
 754            .iter()
 755            .skip(d_rgsw)
 756            .chain(rgsw0_m.iter().skip(d_rgsw)),
 757        rlwe_dash_space_nsm_parta
 758            .iter_mut()
 759            .chain(rlwe_dash_space_m_parta.iter_mut()),
 760        rlwe_dash_space_nsm_partb
 761            .iter_mut()
 762            .chain(rlwe_dash_space_m_partb.iter_mut()),
 763    )
 764    .for_each(|(rlwe_a, rlwe_b, rlwe_out_a, rlwe_out_b)| {
 765        // Part A
 766        decompose_r(rlwe_a.as_ref(), decomp_r_space.as_mut(), decomposer);
 767        decomp_r_space
 768            .iter_mut()
 769            .for_each(|ri| ntt_op.forward(ri.as_mut()));
 770        routine(
 771            rlwe_out_a.as_mut(),
 772            decomp_r_space,
 773            &rgsw1_nsm[..d_rgsw],
 774            mod_op,
 775        );
 776        routine(
 777            rlwe_out_b.as_mut(),
 778            decomp_r_space,
 779            &rgsw1_nsm[d_rgsw..],
 780            mod_op,
 781        );
 782
 783        // Part B
 784        decompose_r(rlwe_b.as_ref(), decomp_r_space.as_mut(), decomposer);
 785        decomp_r_space
 786            .iter_mut()
 787            .for_each(|ri| ntt_op.forward(ri.as_mut()));
 788        routine(
 789            rlwe_out_a.as_mut(),
 790            decomp_r_space,
 791            &rgsw1_m[..d_rgsw],
 792            mod_op,
 793        );
 794        routine(
 795            rlwe_out_b.as_mut(),
 796            decomp_r_space,
 797            &rgsw1_m[d_rgsw..],
 798            mod_op,
 799        );
 800    });
 801
 802    // copy over RGSW(m0m1) into RGSW(m0)
 803    izip!(rgsw_0.iter_rows_mut(), rgsw_space.iter())
 804        .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
 805
 806    // send back to coefficient domain
 807    rgsw_0
 808        .iter_rows_mut()
 809        .for_each(|ri| ntt_op.backward(ri.as_mut()));
 810}
 811
 812/// Encrypts message m as a RGSW ciphertext.
 813///
 814/// - m_eval: is `m` is evaluation domain
 815/// - out_rgsw: RGSW(m) is stored as single matrix of dimension (d_rgsw * 3,
 816///   ring_size). The matrix has the following structure [RLWE'_A(-sm) ||
 817///   RLWE'_B(-sm) || RLWE'_B(m)]^T and RLWE'_A(m) is generated via seed (where
 818///   p_rng is assumed to be seeded with seed)
 819pub(crate) fn secret_key_encrypt_rgsw<
 820    Mmut: MatrixMut + MatrixEntity,
 821    S,
 822    R: RandomGaussianDist<[Mmut::MatElement], Parameters = Mmut::MatElement>
 823        + RandomUniformDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
 824    PR: RandomUniformDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
 825    ModOp: VectorOps<Element = Mmut::MatElement>,
 826    NttOp: Ntt<Element = Mmut::MatElement>,
 827>(
 828    out_rgsw: &mut Mmut,
 829    m: &[Mmut::MatElement],
 830    gadget_vector: &[Mmut::MatElement],
 831    s: &[S],
 832    mod_op: &ModOp,
 833    ntt_op: &NttOp,
 834    p_rng: &mut PR,
 835    rng: &mut R,
 836) where
 837    <Mmut as Matrix>::R:
 838        RowMut + RowEntity + TryConvertFrom<[S], Parameters = Mmut::MatElement> + Debug,
 839    Mmut::MatElement: Copy + Debug,
 840{
 841    let d = gadget_vector.len();
 842    let q = mod_op.modulus();
 843    let ring_size = s.len();
 844    assert!(out_rgsw.dimension() == (d * 3, ring_size));
 845    assert!(m.as_ref().len() == ring_size);
 846
 847    // RLWE(-sm), RLWE(m)
 848    let (rlwe_dash_nsm, b_rlwe_dash_m) = out_rgsw.split_at_row_mut(d * 2);
 849
 850    let mut s_eval = Mmut::R::try_convert_from(s, &q);
 851    ntt_op.forward(s_eval.as_mut());
 852
 853    let mut scratch_space = Mmut::R::zeros(ring_size);
 854
 855    // RLWE'(-sm)
 856    let (a_rlwe_dash_nsm, b_rlwe_dash_nsm) = rlwe_dash_nsm.split_at_mut(d);
 857    izip!(
 858        a_rlwe_dash_nsm.iter_mut(),
 859        b_rlwe_dash_nsm.iter_mut(),
 860        gadget_vector.iter()
 861    )
 862    .for_each(|(ai, bi, beta_i)| {
 863        // Sample a_i
 864        RandomUniformDist::random_fill(rng, &q, ai.as_mut());
 865
 866        // a_i * s
 867        scratch_space.as_mut().copy_from_slice(ai.as_ref());
 868        ntt_op.forward(scratch_space.as_mut());
 869        mod_op.elwise_mul_mut(scratch_space.as_mut(), s_eval.as_ref());
 870        ntt_op.backward(scratch_space.as_mut());
 871
 872        // b_i = e_i + a_i * s
 873        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
 874        mod_op.elwise_add_mut(bi.as_mut(), scratch_space.as_ref());
 875
 876        // a_i + \beta_i * m
 877        mod_op.elwise_scalar_mul(scratch_space.as_mut(), m.as_ref(), beta_i);
 878        mod_op.elwise_add_mut(ai.as_mut(), scratch_space.as_ref());
 879    });
 880
 881    // RLWE(m)
 882    let mut a_rlwe_dash_m = {
 883        // polynomials of part A of RLWE'(m) are sampled from seed
 884        let mut a = Mmut::zeros(d, ring_size);
 885        a.iter_rows_mut()
 886            .for_each(|ai| RandomUniformDist::random_fill(p_rng, &q, ai.as_mut()));
 887        a
 888    };
 889
 890    izip!(
 891        a_rlwe_dash_m.iter_rows_mut(),
 892        b_rlwe_dash_m.iter_mut(),
 893        gadget_vector.iter()
 894    )
 895    .for_each(|(ai, bi, beta_i)| {
 896        // ai * s
 897        ntt_op.forward(ai.as_mut());
 898        mod_op.elwise_mul_mut(ai.as_mut(), s_eval.as_ref());
 899        ntt_op.backward(ai.as_mut());
 900
 901        // beta_i * m
 902        mod_op.elwise_scalar_mul(scratch_space.as_mut(), m.as_ref(), beta_i);
 903
 904        // Sample e_i
 905        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
 906        // e_i + beta_i * m + ai*s
 907        mod_op.elwise_add_mut(bi.as_mut(), scratch_space.as_ref());
 908        mod_op.elwise_add_mut(bi.as_mut(), ai.as_ref());
 909    });
 910}
 911
 912pub(crate) fn public_key_encrypt_rgsw<
 913    Mmut: MatrixMut + MatrixEntity,
 914    M: Matrix<MatElement = Mmut::MatElement>,
 915    R: RandomGaussianDist<[Mmut::MatElement], Parameters = Mmut::MatElement>
 916        + RandomUniformDist<[u8], Parameters = u8>
 917        + RandomUniformDist<usize, Parameters = usize>,
 918    ModOp: VectorOps<Element = Mmut::MatElement>,
 919    NttOp: Ntt<Element = Mmut::MatElement>,
 920>(
 921    out_rgsw: &mut Mmut,
 922    m: &[M::MatElement],
 923    public_key: &M,
 924    gadget_vector: &[Mmut::MatElement],
 925    mod_op: &ModOp,
 926    ntt_op: &NttOp,
 927    rng: &mut R,
 928) where
 929    <Mmut as Matrix>::R: RowMut + RowEntity + TryConvertFrom<[i32], Parameters = Mmut::MatElement>,
 930    Mmut::MatElement: Copy,
 931{
 932    let ring_size = public_key.dimension().1;
 933    let d = gadget_vector.len();
 934    assert!(public_key.dimension().0 == 2);
 935    assert!(out_rgsw.dimension() == (d * 4, ring_size));
 936
 937    let mut pk_eval = Mmut::zeros(2, ring_size);
 938    izip!(pk_eval.iter_rows_mut(), public_key.iter_rows()).for_each(|(to_i, from_i)| {
 939        to_i.as_mut().copy_from_slice(from_i.as_ref());
 940        ntt_op.forward(to_i.as_mut());
 941    });
 942    let p0 = pk_eval.get_row_slice(0);
 943    let p1 = pk_eval.get_row_slice(1);
 944
 945    let q = mod_op.modulus();
 946
 947    // RGSW(m) = RLWE'(-sm), RLWE(m)
 948    let (rlwe_dash_nsm, rlwe_dash_m) = out_rgsw.split_at_row_mut(2 * d);
 949
 950    // RLWE(-sm)
 951    let (rlwe_dash_nsm_parta, rlwe_dash_nsm_partb) = rlwe_dash_nsm.split_at_mut(d);
 952    izip!(
 953        rlwe_dash_nsm_parta.iter_mut(),
 954        rlwe_dash_nsm_partb.iter_mut(),
 955        gadget_vector.iter()
 956    )
 957    .for_each(|(ai, bi, beta_i)| {
 958        // sample ephemeral secret u_i
 959        let mut u = vec![0i32; ring_size];
 960        fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
 961        let mut u_eval = Mmut::R::try_convert_from(u.as_ref(), &q);
 962        ntt_op.forward(u_eval.as_mut());
 963
 964        let mut u_eval_copy = Mmut::R::zeros(ring_size);
 965        u_eval_copy.as_mut().copy_from_slice(u_eval.as_ref());
 966
 967        // p0 * u
 968        mod_op.elwise_mul_mut(u_eval.as_mut(), p0.as_ref());
 969        // p1 * u
 970        mod_op.elwise_mul_mut(u_eval_copy.as_mut(), p1.as_ref());
 971        ntt_op.backward(u_eval.as_mut());
 972        ntt_op.backward(u_eval_copy.as_mut());
 973
 974        // sample error
 975        RandomGaussianDist::random_fill(rng, &q, ai.as_mut());
 976        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
 977
 978        // a = p0*u+e0
 979        mod_op.elwise_add_mut(ai.as_mut(), u_eval.as_ref());
 980        // b = p1*u+e1
 981        mod_op.elwise_add_mut(bi.as_mut(), u_eval_copy.as_ref());
 982
 983        // a = p0*u + e0 + \beta*m
 984        // use u_eval as scratch
 985        mod_op.elwise_scalar_mul(u_eval.as_mut(), m.as_ref(), beta_i);
 986        mod_op.elwise_add_mut(ai.as_mut(), u_eval.as_ref());
 987    });
 988
 989    // RLWE(m)
 990    let (rlwe_dash_m_parta, rlwe_dash_m_partb) = rlwe_dash_m.split_at_mut(d);
 991    izip!(
 992        rlwe_dash_m_parta.iter_mut(),
 993        rlwe_dash_m_partb.iter_mut(),
 994        gadget_vector.iter()
 995    )
 996    .for_each(|(ai, bi, beta_i)| {
 997        // sample ephemeral secret u_i
 998        let mut u = vec![0i32; ring_size];
 999        fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
1000        let mut u_eval = Mmut::R::try_convert_from(u.as_ref(), &q);
1001        ntt_op.forward(u_eval.as_mut());
1002
1003        let mut u_eval_copy = Mmut::R::zeros(ring_size);
1004        u_eval_copy.as_mut().copy_from_slice(u_eval.as_ref());
1005
1006        // p0 * u
1007        mod_op.elwise_mul_mut(u_eval.as_mut(), p0.as_ref());
1008        // p1 * u
1009        mod_op.elwise_mul_mut(u_eval_copy.as_mut(), p1.as_ref());
1010        ntt_op.backward(u_eval.as_mut());
1011        ntt_op.backward(u_eval_copy.as_mut());
1012
1013        // sample error
1014        RandomGaussianDist::random_fill(rng, &q, ai.as_mut());
1015        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
1016
1017        // a = p0*u+e0
1018        mod_op.elwise_add_mut(ai.as_mut(), u_eval.as_ref());
1019        // b = p1*u+e1
1020        mod_op.elwise_add_mut(bi.as_mut(), u_eval_copy.as_ref());
1021
1022        // b = p1*u + e0 + \beta*m
1023        // use u_eval as scratch
1024        mod_op.elwise_scalar_mul(u_eval.as_mut(), m.as_ref(), beta_i);
1025        mod_op.elwise_add_mut(bi.as_mut(), u_eval.as_ref());
1026    });
1027}
1028
1029/// Generates RLWE Key switching key to key switch ciphertext RLWE_{from_s}(m)
1030/// to RLWE_{to_s}(m).
1031///
1032/// Key switching equals
1033///     \sum decompose(c_1)_i * RLWE_{to_s}(\beta^i -from_s)
1034/// Hence, key switchin key equals RLWE'(-from_s) = RLWE(-from_s), RLWE(beta^1
1035/// -from_s), ..., RLWE(beta^{d-1} -from_s).
1036///
1037/// - ksk_out: Output Key switching key. Key switching key stores only part B
1038///   polynomials of ksk RLWE ciphertexts (i.e. RLWE'_B(-from_s)) in coefficient
1039///   domain
1040/// - neg_from_s: Negative of secret polynomial to key switch from
1041/// - to_s: secret polynomial to key switch to.
1042pub(crate) fn rlwe_ksk_gen<
1043    Mmut: MatrixMut + MatrixEntity,
1044    ModOp: ArithmeticOps<Element = Mmut::MatElement> + VectorOps<Element = Mmut::MatElement>,
1045    NttOp: Ntt<Element = Mmut::MatElement>,
1046    R: RandomGaussianDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
1047    PR: RandomUniformDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
1048>(
1049    ksk_out: &mut Mmut,
1050    neg_from_s: Mmut::R,
1051    mut to_s: Mmut::R,
1052    gadget_vector: &[Mmut::MatElement],
1053    mod_op: &ModOp,
1054    ntt_op: &NttOp,
1055    p_rng: &mut PR,
1056    rng: &mut R,
1057) where
1058    <Mmut as Matrix>::R: RowMut,
1059{
1060    let ring_size = neg_from_s.as_ref().len();
1061    let d = gadget_vector.len();
1062    assert!(ksk_out.dimension() == (d, ring_size));
1063
1064    let q = ArithmeticOps::modulus(mod_op);
1065
1066    ntt_op.forward(to_s.as_mut());
1067
1068    // RLWE'_{to_s}(-from_s)
1069    let mut part_a = {
1070        let mut a = Mmut::zeros(d, ring_size);
1071        a.iter_rows_mut()
1072            .for_each(|ai| RandomUniformDist::random_fill(p_rng, &q, ai.as_mut()));
1073        a
1074    };
1075    izip!(
1076        part_a.iter_rows_mut(),
1077        ksk_out.iter_rows_mut(),
1078        gadget_vector.iter(),
1079    )
1080    .for_each(|(ai, bi, beta_i)| {
1081        // si * ai
1082        ntt_op.forward(ai.as_mut());
1083        mod_op.elwise_mul_mut(ai.as_mut(), to_s.as_ref());
1084        ntt_op.backward(ai.as_mut());
1085
1086        // ei + to_s*ai
1087        RandomGaussianDist::random_fill(rng, &q, bi.as_mut());
1088        mod_op.elwise_add_mut(bi.as_mut(), ai.as_ref());
1089
1090        // beta_i * -from_s
1091        // use ai as scratch space
1092        mod_op.elwise_scalar_mul(ai.as_mut(), neg_from_s.as_ref(), beta_i);
1093
1094        // bi = ei + to_s*ai + beta_i*-from_s
1095        mod_op.elwise_add_mut(bi.as_mut(), ai.as_ref());
1096    });
1097}
1098
1099pub(crate) fn galois_key_gen<
1100    Mmut: MatrixMut + MatrixEntity,
1101    ModOp: ArithmeticOps<Element = Mmut::MatElement> + VectorOps<Element = Mmut::MatElement>,
1102    NttOp: Ntt<Element = Mmut::MatElement>,
1103    S,
1104    R: RandomGaussianDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
1105    PR: RandomUniformDist<[Mmut::MatElement], Parameters = Mmut::MatElement>,
1106>(
1107    ksk_out: &mut Mmut,
1108    s: &[S],
1109    auto_k: isize,
1110    gadget_vector: &[Mmut::MatElement],
1111    mod_op: &ModOp,
1112    ntt_op: &NttOp,
1113    p_rng: &mut PR,
1114    rng: &mut R,
1115) where
1116    <Mmut as Matrix>::R: RowMut,
1117    Mmut::R: TryConvertFrom<[S], Parameters = Mmut::MatElement> + RowEntity,
1118    Mmut::MatElement: Copy + Sub<Output = Mmut::MatElement>,
1119{
1120    let ring_size = s.len();
1121    let (auto_map_index, auto_map_sign) = generate_auto_map(ring_size, auto_k);
1122
1123    let q = ArithmeticOps::modulus(mod_op);
1124
1125    // s(X) -> -s(X^k)
1126    let s = Mmut::R::try_convert_from(s, &q);
1127    let mut neg_s_auto = Mmut::R::zeros(s.as_ref().len());
1128    izip!(s.as_ref(), auto_map_index.iter(), auto_map_sign.iter()).for_each(
1129        |(el, to_index, sign)| {
1130            // if sign is +ve (true), then negate because we need -s(X) (i.e. do the
1131            // opposite than the usual case)
1132            if *sign {
1133                neg_s_auto.as_mut()[*to_index] = q - *el;
1134            } else {
1135                neg_s_auto.as_mut()[*to_index] = *el;
1136            }
1137        },
1138    );
1139
1140    // Ksk from -s(X^k) to s(X)
1141    rlwe_ksk_gen(
1142        ksk_out,
1143        neg_s_auto,
1144        s,
1145        gadget_vector,
1146        mod_op,
1147        ntt_op,
1148        p_rng,
1149        rng,
1150    );
1151}
1152
1153/// Encrypt polynomial m(X) as RLWE ciphertext.
1154///
1155/// - rlwe_out: returned RLWE ciphertext RLWE(m) in coefficient domain. RLWE
1156///   ciphertext is a matirx with first row consiting polynomial `a` and the
1157///   second rows consting polynomial `b`
1158pub(crate) fn secret_key_encrypt_rlwe<
1159    Ro: Row + RowMut + RowEntity,
1160    ModOp: VectorOps<Element = Ro::Element>,
1161    NttOp: Ntt<Element = Ro::Element>,
1162    S,
1163    R: RandomGaussianDist<[Ro::Element], Parameters = Ro::Element>,
1164    PR: RandomUniformDist<[Ro::Element], Parameters = Ro::Element>,
1165>(
1166    m: &[Ro::Element],
1167    b_rlwe_out: &mut Ro,
1168    s: &[S],
1169    mod_op: &ModOp,
1170    ntt_op: &NttOp,
1171    p_rng: &mut PR,
1172    rng: &mut R,
1173) where
1174    Ro: TryConvertFrom<[S], Parameters = Ro::Element> + Debug,
1175{
1176    let ring_size = s.len();
1177    assert!(m.as_ref().len() == ring_size);
1178    assert!(b_rlwe_out.as_ref().len() == ring_size);
1179
1180    let q = mod_op.modulus();
1181
1182    // sample a
1183    let mut a = {
1184        let mut a = Ro::zeros(ring_size);
1185        RandomUniformDist::random_fill(p_rng, &q, a.as_mut());
1186        a
1187    };
1188
1189    // s * a
1190    let mut sa = Ro::try_convert_from(s, &q);
1191    ntt_op.forward(sa.as_mut());
1192    ntt_op.forward(a.as_mut());
1193    mod_op.elwise_mul_mut(sa.as_mut(), a.as_ref());
1194    ntt_op.backward(sa.as_mut());
1195
1196    // sample e
1197    RandomGaussianDist::random_fill(rng, &q, b_rlwe_out.as_mut());
1198    mod_op.elwise_add_mut(b_rlwe_out.as_mut(), m.as_ref());
1199    mod_op.elwise_add_mut(b_rlwe_out.as_mut(), sa.as_ref());
1200}
1201
1202pub(crate) fn public_key_encrypt_rlwe<
1203    M: Matrix,
1204    Mmut: MatrixMut<MatElement = M::MatElement>,
1205    ModOp: VectorOps<Element = M::MatElement>,
1206    NttOp: Ntt<Element = M::MatElement>,
1207    S,
1208    R: RandomGaussianDist<[M::MatElement], Parameters = M::MatElement>
1209        + RandomUniformDist<[M::MatElement], Parameters = M::MatElement>
1210        + RandomUniformDist<[u8], Parameters = u8>
1211        + RandomUniformDist<usize, Parameters = usize>,
1212>(
1213    rlwe_out: &mut Mmut,
1214    pk: &M,
1215    m: &[M::MatElement],
1216    mod_op: &ModOp,
1217    ntt_op: &NttOp,
1218    rng: &mut R,
1219) where
1220    <Mmut as Matrix>::R: RowMut + TryConvertFrom<[S], Parameters = M::MatElement> + RowEntity,
1221    M::MatElement: Copy,
1222    S: Zero + Signed + Copy,
1223{
1224    let ring_size = m.len();
1225    assert!(rlwe_out.dimension() == (2, ring_size));
1226
1227    let q = mod_op.modulus();
1228
1229    let mut u = vec![S::zero(); ring_size];
1230    fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
1231    let mut u = Mmut::R::try_convert_from(&u, &q);
1232    ntt_op.forward(u.as_mut());
1233
1234    let mut ua = Mmut::R::zeros(ring_size);
1235    ua.as_mut().copy_from_slice(pk.get_row_slice(0));
1236    let mut ub = Mmut::R::zeros(ring_size);
1237    ub.as_mut().copy_from_slice(pk.get_row_slice(1));
1238
1239    // a*u
1240    ntt_op.forward(ua.as_mut());
1241    mod_op.elwise_mul_mut(ua.as_mut(), u.as_ref());
1242    ntt_op.backward(ua.as_mut());
1243
1244    // b*u
1245    ntt_op.forward(ub.as_mut());
1246    mod_op.elwise_mul_mut(ub.as_mut(), u.as_ref());
1247    ntt_op.backward(ub.as_mut());
1248
1249    // sample error
1250    rlwe_out.iter_rows_mut().for_each(|ri| {
1251        RandomGaussianDist::random_fill(rng, &q, ri.as_mut());
1252    });
1253
1254    // a*u + e0
1255    mod_op.elwise_add_mut(rlwe_out.get_row_mut(0), ua.as_ref());
1256    // b*u + e1
1257    mod_op.elwise_add_mut(rlwe_out.get_row_mut(1), ub.as_ref());
1258
1259    // b*u + e1 + m
1260    mod_op.elwise_add_mut(rlwe_out.get_row_mut(1), m);
1261}
1262
1263/// Generates RLWE public key
1264pub(crate) fn gen_rlwe_public_key<
1265    Ro: RowMut + RowEntity,
1266    S,
1267    ModOp: VectorOps<Element = Ro::Element>,
1268    NttOp: Ntt<Element = Ro::Element>,
1269    PRng: RandomUniformDist<[Ro::Element], Parameters = Ro::Element>,
1270    Rng: RandomGaussianDist<[Ro::Element], Parameters = Ro::Element>,
1271>(
1272    part_b_out: &mut Ro,
1273    s: &[S],
1274    ntt_op: &NttOp,
1275    mod_op: &ModOp,
1276    p_rng: &mut PRng,
1277    rng: &mut Rng,
1278) where
1279    Ro: TryConvertFrom<[S], Parameters = Ro::Element>,
1280{
1281    let ring_size = s.len();
1282    assert!(part_b_out.as_ref().len() == ring_size);
1283
1284    let q = mod_op.modulus();
1285
1286    // sample a
1287    let mut a = {
1288        let mut tmp = Ro::zeros(ring_size);
1289        RandomUniformDist::random_fill(p_rng, &q, tmp.as_mut());
1290        tmp
1291    };
1292    ntt_op.forward(a.as_mut());
1293
1294    // s*a
1295    let mut sa = Ro::try_convert_from(s, &q);
1296    ntt_op.forward(sa.as_mut());
1297    mod_op.elwise_mul_mut(sa.as_mut(), a.as_ref());
1298    ntt_op.backward(sa.as_mut());
1299
1300    // s*a + e
1301    RandomGaussianDist::random_fill(rng, &q, part_b_out.as_mut());
1302    mod_op.elwise_add_mut(part_b_out.as_mut(), sa.as_ref());
1303}
1304
1305/// Decrypts degree 1 RLWE ciphertext RLWE(m) and returns m
1306///
1307/// - rlwe_ct: input degree 1 ciphertext RLWE(m).
1308pub(crate) fn decrypt_rlwe<
1309    R: RowMut,
1310    M: Matrix<MatElement = R::Element>,
1311    ModOp: VectorOps<Element = R::Element>,
1312    NttOp: Ntt<Element = R::Element>,
1313    S,
1314>(
1315    rlwe_ct: &M,
1316    s: &[S],
1317    m_out: &mut R,
1318    ntt_op: &NttOp,
1319    mod_op: &ModOp,
1320) where
1321    R: TryConvertFrom<[S], Parameters = R::Element>,
1322    R::Element: Copy,
1323{
1324    let ring_size = s.len();
1325    assert!(rlwe_ct.dimension() == (2, ring_size));
1326    assert!(m_out.as_ref().len() == ring_size);
1327
1328    // transform a to evluation form
1329    m_out.as_mut().copy_from_slice(rlwe_ct.get_row_slice(0));
1330    ntt_op.forward(m_out.as_mut());
1331
1332    // -s*a
1333    let mut s = R::try_convert_from(&s, &mod_op.modulus());
1334    ntt_op.forward(s.as_mut());
1335    mod_op.elwise_mul_mut(m_out.as_mut(), s.as_ref());
1336    mod_op.elwise_neg_mut(m_out.as_mut());
1337    ntt_op.backward(m_out.as_mut());
1338
1339    // m+e = b - s*a
1340    mod_op.elwise_add_mut(m_out.as_mut(), rlwe_ct.get_row_slice(1));
1341}
1342
1343// Measures noise in degree 1 RLWE ciphertext against encoded ideal message
1344// encoded_m
1345pub(crate) fn measure_noise<
1346    Mmut: MatrixMut + Matrix,
1347    ModOp: VectorOps<Element = Mmut::MatElement>,
1348    NttOp: Ntt<Element = Mmut::MatElement>,
1349    S,
1350>(
1351    rlwe_ct: &Mmut,
1352    encoded_m_ideal: &Mmut::R,
1353    ntt_op: &NttOp,
1354    mod_op: &ModOp,
1355    s: &[S],
1356) -> f64
1357where
1358    <Mmut as Matrix>::R: RowMut,
1359    Mmut::R: RowEntity + TryConvertFrom<[S], Parameters = Mmut::MatElement>,
1360    Mmut::MatElement: PrimInt + ToPrimitive + Debug,
1361{
1362    let ring_size = s.len();
1363    assert!(rlwe_ct.dimension() == (2, ring_size));
1364    assert!(encoded_m_ideal.as_ref().len() == ring_size);
1365
1366    // -(s * a)
1367    let q = VectorOps::modulus(mod_op);
1368    let mut s = Mmut::R::try_convert_from(s, &q);
1369    ntt_op.forward(s.as_mut());
1370    let mut a = Mmut::R::zeros(ring_size);
1371    a.as_mut().copy_from_slice(rlwe_ct.get_row_slice(0));
1372    ntt_op.forward(a.as_mut());
1373    mod_op.elwise_mul_mut(s.as_mut(), a.as_ref());
1374    mod_op.elwise_neg_mut(s.as_mut());
1375    ntt_op.backward(s.as_mut());
1376
1377    // m+e = b - s*a
1378    let mut m_plus_e = s;
1379    mod_op.elwise_add_mut(m_plus_e.as_mut(), rlwe_ct.get_row_slice(1));
1380
1381    // difference
1382    mod_op.elwise_sub_mut(m_plus_e.as_mut(), encoded_m_ideal.as_ref());
1383
1384    let mut max_diff_bits = f64::MIN;
1385    m_plus_e.as_ref().iter().for_each(|v| {
1386        let mut v = *v;
1387        if v >= (q >> 1) {
1388            // v is -ve
1389            v = q - v;
1390        }
1391
1392        let bits = (v.to_f64().unwrap()).log2();
1393
1394        if max_diff_bits < bits {
1395            max_diff_bits = bits;
1396        }
1397    });
1398
1399    return max_diff_bits;
1400}
1401
1402#[cfg(test)]
1403pub(crate) mod tests {
1404    use std::{ops::Mul, vec};
1405
1406    use itertools::{izip, Itertools};
1407    use rand::{thread_rng, Rng};
1408
1409    use crate::{
1410        backend::{ModInit, ModularOpsU64, VectorOps},
1411        decomposer::{gadget_vector, DefaultDecomposer},
1412        ntt::{self, Ntt, NttBackendU64, NttInit},
1413        random::{DefaultSecureRng, NewWithSeed, RandomUniformDist},
1414        rgsw::{
1415            gen_rlwe_public_key, measure_noise, public_key_encrypt_rgsw, AutoKeyEvaluationDomain,
1416            RgswCiphertext, RgswCiphertextEvaluationDomain, RlweCiphertext, RlwePublicKey,
1417            SeededAutoKey, SeededRgswCiphertext, SeededRlweCiphertext, SeededRlwePublicKey,
1418        },
1419        utils::{generate_prime, negacyclic_mul, TryConvertFrom},
1420        Matrix, Secret,
1421    };
1422
1423    use super::{
1424        decrypt_rlwe, galois_auto, galois_key_gen, generate_auto_map, public_key_encrypt_rlwe,
1425        rgsw_by_rgsw_inplace, rlwe_by_rgsw, secret_key_encrypt_rgsw, secret_key_encrypt_rlwe,
1426        RlweSecret,
1427    };
1428
1429    #[test]
1430    fn rlwe_encrypt_decryption() {
1431        let logq = 50;
1432        let logp = 2;
1433        let ring_size = 1 << 4;
1434        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
1435        let p = 1u64 << logp;
1436
1437        let mut rng = DefaultSecureRng::new();
1438
1439        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1440
1441        // sample m0
1442        let mut m0 = vec![0u64; ring_size as usize];
1443        RandomUniformDist::<[u64]>::random_fill(&mut rng, &(1u64 << logp), m0.as_mut_slice());
1444
1445        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1446        let mod_op = ModularOpsU64::new(q);
1447
1448        // encrypt m0
1449        let mut rlwe_seed = [0u8; 32];
1450        rng.fill_bytes(&mut rlwe_seed);
1451        let mut seeded_rlwe_in_ct =
1452            SeededRlweCiphertext::<_, [u8; 32]>::empty(ring_size as usize, rlwe_seed, q);
1453        let mut p_rng = DefaultSecureRng::new_with_seed(rlwe_seed);
1454        let encoded_m = m0
1455            .iter()
1456            .map(|v| (((*v as f64) * q as f64) / (p as f64)).round() as u64)
1457            .collect_vec();
1458        secret_key_encrypt_rlwe(
1459            &encoded_m,
1460            &mut seeded_rlwe_in_ct.data,
1461            s.values(),
1462            &mod_op,
1463            &ntt_op,
1464            &mut p_rng,
1465            &mut rng,
1466        );
1467        let rlwe_in_ct =
1468            RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_rlwe_in_ct);
1469
1470        let mut encoded_m_back = vec![0u64; ring_size as usize];
1471        decrypt_rlwe(
1472            &rlwe_in_ct,
1473            s.values(),
1474            &mut encoded_m_back,
1475            &ntt_op,
1476            &mod_op,
1477        );
1478        let m_back = encoded_m_back
1479            .iter()
1480            .map(|v| (((*v as f64 * p as f64) / q as f64).round() as u64) % p)
1481            .collect_vec();
1482        assert_eq!(m0, m_back);
1483
1484        let noise = measure_noise(&rlwe_in_ct, &encoded_m, &ntt_op, &mod_op, s.values());
1485        println!("Noise: {noise}");
1486    }
1487
1488    #[test]
1489    fn rlwe_by_rgsw_works() {
1490        let logq = 50;
1491        let logp = 2;
1492        let ring_size = 1 << 9;
1493        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
1494        let p = 1u64 << logp;
1495        let d_rgsw = 10;
1496        let logb = 5;
1497
1498        let mut rng = DefaultSecureRng::new_seeded([0u8; 32]);
1499
1500        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1501
1502        let mut m0 = vec![0u64; ring_size as usize];
1503        RandomUniformDist::<[u64]>::random_fill(&mut rng, &(1u64 << logp), m0.as_mut_slice());
1504        let mut m1 = vec![0u64; ring_size as usize];
1505        m1[thread_rng().gen_range(0..ring_size) as usize] = 1;
1506
1507        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1508        let mod_op = ModularOpsU64::new(q);
1509        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
1510
1511        // Encrypt m1 as RGSW(m1)
1512        let rgsw_ct = {
1513            //TODO(Jay): Figure out better way to test secret key and public key variant of
1514            // RGSW ciphertext encryption within the same test
1515
1516            if true {
1517                // Encryption m1 as RGSW(m1) using secret key
1518                _sk_encrypt_rgsw(&m1, s.values(), &gadget_vector, &mod_op, &ntt_op)
1519            } else {
1520                // Encrypt m1 as RGSW(m1) as public key
1521
1522                // first create public key
1523                let mut pk_seed = [0u8; 32];
1524                rng.fill_bytes(&mut pk_seed);
1525                let mut pk_prng = DefaultSecureRng::new_seeded(pk_seed);
1526                let mut seeded_pk =
1527                    SeededRlwePublicKey::<Vec<u64>, _>::empty(ring_size as usize, pk_seed, q);
1528                gen_rlwe_public_key(
1529                    &mut seeded_pk.data,
1530                    s.values(),
1531                    &ntt_op,
1532                    &mod_op,
1533                    &mut pk_prng,
1534                    &mut rng,
1535                );
1536                let pk = RlwePublicKey::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_pk);
1537
1538                let rgsw_ct = _pk_encrypt_rgsw(&m1, &pk, &gadget_vector, &mod_op, &ntt_op);
1539                RgswCiphertextEvaluationDomain::<_, DefaultSecureRng, NttBackendU64>::from(
1540                    &RgswCiphertext {
1541                        data: rgsw_ct.data,
1542                        modulus: q,
1543                    },
1544                )
1545            }
1546        };
1547
1548        // Encrypt m0 as RLWE(m0)
1549        let mut rlwe_in_ct = {
1550            let mut rlwe_seed = [0u8; 32];
1551            rng.fill_bytes(&mut rlwe_seed);
1552            let mut seeded_rlwe_in_ct =
1553                SeededRlweCiphertext::<_, [u8; 32]>::empty(ring_size as usize, rlwe_seed, q);
1554            let mut p_rng = DefaultSecureRng::new_seeded(rlwe_seed);
1555            let encoded_m = m0
1556                .iter()
1557                .map(|v| (((*v as f64) * q as f64) / (p as f64)).round() as u64)
1558                .collect_vec();
1559            secret_key_encrypt_rlwe(
1560                &encoded_m,
1561                &mut seeded_rlwe_in_ct.data,
1562                s.values(),
1563                &mod_op,
1564                &ntt_op,
1565                &mut p_rng,
1566                &mut rng,
1567            );
1568
1569            RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_rlwe_in_ct)
1570        };
1571
1572        // RLWE(m0m1) = RLWE(m0) x RGSW(m1)
1573        let mut scratch_space = vec![vec![0u64; ring_size as usize]; d_rgsw + 2];
1574        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
1575        rlwe_by_rgsw(
1576            &mut rlwe_in_ct,
1577            &rgsw_ct.data,
1578            &mut scratch_space,
1579            &decomposer,
1580            &ntt_op,
1581            &mod_op,
1582        );
1583
1584        // Decrypt RLWE(m0m1)
1585        let mut encoded_m0m1_back = vec![0u64; ring_size as usize];
1586        decrypt_rlwe(
1587            &rlwe_in_ct,
1588            s.values(),
1589            &mut encoded_m0m1_back,
1590            &ntt_op,
1591            &mod_op,
1592        );
1593        let m0m1_back = encoded_m0m1_back
1594            .iter()
1595            .map(|v| (((*v as f64 * p as f64) / (q as f64)).round() as u64) % p)
1596            .collect_vec();
1597
1598        let mul_mod = |v0: &u64, v1: &u64| (v0 * v1) % p;
1599        let m0m1 = negacyclic_mul(&m0, &m1, mul_mod, p);
1600
1601        {
1602            // measure noise
1603            let encoded_m_ideal = m0m1
1604                .iter()
1605                .map(|v| (((*v as f64) * q as f64) / (p as f64)).round() as u64)
1606                .collect_vec();
1607
1608            let noise = measure_noise(&rlwe_in_ct, &encoded_m_ideal, &ntt_op, &mod_op, s.values());
1609            println!("Noise RLWE(m0m1)(= RLWE(m0)xRGSW(m1)) : {noise}");
1610        }
1611
1612        assert!(
1613            m0m1 == m0m1_back,
1614            "Expected {:?} \n Got {:?}",
1615            m0m1,
1616            m0m1_back
1617        );
1618    }
1619
1620    fn _secret_encrypt_rlwe(
1621        m: &[u64],
1622        s: &[i32],
1623        ntt_op: &NttBackendU64,
1624        mod_op: &ModularOpsU64,
1625    ) -> RlweCiphertext<Vec<Vec<u64>>, DefaultSecureRng> {
1626        let ring_size = m.len();
1627        let q = mod_op.modulus();
1628        assert!(s.len() == ring_size);
1629
1630        let mut rng = DefaultSecureRng::new();
1631        let mut rlwe_seed = [0u8; 32];
1632        rng.fill_bytes(&mut rlwe_seed);
1633        let mut seeded_rlwe_ct =
1634            SeededRlweCiphertext::<_, [u8; 32]>::empty(ring_size as usize, rlwe_seed, q);
1635        let mut p_rng = DefaultSecureRng::new_seeded(rlwe_seed);
1636        secret_key_encrypt_rlwe(
1637            &m,
1638            &mut seeded_rlwe_ct.data,
1639            s,
1640            mod_op,
1641            ntt_op,
1642            &mut p_rng,
1643            &mut rng,
1644        );
1645
1646        RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_rlwe_ct)
1647    }
1648
1649    #[test]
1650    fn rlwe_by_rgsw_noise_growth() {
1651        let logq = 28;
1652        let ring_size = 1 << 10;
1653        let q = generate_prime(logq, ring_size * 2, 1u64 << logq).unwrap();
1654        let d_rgsw = 2;
1655        let logb = 7;
1656
1657        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1658
1659        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1660        let mod_op = ModularOpsU64::new(q);
1661        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
1662        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
1663
1664        let mul_mod = |v0: &u64, v1: &u64| ((*v0 as u128 * *v1 as u128) % (q as u128)) as u64;
1665
1666        let mut carry_m = vec![0u64; ring_size as usize];
1667        carry_m[thread_rng().gen_range(0..ring_size) as usize] = 1;
1668        let mut rlwe = _secret_encrypt_rlwe(&carry_m, s.values(), &ntt_op, &mod_op);
1669
1670        let mut scratch_matrix_dplus2_ring = vec![vec![0u64; ring_size as usize]; d_rgsw + 2];
1671        for i in 0..1000usize {
1672            // Encrypt monomial as RGSW
1673            let mut m = vec![0u64; ring_size as usize];
1674            m[thread_rng().gen_range(0..ring_size) as usize] = if i & 1 == 1 { 1 } else { q - 1 };
1675            let rgsw_ct = _sk_encrypt_rgsw(&m, s.values(), &gadget_vector, &mod_op, &ntt_op);
1676
1677            // RLWE(carry_m * m) = RLWE(carry_m) x RGSW(m)
1678            rlwe_by_rgsw(
1679                &mut rlwe,
1680                &rgsw_ct.data,
1681                &mut scratch_matrix_dplus2_ring,
1682                &decomposer,
1683                &ntt_op,
1684                &mod_op,
1685            );
1686
1687            carry_m = negacyclic_mul(&carry_m, &m, mul_mod, q);
1688            let noise = measure_noise(&rlwe, &carry_m, &ntt_op, &mod_op, s.values());
1689
1690            println!("Noise RLWE(carry_m) after {i}^th iteration: {noise}");
1691        }
1692    }
1693
1694    // Encrypt m as RGSW ciphertext RGSW(m) using supplied public key
1695    fn _pk_encrypt_rgsw(
1696        m: &[u64],
1697        public_key: &RlwePublicKey<Vec<Vec<u64>>, DefaultSecureRng>,
1698        gadget_vector: &[u64],
1699        mod_op: &ModularOpsU64,
1700        ntt_op: &NttBackendU64,
1701    ) -> RgswCiphertext<Vec<Vec<u64>>> {
1702        let (_, ring_size) = Matrix::dimension(&public_key.data);
1703        let d_rgsw = gadget_vector.len();
1704
1705        let mut rng = DefaultSecureRng::new();
1706
1707        assert!(m.len() == ring_size);
1708
1709        // public key encrypt RGSW(m1)
1710        let mut rgsw_ct = vec![vec![0u64; ring_size]; d_rgsw * 4];
1711        public_key_encrypt_rgsw(
1712            &mut rgsw_ct,
1713            m,
1714            &public_key.data,
1715            gadget_vector,
1716            mod_op,
1717            ntt_op,
1718            &mut rng,
1719        );
1720
1721        RgswCiphertext {
1722            data: rgsw_ct,
1723            modulus: mod_op.modulus(),
1724        }
1725    }
1726
1727    /// Encrypts m as RGSW ciphertext RGSW(m) using supplied secret key. Returns
1728    /// unseeded RGSW ciphertext in coefficient domain
1729    fn _sk_encrypt_rgsw(
1730        m: &[u64],
1731        s: &[i32],
1732        gadget_vector: &[u64],
1733        mod_op: &ModularOpsU64,
1734        ntt_op: &NttBackendU64,
1735    ) -> RgswCiphertextEvaluationDomain<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64> {
1736        let ring_size = s.len();
1737        assert!(m.len() == s.len());
1738
1739        let d_rgsw = gadget_vector.len();
1740        let q = mod_op.modulus();
1741
1742        let mut rng = DefaultSecureRng::new();
1743        let mut rgsw_seed = [0u8; 32];
1744        rng.fill_bytes(&mut rgsw_seed);
1745        let mut seeded_rgsw_ct = SeededRgswCiphertext::<Vec<Vec<u64>>, [u8; 32]>::empty(
1746            ring_size as usize,
1747            d_rgsw,
1748            rgsw_seed,
1749            q,
1750        );
1751        let mut p_rng = DefaultSecureRng::new_seeded(rgsw_seed);
1752        secret_key_encrypt_rgsw(
1753            &mut seeded_rgsw_ct.data,
1754            m,
1755            &gadget_vector,
1756            s,
1757            mod_op,
1758            ntt_op,
1759            &mut p_rng,
1760            &mut rng,
1761        );
1762
1763        RgswCiphertextEvaluationDomain::<_, DefaultSecureRng, NttBackendU64>::from(&seeded_rgsw_ct)
1764    }
1765
1766    /// Prints noise in RGSW ciphertext RGSW(m).
1767    ///
1768    /// - rgsw_ct: RGSW ciphertext in coefficient domain
1769    pub(crate) fn _measure_noise_rgsw(
1770        rgsw_ct: &[Vec<u64>],
1771        m: &[u64],
1772        s: &[i32],
1773        gadget_vector: &[u64],
1774        q: u64,
1775    ) {
1776        let d_rgsw = gadget_vector.len();
1777        let ring_size = s.len();
1778        assert!(Matrix::dimension(&rgsw_ct) == (d_rgsw * 2 * 2, ring_size));
1779        assert!(m.len() == ring_size);
1780
1781        let mod_op = ModularOpsU64::new(q);
1782        let ntt_op = NttBackendU64::new(q, ring_size);
1783
1784        let mul_mod = |a: &u64, b: &u64| ((*a as u128 * *b as u128) % q as u128) as u64;
1785        let s_poly = Vec::<u64>::try_convert_from(s, &q);
1786        let mut neg_s = s_poly.clone();
1787        mod_op.elwise_neg_mut(neg_s.as_mut());
1788        let neg_sm0m1 = negacyclic_mul(&neg_s, &m, mul_mod, q);
1789        for i in 0..2 {
1790            for j in 0..d_rgsw {
1791                let ideal_m = {
1792                    if i == 0 {
1793                        // RLWE(\beta^j -s * m)
1794                        let mut beta_neg_sm0m1 = vec![0u64; ring_size as usize];
1795                        mod_op.elwise_scalar_mul(
1796                            beta_neg_sm0m1.as_mut(),
1797                            &neg_sm0m1,
1798                            &gadget_vector[j],
1799                        );
1800                        beta_neg_sm0m1
1801                    } else {
1802                        // RLWE(\beta^j  m)
1803                        let mut beta_m0m1 = vec![0u64; ring_size as usize];
1804                        mod_op.elwise_scalar_mul(beta_m0m1.as_mut(), &m, &gadget_vector[j]);
1805                        beta_m0m1
1806                    }
1807                };
1808
1809                let mut rlwe = vec![vec![0u64; ring_size as usize]; 2];
1810                rlwe[0].copy_from_slice(rgsw_ct.get_row_slice((i * 2 * d_rgsw) + j));
1811                rlwe[1].copy_from_slice(rgsw_ct.get_row_slice((i * 2 * d_rgsw) + d_rgsw + j));
1812                let noise = measure_noise(&rlwe, &ideal_m, &ntt_op, &mod_op, s);
1813
1814                if i == 0 {
1815                    println!(r"Noise RLWE(\beta^{j} -sm0m1): {noise}");
1816                } else {
1817                    println!(r"Noise RLWE(\beta^{j} m0m1): {noise}");
1818                }
1819            }
1820            // m0m1
1821        }
1822    }
1823
1824    #[test]
1825    fn pk_rgsw_by_rgsw() {
1826        let logq = 60;
1827        let logp = 2;
1828        let ring_size = 1 << 11;
1829        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
1830        let p = 1u64 << logp;
1831        let d_rgsw = 3;
1832        let logb = 15;
1833
1834        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1835
1836        let mut rng = DefaultSecureRng::new();
1837        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1838        let mod_op = ModularOpsU64::new(q);
1839        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
1840        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
1841
1842        let mul_mod = |a: &u64, b: &u64| ((*a as u128 * *b as u128) % q as u128) as u64;
1843
1844        // Public Key
1845        let public_key = {
1846            let mut pk_seed = [0u8; 32];
1847            rng.fill_bytes(&mut pk_seed);
1848            let mut pk_prng = DefaultSecureRng::new_seeded(pk_seed);
1849            let mut seeded_pk =
1850                SeededRlwePublicKey::<Vec<u64>, _>::empty(ring_size as usize, pk_seed, q);
1851            gen_rlwe_public_key(
1852                &mut seeded_pk.data,
1853                s.values(),
1854                &ntt_op,
1855                &mod_op,
1856                &mut pk_prng,
1857                &mut rng,
1858            );
1859            RlwePublicKey::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_pk)
1860        };
1861
1862        let mut carry_m = vec![0u64; ring_size as usize];
1863        carry_m[thread_rng().gen_range(0..ring_size) as usize] = 1;
1864
1865        // RGSW(carry_m)
1866        let mut rgsw_carrym =
1867            _pk_encrypt_rgsw(&carry_m, &public_key, &gadget_vector, &mod_op, &ntt_op);
1868        // let mut rgsw_carrym = {
1869        //     let mut rgsw_eval =
1870        //         _sk_encrypt_rgsw(&carry_m, s.values(), &gadget_vector, &mod_op,
1871        // &ntt_op);     rgsw_eval
1872        //         .data
1873        //         .iter_mut()
1874        //         .for_each(|ri| ntt_op.backward(ri.as_mut()));
1875        //     rgsw_eval.data
1876        // };
1877
1878        println!("########### Noise RGSW(carrym) at start ###########");
1879        _measure_noise_rgsw(&rgsw_carrym.data, &carry_m, s.values(), &gadget_vector, q);
1880
1881        let mut scratch_matrix_d_plus_rgsw_by_ring =
1882            vec![vec![0u64; ring_size as usize]; d_rgsw + (d_rgsw * 4)];
1883
1884        for i in 0..10 {
1885            let mut m = vec![0u64; ring_size as usize];
1886            m[thread_rng().gen_range(0..ring_size) as usize] = q - 1;
1887            let rgsw_m = {
1888                RgswCiphertextEvaluationDomain::<_, DefaultSecureRng, NttBackendU64>::from(
1889                    &_pk_encrypt_rgsw(&m, &public_key, &gadget_vector, &mod_op, &ntt_op),
1890                )
1891            };
1892
1893            rgsw_by_rgsw_inplace(
1894                &mut rgsw_carrym.data,
1895                &rgsw_m.data,
1896                &decomposer,
1897                &mut scratch_matrix_d_plus_rgsw_by_ring,
1898                &ntt_op,
1899                &mod_op,
1900            );
1901
1902            // measure noise
1903            carry_m = negacyclic_mul(&carry_m, &m, mul_mod, q);
1904            println!("########### Noise RGSW(carrym) in {i}^th loop ###########");
1905            _measure_noise_rgsw(&rgsw_carrym.data, &carry_m, s.values(), &gadget_vector, q);
1906        }
1907
1908        // {
1909        //     // RLWE(m) x RGSW(carry_m)
1910        //     let mut m = vec![0u64; ring_size as usize];
1911        //     RandomUniformDist::random_fill(&mut rng, &q, m.as_mut_slice());
1912        //     let mut rlwe_ct = RlweCiphertext::<_,
1913        // DefaultSecureRng>::from_raw(         vec![vec![0u64;
1914        // ring_size as usize]; 2],         false,
1915        //     );
1916        //     let mut scratch_matrix_dplus2_ring = vec![vec![0u64; ring_size as
1917        // usize]; d_rgsw + 2];     public_key_encrypt_rlwe(
1918        //         &mut rlwe_ct,
1919        //         &public_key.data,
1920        //         &m,
1921        //         &mod_op,
1922        //         &ntt_op,
1923        //         &mut rng,
1924        //     );
1925        //     rlwe_by_rgsw(
1926        //         &mut rlwe_ct,
1927        //         &RgswCiphertextEvaluationDomain::<_, DefaultSecureRng,
1928        // NttBackendU64>::from(             &rgsw_carrym,
1929        //         )
1930        //         .data,
1931        //         &mut scratch_matrix_dplus2_ring,
1932        //         &decomposer,
1933        //         &ntt_op,
1934        //         &mod_op,
1935        //     );
1936        //     let m_expected = negacyclic_mul(&carry_m, &m, mul_mod, q);
1937        //     let noise = measure_noise(&rlwe_ct, &m_expected, &ntt_op,
1938        // &mod_op, s.values());     println!("RLWE(m) x RGSW(carry_m):
1939        // {noise}"); }
1940    }
1941
1942    #[test]
1943    fn sk_rgsw_by_rgsw() {
1944        let logq = 60;
1945        let logp = 2;
1946        let ring_size = 1 << 11;
1947        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
1948        let p = 1u64 << logp;
1949        let d_rgsw = 3;
1950        let logb = 15;
1951
1952        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
1953
1954        let mut rng = DefaultSecureRng::new();
1955        let ntt_op = NttBackendU64::new(q, ring_size as usize);
1956        let mod_op = ModularOpsU64::new(q);
1957        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
1958        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
1959        let mul_mod = |a: &u64, b: &u64| ((*a as u128 * *b as u128) % q as u128) as u64;
1960
1961        let mut carry_m = vec![0u64; ring_size as usize];
1962        carry_m[thread_rng().gen_range(0..ring_size) as usize] = 1;
1963
1964        // RGSW(carry_m)
1965        let mut rgsw_carrym = {
1966            let mut rgsw_eval =
1967                _sk_encrypt_rgsw(&carry_m, s.values(), &gadget_vector, &mod_op, &ntt_op);
1968            rgsw_eval
1969                .data
1970                .iter_mut()
1971                .for_each(|ri| ntt_op.backward(ri.as_mut()));
1972            rgsw_eval.data
1973        };
1974        println!("########### Noise RGSW(carrym) at start ###########");
1975        _measure_noise_rgsw(&rgsw_carrym, &carry_m, s.values(), &gadget_vector, q);
1976
1977        let mut scratch_matrix_d_plus_rgsw_by_ring =
1978            vec![vec![0u64; ring_size as usize]; d_rgsw + (d_rgsw * 4)];
1979
1980        for i in 0..10 {
1981            let mut m = vec![0u64; ring_size as usize];
1982            m[thread_rng().gen_range(0..ring_size) as usize] = if (i & 1) == 1 { q - 1 } else { 1 };
1983            let rgsw_m = _sk_encrypt_rgsw(&m, s.values(), &gadget_vector, &mod_op, &ntt_op);
1984            rgsw_by_rgsw_inplace(
1985                &mut rgsw_carrym,
1986                &rgsw_m.data,
1987                &decomposer,
1988                &mut scratch_matrix_d_plus_rgsw_by_ring,
1989                &ntt_op,
1990                &mod_op,
1991            );
1992
1993            // measure noise
1994            carry_m = negacyclic_mul(&carry_m, &m, mul_mod, q);
1995            println!("########### Noise RGSW(carrym) in {i}^th loop ###########");
1996            _measure_noise_rgsw(&rgsw_carrym, &carry_m, s.values(), &gadget_vector, q);
1997        }
1998    }
1999
2000    #[test]
2001    fn galois_auto_works() {
2002        let logq = 50;
2003        let ring_size = 1 << 4;
2004        let q = generate_prime(logq, 2 * ring_size, 1u64 << logq).unwrap();
2005        let logp = 3;
2006        let p = 1u64 << logp;
2007        let d_rgsw = 10;
2008        let logb = 5;
2009
2010        let mut rng = DefaultSecureRng::new();
2011        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
2012
2013        let mut m = vec![0u64; ring_size as usize];
2014        RandomUniformDist::random_fill(&mut rng, &p, m.as_mut_slice());
2015        let encoded_m = m
2016            .iter()
2017            .map(|v| (((*v as f64 * q as f64) / (p as f64)).round() as u64))
2018            .collect_vec();
2019
2020        let ntt_op = NttBackendU64::new(q, ring_size as usize);
2021        let mod_op = ModularOpsU64::new(q);
2022
2023        // RLWE_{s}(m)
2024        let mut seed_rlwe = [0u8; 32];
2025        rng.fill_bytes(&mut seed_rlwe);
2026        let mut seeded_rlwe_m = SeededRlweCiphertext::empty(ring_size as usize, seed_rlwe, q);
2027        let mut p_rng = DefaultSecureRng::new_seeded(seed_rlwe);
2028        secret_key_encrypt_rlwe(
2029            &encoded_m,
2030            &mut seeded_rlwe_m.data,
2031            s.values(),
2032            &mod_op,
2033            &ntt_op,
2034            &mut p_rng,
2035            &mut rng,
2036        );
2037        let mut rlwe_m = RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_rlwe_m);
2038
2039        let auto_k = -5;
2040
2041        // Generate galois key to key switch from s^k to s
2042        let mut seed_auto = [0u8; 32];
2043        rng.fill_bytes(&mut seed_auto);
2044        let mut seeded_auto_key = SeededAutoKey::empty(ring_size as usize, d_rgsw, seed_auto, q);
2045        let mut p_rng = DefaultSecureRng::new_seeded(seed_auto);
2046        let gadget_vector = gadget_vector(logq, logb, d_rgsw);
2047        galois_key_gen(
2048            &mut seeded_auto_key.data,
2049            s.values(),
2050            auto_k,
2051            &gadget_vector,
2052            &mod_op,
2053            &ntt_op,
2054            &mut p_rng,
2055            &mut rng,
2056        );
2057        let auto_key =
2058            AutoKeyEvaluationDomain::<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64>::from(
2059                &seeded_auto_key,
2060            );
2061
2062        // Send RLWE_{s}(m) -> RLWE_{s}(m^k)
2063        let mut scratch_space = vec![vec![0u64; ring_size as usize]; d_rgsw + 2];
2064        let (auto_map_index, auto_map_sign) = generate_auto_map(ring_size as usize, auto_k);
2065        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
2066        galois_auto(
2067            &mut rlwe_m,
2068            &auto_key.data,
2069            &mut scratch_space,
2070            &auto_map_index,
2071            &auto_map_sign,
2072            &mod_op,
2073            &ntt_op,
2074            &decomposer,
2075        );
2076
2077        let rlwe_m_k = rlwe_m;
2078
2079        // Decrypt RLWE_{s}(m^k) and check
2080        let mut encoded_m_k_back = vec![0u64; ring_size as usize];
2081        decrypt_rlwe(
2082            &rlwe_m_k,
2083            s.values(),
2084            &mut encoded_m_k_back,
2085            &ntt_op,
2086            &mod_op,
2087        );
2088        let m_k_back = encoded_m_k_back
2089            .iter()
2090            .map(|v| (((*v as f64 * p as f64) / q as f64).round() as u64) % p)
2091            .collect_vec();
2092
2093        let mut m_k = vec![0u64; ring_size as usize];
2094        // Send \delta m -> \delta m^k
2095        izip!(m.iter(), auto_map_index.iter(), auto_map_sign.iter()).for_each(
2096            |(v, to_index, sign)| {
2097                if !*sign {
2098                    m_k[*to_index] = (p - *v) % p;
2099                } else {
2100                    m_k[*to_index] = *v;
2101                }
2102            },
2103        );
2104
2105        {
2106            let encoded_m_k = m_k
2107                .iter()
2108                .map(|v| ((*v as f64 * q as f64) / p as f64).round() as u64)
2109                .collect_vec();
2110
2111            let noise = measure_noise(&rlwe_m_k, &encoded_m_k, &ntt_op, &mod_op, s.values());
2112            println!("Ksk noise: {noise}");
2113        }
2114
2115        assert_eq!(m_k_back, m_k);
2116    }
2117}