arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/rgsw/mod.rs
src/rgsw/mod.rsRaw
  1mod keygen;
  2mod runtime;
  3
  4pub(crate) use keygen::*;
  5pub(crate) use runtime::*;
  6
  7#[cfg(test)]
  8pub(crate) mod tests {
  9    use std::{fmt::Debug, marker::PhantomData, vec};
 10
 11    use itertools::{izip, Itertools};
 12    use rand::{thread_rng, Rng};
 13
 14    use crate::{
 15        backend::{GetModulus, ModInit, ModularOpsU64, Modulus, VectorOps},
 16        decomposer::{Decomposer, DefaultDecomposer, RlweDecomposer},
 17        ntt::{Ntt, NttBackendU64, NttInit},
 18        random::{DefaultSecureRng, NewWithSeed, RandomFillUniformInModulus},
 19        rgsw::{
 20            rlwe_auto_scratch_rows, rlwe_auto_shoup, rlwe_by_rgsw_shoup, rlwe_x_rgsw_scratch_rows,
 21            RgswCiphertextRef, RlweCiphertextMutRef, RlweKskRef, RuntimeScratchMutRef,
 22        },
 23        utils::{
 24            fill_random_ternary_secret_with_hamming_weight, generate_prime, negacyclic_mul,
 25            tests::Stats, ToShoup, TryConvertFrom1, WithLocal,
 26        },
 27        Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret,
 28    };
 29
 30    use super::{
 31        keygen::{
 32            decrypt_rlwe, generate_auto_map, rlwe_public_key, secret_key_encrypt_rgsw,
 33            seeded_auto_key_gen, seeded_secret_key_encrypt_rlwe,
 34        },
 35        rgsw_x_rgsw_scratch_rows,
 36        runtime::{rgsw_by_rgsw_inplace, rlwe_auto, rlwe_by_rgsw},
 37        RgswCiphertextMutRef,
 38    };
 39
 40    struct SeededAutoKey<M, S, Mod>
 41    where
 42        M: Matrix,
 43    {
 44        data: M,
 45        seed: S,
 46        modulus: Mod,
 47    }
 48
 49    impl<M: Matrix + MatrixEntity, S, Mod: Modulus<Element = M::MatElement>> SeededAutoKey<M, S, Mod> {
 50        fn empty<D: Decomposer>(
 51            ring_size: usize,
 52            auto_decomposer: &D,
 53            seed: S,
 54            modulus: Mod,
 55        ) -> Self {
 56            SeededAutoKey {
 57                data: M::zeros(auto_decomposer.decomposition_count().0, ring_size),
 58                seed,
 59                modulus,
 60            }
 61        }
 62    }
 63
 64    struct AutoKeyEvaluationDomain<M: Matrix, R, N> {
 65        data: M,
 66        _phantom: PhantomData<(R, N)>,
 67    }
 68
 69    impl<
 70            M: MatrixMut + MatrixEntity,
 71            Mod: Modulus<Element = M::MatElement> + Clone,
 72            R: RandomFillUniformInModulus<[M::MatElement], Mod> + NewWithSeed,
 73            N: NttInit<Mod> + Ntt<Element = M::MatElement>,
 74        > From<&SeededAutoKey<M, R::Seed, Mod>> for AutoKeyEvaluationDomain<M, R, N>
 75    where
 76        <M as Matrix>::R: RowMut,
 77        M::MatElement: Copy,
 78
 79        R::Seed: Clone,
 80    {
 81        fn from(value: &SeededAutoKey<M, R::Seed, Mod>) -> Self {
 82            let (d, ring_size) = value.data.dimension();
 83            let mut data = M::zeros(2 * d, ring_size);
 84
 85            // sample RLWE'_A(-s(X^k))
 86            let mut p_rng = R::new_with_seed(value.seed.clone());
 87            data.iter_rows_mut().take(d).for_each(|r| {
 88                RandomFillUniformInModulus::random_fill(&mut p_rng, &value.modulus, r.as_mut());
 89            });
 90
 91            // copy over RLWE'_B(-s(X^k))
 92            izip!(data.iter_rows_mut().skip(d), value.data.iter_rows()).for_each(
 93                |(to_r, from_r)| {
 94                    to_r.as_mut().copy_from_slice(from_r.as_ref());
 95                },
 96            );
 97
 98            // send RLWE'(-s(X^k)) polynomials to evaluation domain
 99            let ntt_op = N::new(&value.modulus, ring_size);
100            data.iter_rows_mut()
101                .for_each(|r| ntt_op.forward(r.as_mut()));
102
103            AutoKeyEvaluationDomain {
104                data,
105                _phantom: PhantomData,
106            }
107        }
108    }
109
110    struct RgswCiphertext<M: Matrix, Mod> {
111        /// Rgsw ciphertext polynomials
112        data: M,
113        modulus: Mod,
114        /// Decomposition for RLWE part A
115        d_a: usize,
116        /// Decomposition for RLWE part B
117        d_b: usize,
118    }
119
120    impl<M: MatrixEntity, Mod: Modulus<Element = M::MatElement>> RgswCiphertext<M, Mod> {
121        pub(crate) fn empty<D: RlweDecomposer>(
122            ring_size: usize,
123            decomposer: &D,
124            modulus: Mod,
125        ) -> RgswCiphertext<M, Mod> {
126            RgswCiphertext {
127                data: M::zeros(
128                    decomposer.a().decomposition_count().0 * 2
129                        + decomposer.b().decomposition_count().0 * 2,
130                    ring_size,
131                ),
132                d_a: decomposer.a().decomposition_count().0,
133                d_b: decomposer.b().decomposition_count().0,
134                modulus,
135            }
136        }
137    }
138
139    pub struct SeededRgswCiphertext<M, S, Mod>
140    where
141        M: Matrix,
142    {
143        pub(crate) data: M,
144        seed: S,
145        modulus: Mod,
146        /// Decomposition for RLWE part A
147        d_a: usize,
148        /// Decomposition for RLWE part B
149        d_b: usize,
150    }
151
152    impl<M: Matrix + MatrixEntity, S, Mod> SeededRgswCiphertext<M, S, Mod> {
153        pub(crate) fn empty<D: RlweDecomposer>(
154            ring_size: usize,
155            decomposer: &D,
156            seed: S,
157            modulus: Mod,
158        ) -> SeededRgswCiphertext<M, S, Mod> {
159            SeededRgswCiphertext {
160                data: M::zeros(
161                    decomposer.a().decomposition_count().0 * 2
162                        + decomposer.b().decomposition_count().0,
163                    ring_size,
164                ),
165                seed,
166                modulus,
167                d_a: decomposer.a().decomposition_count().0,
168                d_b: decomposer.b().decomposition_count().0,
169            }
170        }
171    }
172
173    impl<M: Debug + Matrix, S: Debug, Mod: Debug> Debug for SeededRgswCiphertext<M, S, Mod>
174    where
175        M::MatElement: Debug,
176    {
177        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
178            f.debug_struct("SeededRgswCiphertext")
179                .field("data", &self.data)
180                .field("seed", &self.seed)
181                .field("modulus", &self.modulus)
182                .finish()
183        }
184    }
185
186    pub struct RgswCiphertextEvaluationDomain<M, Mod, R, N> {
187        pub(crate) data: M,
188        modulus: Mod,
189        _phantom: PhantomData<(R, N)>,
190    }
191
192    impl<
193            M: MatrixMut + MatrixEntity,
194            Mod: Modulus<Element = M::MatElement> + Clone,
195            R: NewWithSeed + RandomFillUniformInModulus<[M::MatElement], Mod>,
196            N: NttInit<Mod> + Ntt<Element = M::MatElement> + Debug,
197        > From<&SeededRgswCiphertext<M, R::Seed, Mod>>
198        for RgswCiphertextEvaluationDomain<M, Mod, R, N>
199    where
200        <M as Matrix>::R: RowMut,
201        M::MatElement: Copy,
202        R::Seed: Clone,
203        M: Debug,
204    {
205        fn from(value: &SeededRgswCiphertext<M, R::Seed, Mod>) -> Self {
206            let mut data = M::zeros(value.d_a * 2 + value.d_b * 2, value.data.dimension().1);
207
208            // copy RLWE'(-sm)
209            izip!(
210                data.iter_rows_mut().take(value.d_a * 2),
211                value.data.iter_rows().take(value.d_a * 2)
212            )
213            .for_each(|(to_ri, from_ri)| {
214                to_ri.as_mut().copy_from_slice(from_ri.as_ref());
215            });
216
217            // sample A polynomials of RLWE'(m) - RLWE'A(m)
218            let mut p_rng = R::new_with_seed(value.seed.clone());
219            izip!(data.iter_rows_mut().skip(value.d_a * 2).take(value.d_b * 1))
220                .for_each(|ri| p_rng.random_fill(&value.modulus, ri.as_mut()));
221
222            // RLWE'_B(m)
223            izip!(
224                data.iter_rows_mut().skip(value.d_a * 2 + value.d_b),
225                value.data.iter_rows().skip(value.d_a * 2)
226            )
227            .for_each(|(to_ri, from_ri)| {
228                to_ri.as_mut().copy_from_slice(from_ri.as_ref());
229            });
230
231            // Send polynomials to evaluation domain
232            let ring_size = data.dimension().1;
233            let nttop = N::new(&value.modulus, ring_size);
234            data.iter_rows_mut()
235                .for_each(|ri| nttop.forward(ri.as_mut()));
236
237            Self {
238                data: data,
239                modulus: value.modulus.clone(),
240                _phantom: PhantomData,
241            }
242        }
243    }
244
245    impl<
246            M: MatrixMut + MatrixEntity,
247            Mod: Modulus<Element = M::MatElement> + Clone,
248            R,
249            N: NttInit<Mod> + Ntt<Element = M::MatElement>,
250        > From<&RgswCiphertext<M, Mod>> for RgswCiphertextEvaluationDomain<M, Mod, R, N>
251    where
252        <M as Matrix>::R: RowMut,
253        M::MatElement: Copy,
254        M: Debug,
255    {
256        fn from(value: &RgswCiphertext<M, Mod>) -> Self {
257            let mut data = M::zeros(value.d_a * 2 + value.d_b * 2, value.data.dimension().1);
258
259            // copy RLWE'(-sm)
260            izip!(
261                data.iter_rows_mut().take(value.d_a * 2),
262                value.data.iter_rows().take(value.d_a * 2)
263            )
264            .for_each(|(to_ri, from_ri)| {
265                to_ri.as_mut().copy_from_slice(from_ri.as_ref());
266            });
267
268            // copy RLWE'(m)
269            izip!(
270                data.iter_rows_mut().skip(value.d_a * 2),
271                value.data.iter_rows().skip(value.d_a * 2)
272            )
273            .for_each(|(to_ri, from_ri)| {
274                to_ri.as_mut().copy_from_slice(from_ri.as_ref());
275            });
276
277            // Send polynomials to evaluation domain
278            let ring_size = data.dimension().1;
279            let nttop = N::new(&value.modulus, ring_size);
280            data.iter_rows_mut()
281                .for_each(|ri| nttop.forward(ri.as_mut()));
282
283            Self {
284                data: data,
285                modulus: value.modulus.clone(),
286                _phantom: PhantomData,
287            }
288        }
289    }
290
291    impl<M: Debug, Mod: Debug, R, N> Debug for RgswCiphertextEvaluationDomain<M, Mod, R, N> {
292        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
293            f.debug_struct("RgswCiphertextEvaluationDomain")
294                .field("data", &self.data)
295                .field("modulus", &self.modulus)
296                .field("_phantom", &self._phantom)
297                .finish()
298        }
299    }
300
301    struct SeededRlweCiphertext<R, S, Mod> {
302        data: R,
303        seed: S,
304        modulus: Mod,
305    }
306
307    impl<R: RowEntity, S, Mod> SeededRlweCiphertext<R, S, Mod> {
308        fn empty(ring_size: usize, seed: S, modulus: Mod) -> Self {
309            SeededRlweCiphertext {
310                data: R::zeros(ring_size),
311                seed,
312                modulus,
313            }
314        }
315    }
316
317    pub struct RlweCiphertext<M, Rng> {
318        data: M,
319        _phatom: PhantomData<Rng>,
320    }
321
322    impl<
323            R: Row,
324            M: MatrixEntity<R = R, MatElement = R::Element> + MatrixMut,
325            Rng: NewWithSeed + RandomFillUniformInModulus<[M::MatElement], Mod>,
326            Mod: Modulus<Element = R::Element>,
327        > From<&SeededRlweCiphertext<R, Rng::Seed, Mod>> for RlweCiphertext<M, Rng>
328    where
329        Rng::Seed: Clone,
330        <M as Matrix>::R: RowMut,
331        R::Element: Copy,
332    {
333        fn from(value: &SeededRlweCiphertext<R, Rng::Seed, Mod>) -> Self {
334            let mut data = M::zeros(2, value.data.as_ref().len());
335
336            // sample a
337            let mut p_rng = Rng::new_with_seed(value.seed.clone());
338            RandomFillUniformInModulus::random_fill(
339                &mut p_rng,
340                &value.modulus,
341                data.get_row_mut(0),
342            );
343
344            data.get_row_mut(1).copy_from_slice(value.data.as_ref());
345
346            RlweCiphertext {
347                data,
348                _phatom: PhantomData,
349            }
350        }
351    }
352
353    struct SeededRlwePublicKey<Ro: Row, S> {
354        data: Ro,
355        seed: S,
356        modulus: Ro::Element,
357    }
358
359    impl<Ro: RowEntity, S> SeededRlwePublicKey<Ro, S> {
360        pub(crate) fn empty(ring_size: usize, seed: S, modulus: Ro::Element) -> Self {
361            Self {
362                data: Ro::zeros(ring_size),
363                seed,
364                modulus,
365            }
366        }
367    }
368
369    struct RlwePublicKey<M, R> {
370        data: M,
371        _phantom: PhantomData<R>,
372    }
373
374    impl<
375            M: MatrixMut + MatrixEntity,
376            Rng: NewWithSeed + RandomFillUniformInModulus<[M::MatElement], M::MatElement>,
377        > From<&SeededRlwePublicKey<M::R, Rng::Seed>> for RlwePublicKey<M, Rng>
378    where
379        <M as Matrix>::R: RowMut,
380        M::MatElement: Copy,
381        Rng::Seed: Copy,
382    {
383        fn from(value: &SeededRlwePublicKey<M::R, Rng::Seed>) -> Self {
384            let mut data = M::zeros(2, value.data.as_ref().len());
385
386            // sample a
387            let mut p_rng = Rng::new_with_seed(value.seed);
388            RandomFillUniformInModulus::random_fill(
389                &mut p_rng,
390                &value.modulus,
391                data.get_row_mut(0),
392            );
393
394            // copy over b
395            data.get_row_mut(1).copy_from_slice(value.data.as_ref());
396
397            Self {
398                data,
399                _phantom: PhantomData,
400            }
401        }
402    }
403
404    #[derive(Clone)]
405    struct RlweSecret {
406        pub(crate) values: Vec<i32>,
407    }
408
409    impl Secret for RlweSecret {
410        type Element = i32;
411        fn values(&self) -> &[Self::Element] {
412            &self.values
413        }
414    }
415
416    impl RlweSecret {
417        pub fn random(hw: usize, n: usize) -> RlweSecret {
418            DefaultSecureRng::with_local_mut(|rng| {
419                let mut out = vec![0i32; n];
420                fill_random_ternary_secret_with_hamming_weight(&mut out, hw, rng);
421
422                RlweSecret { values: out }
423            })
424        }
425    }
426
427    fn random_seed() -> [u8; 32] {
428        let mut rng = DefaultSecureRng::new();
429        let mut seed = [0u8; 32];
430        rng.fill_bytes(&mut seed);
431        seed
432    }
433
434    /// Encrypts m as RGSW ciphertext RGSW(m) using supplied secret key. Returns
435    /// seeded RGSW ciphertext in coefficient domain
436    fn sk_encrypt_rgsw<T: Modulus<Element = u64> + Clone>(
437        m: &[u64],
438        s: &[i32],
439        decomposer: &(DefaultDecomposer<u64>, DefaultDecomposer<u64>),
440        mod_op: &ModularOpsU64<T>,
441        ntt_op: &NttBackendU64,
442    ) -> SeededRgswCiphertext<Vec<Vec<u64>>, [u8; 32], T> {
443        let ring_size = s.len();
444        assert!(m.len() == s.len());
445
446        let mut rng = DefaultSecureRng::new();
447
448        let q = mod_op.modulus();
449        let rgsw_seed = random_seed();
450        let mut seeded_rgsw_ct = SeededRgswCiphertext::<Vec<Vec<u64>>, [u8; 32], T>::empty(
451            ring_size as usize,
452            decomposer,
453            rgsw_seed,
454            q.clone(),
455        );
456        let mut p_rng = DefaultSecureRng::new_seeded(rgsw_seed);
457        secret_key_encrypt_rgsw(
458            &mut seeded_rgsw_ct.data,
459            m,
460            &decomposer.a().gadget_vector(),
461            &decomposer.b().gadget_vector(),
462            s,
463            mod_op,
464            ntt_op,
465            &mut p_rng,
466            &mut rng,
467        );
468        seeded_rgsw_ct
469    }
470
471    #[test]
472    fn rlwe_encrypt_decryption() {
473        let logq = 50;
474        let logp = 2;
475        let ring_size = 1 << 4;
476        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
477        let p = 1u64 << logp;
478
479        let mut rng = DefaultSecureRng::new();
480
481        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
482
483        // sample m0
484        let mut m0 = vec![0u64; ring_size as usize];
485        RandomFillUniformInModulus::<[u64], u64>::random_fill(
486            &mut rng,
487            &(1u64 << logp),
488            m0.as_mut_slice(),
489        );
490
491        let ntt_op = NttBackendU64::new(&q, ring_size as usize);
492        let mod_op = ModularOpsU64::new(q);
493
494        // encrypt m0
495        let encoded_m = m0
496            .iter()
497            .map(|v| (((*v as f64) * q as f64) / (p as f64)).round() as u64)
498            .collect_vec();
499        let seed = random_seed();
500        let mut rlwe_in_ct =
501            SeededRlweCiphertext::<Vec<u64>, _, _>::empty(ring_size as usize, seed, q);
502        let mut p_rng = DefaultSecureRng::new_seeded(seed);
503        seeded_secret_key_encrypt_rlwe(
504            &encoded_m,
505            &mut rlwe_in_ct.data,
506            s.values(),
507            &mod_op,
508            &ntt_op,
509            &mut p_rng,
510            &mut rng,
511        );
512        let rlwe_in_ct = RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&rlwe_in_ct);
513
514        let mut encoded_m_back = vec![0u64; ring_size as usize];
515        decrypt_rlwe(
516            &rlwe_in_ct.data,
517            s.values(),
518            &mut encoded_m_back,
519            &ntt_op,
520            &mod_op,
521        );
522        let m_back = encoded_m_back
523            .iter()
524            .map(|v| (((*v as f64 * p as f64) / q as f64).round() as u64) % p)
525            .collect_vec();
526        assert_eq!(m0, m_back);
527    }
528
529    #[test]
530    fn rlwe_by_rgsw_works() {
531        let logq = 50;
532        let logp = 2;
533        let ring_size = 1 << 4;
534        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
535        let p: u64 = 1u64 << logp;
536
537        let mut rng = DefaultSecureRng::new_seeded([0u8; 32]);
538
539        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
540
541        let mut m0 = vec![0u64; ring_size as usize];
542        RandomFillUniformInModulus::<[u64], _>::random_fill(
543            &mut rng,
544            &(1u64 << logp),
545            m0.as_mut_slice(),
546        );
547        let mut m1 = vec![0u64; ring_size as usize];
548        m1[thread_rng().gen_range(0..ring_size) as usize] = 1;
549
550        let ntt_op = NttBackendU64::new(&q, ring_size as usize);
551        let mod_op = ModularOpsU64::new(q);
552        let d_rgsw = 10;
553        let logb = 5;
554        let decomposer = (
555            DefaultDecomposer::new(q, logb, d_rgsw),
556            DefaultDecomposer::new(q, logb, d_rgsw),
557        );
558
559        // create public key
560        let pk_seed = random_seed();
561        let mut pk_prng = DefaultSecureRng::new_seeded(pk_seed);
562        let mut seeded_pk =
563            SeededRlwePublicKey::<Vec<u64>, _>::empty(ring_size as usize, pk_seed, q);
564        rlwe_public_key(
565            &mut seeded_pk.data,
566            s.values(),
567            &ntt_op,
568            &mod_op,
569            &mut pk_prng,
570            &mut rng,
571        );
572        // let pk = RlwePublicKey::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_pk);
573
574        // Encrypt m1 as RGSW(m1)
575        let rgsw_ct = {
576            // Encryption m1 as RGSW(m1) using secret key
577            let seeded_rgsw_ct = sk_encrypt_rgsw(&m1, s.values(), &decomposer, &mod_op, &ntt_op);
578            RgswCiphertextEvaluationDomain::<Vec<Vec<u64>>, _,DefaultSecureRng, NttBackendU64>::from(&seeded_rgsw_ct)
579        };
580
581        // Encrypt m0 as RLWE(m0)
582        let mut rlwe_in_ct = {
583            let encoded_m = m0
584                .iter()
585                .map(|v| (((*v as f64) * q as f64) / (p as f64)).round() as u64)
586                .collect_vec();
587
588            let seed = random_seed();
589            let mut p_rng = DefaultSecureRng::new_seeded(seed);
590            let mut seeded_rlwe = SeededRlweCiphertext::empty(ring_size as usize, seed, q);
591            seeded_secret_key_encrypt_rlwe(
592                &encoded_m,
593                &mut seeded_rlwe.data,
594                s.values(),
595                &mod_op,
596                &ntt_op,
597                &mut p_rng,
598                &mut rng,
599            );
600            RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_rlwe)
601        };
602
603        // RLWE(m0m1) = RLWE(m0) x RGSW(m1)
604        let mut scratch_space =
605            vec![vec![0u64; ring_size as usize]; rlwe_x_rgsw_scratch_rows(&decomposer)];
606
607        // rlwe x rgsw with with soup repr
608        let rlwe_in_ct_shoup = {
609            let mut rlwe_in_ct_shoup = rlwe_in_ct.data.clone();
610
611            let rgsw_ct_shoup = ToShoup::to_shoup(&rgsw_ct.data, q);
612
613            rlwe_by_rgsw_shoup(
614                &mut RlweCiphertextMutRef::new(rlwe_in_ct_shoup.as_mut()),
615                &RgswCiphertextRef::new(
616                    rgsw_ct.data.as_ref(),
617                    decomposer.a().decomposition_count().0,
618                    decomposer.b().decomposition_count().0,
619                ),
620                &RgswCiphertextRef::new(
621                    rgsw_ct_shoup.as_ref(),
622                    decomposer.a().decomposition_count().0,
623                    decomposer.b().decomposition_count().0,
624                ),
625                &mut RuntimeScratchMutRef::new(scratch_space.as_mut()),
626                &decomposer,
627                &ntt_op,
628                &mod_op,
629                false,
630            );
631
632            rlwe_in_ct_shoup
633        };
634
635        // rlwe x rgsw normal
636        {
637            rlwe_by_rgsw(
638                &mut RlweCiphertextMutRef::new(rlwe_in_ct.data.as_mut()),
639                &RgswCiphertextRef::new(
640                    rgsw_ct.data.as_ref(),
641                    decomposer.a().decomposition_count().0,
642                    decomposer.b().decomposition_count().0,
643                ),
644                &mut RuntimeScratchMutRef::new(scratch_space.as_mut()),
645                &decomposer,
646                &ntt_op,
647                &mod_op,
648                false,
649            );
650        }
651
652        // output from both functions must be equal
653        assert_eq!(rlwe_in_ct.data, rlwe_in_ct_shoup);
654
655        // Decrypt RLWE(m0m1)
656        let mut encoded_m0m1_back = vec![0u64; ring_size as usize];
657        decrypt_rlwe(
658            &rlwe_in_ct_shoup,
659            s.values(),
660            &mut encoded_m0m1_back,
661            &ntt_op,
662            &mod_op,
663        );
664        let m0m1_back = encoded_m0m1_back
665            .iter()
666            .map(|v| (((*v as f64 * p as f64) / (q as f64)).round() as u64) % p)
667            .collect_vec();
668
669        let mul_mod = |v0: &u64, v1: &u64| (v0 * v1) % p;
670        let m0m1 = negacyclic_mul(&m0, &m1, mul_mod, p);
671
672        // {
673        //     // measure noise
674        //     let encoded_m_ideal = m0m1
675        //         .iter()
676        //         .map(|v| (((*v as f64) * q as f64) / (p as f64)).round() as u64)
677        //         .collect_vec();
678
679        //     let noise = measure_noise(&rlwe_in_ct, &encoded_m_ideal, &ntt_op,
680        // &mod_op, s.values());     println!("Noise RLWE(m0m1)(=
681        // RLWE(m0)xRGSW(m1)) : {noise}"); }
682
683        assert!(
684            m0m1 == m0m1_back,
685            "Expected {:?} \n Got {:?}",
686            m0m1,
687            m0m1_back
688        );
689    }
690
691    #[test]
692    fn rlwe_auto_works() {
693        let logq = 55;
694        let ring_size = 1 << 11;
695        let q = generate_prime(logq, 2 * ring_size, 1u64 << logq).unwrap();
696        let logp = 3;
697        let p = 1u64 << logp;
698        let d_rgsw = 5;
699        let logb = 11;
700
701        let mut rng = DefaultSecureRng::new();
702        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
703
704        let mut m = vec![0u64; ring_size as usize];
705        RandomFillUniformInModulus::random_fill(&mut rng, &p, m.as_mut_slice());
706        let encoded_m = m
707            .iter()
708            .map(|v| (((*v as f64 * q as f64) / (p as f64)).round() as u64))
709            .collect_vec();
710
711        let ntt_op = NttBackendU64::new(&q, ring_size as usize);
712        let mod_op = ModularOpsU64::new(q);
713
714        // RLWE_{s}(m)
715        let seed_rlwe = random_seed();
716        let mut seeded_rlwe_m = SeededRlweCiphertext::empty(ring_size as usize, seed_rlwe, q);
717        let mut p_rng = DefaultSecureRng::new_seeded(seed_rlwe);
718        seeded_secret_key_encrypt_rlwe(
719            &encoded_m,
720            &mut seeded_rlwe_m.data,
721            s.values(),
722            &mod_op,
723            &ntt_op,
724            &mut p_rng,
725            &mut rng,
726        );
727        let mut rlwe_m = RlweCiphertext::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_rlwe_m);
728
729        let auto_k = -125;
730
731        // Generate auto key to key switch from s^k to s
732        let decomposer = DefaultDecomposer::new(q, logb, d_rgsw);
733        let seed_auto = random_seed();
734        let mut seeded_auto_key =
735            SeededAutoKey::empty(ring_size as usize, &decomposer, seed_auto, q);
736        let mut p_rng = DefaultSecureRng::new_seeded(seed_auto);
737        let gadget_vector = decomposer.gadget_vector();
738        seeded_auto_key_gen(
739            &mut seeded_auto_key.data,
740            s.values(),
741            auto_k,
742            &gadget_vector,
743            &mod_op,
744            &ntt_op,
745            &mut p_rng,
746            &mut rng,
747        );
748        let auto_key =
749            AutoKeyEvaluationDomain::<Vec<Vec<u64>>, DefaultSecureRng, NttBackendU64>::from(
750                &seeded_auto_key,
751            );
752
753        // Send RLWE_{s}(m) -> RLWE_{s}(m^k)
754        let mut scratch_space =
755            vec![vec![0; ring_size as usize]; rlwe_auto_scratch_rows(&decomposer)];
756        let (auto_map_index, auto_map_sign) = generate_auto_map(ring_size as usize, auto_k);
757
758        // galois auto with auto key in shoup repr
759        let rlwe_m_shoup = {
760            let auto_key_shoup = ToShoup::to_shoup(&auto_key.data, q);
761            let mut rlwe_m_shoup = rlwe_m.data.clone();
762            rlwe_auto_shoup(
763                &mut RlweCiphertextMutRef::new(&mut rlwe_m_shoup),
764                &RlweKskRef::new(&auto_key.data, decomposer.decomposition_count().0),
765                &RlweKskRef::new(&auto_key_shoup, decomposer.decomposition_count().0),
766                &mut RuntimeScratchMutRef::new(&mut scratch_space),
767                &auto_map_index,
768                &auto_map_sign,
769                &mod_op,
770                &ntt_op,
771                &decomposer,
772                false,
773            );
774            rlwe_m_shoup
775        };
776
777        // normal galois auto
778        {
779            rlwe_auto(
780                &mut RlweCiphertextMutRef::new(rlwe_m.data.as_mut()),
781                &RlweKskRef::new(auto_key.data.as_ref(), decomposer.decomposition_count().0),
782                &mut RuntimeScratchMutRef::new(scratch_space.as_mut()),
783                &auto_map_index,
784                &auto_map_sign,
785                &mod_op,
786                &ntt_op,
787                &decomposer,
788                false,
789            );
790        }
791
792        // rlwe out from both functions must be same
793        assert_eq!(rlwe_m.data, rlwe_m_shoup);
794
795        let rlwe_m_k = rlwe_m;
796
797        // Decrypt RLWE_{s}(m^k) and check
798        let mut encoded_m_k_back = vec![0u64; ring_size as usize];
799        decrypt_rlwe(
800            &rlwe_m_k.data,
801            s.values(),
802            &mut encoded_m_k_back,
803            &ntt_op,
804            &mod_op,
805        );
806        let m_k_back = encoded_m_k_back
807            .iter()
808            .map(|v| (((*v as f64 * p as f64) / q as f64).round() as u64) % p)
809            .collect_vec();
810
811        let mut m_k = vec![0u64; ring_size as usize];
812        // Send \delta m -> \delta m^k
813        izip!(m.iter(), auto_map_index.iter(), auto_map_sign.iter()).for_each(
814            |(v, to_index, sign)| {
815                if !*sign {
816                    m_k[*to_index] = (p - *v) % p;
817                } else {
818                    m_k[*to_index] = *v;
819                }
820            },
821        );
822
823        // {
824        //     let encoded_m_k = m_k
825        //         .iter()
826        //         .map(|v| ((*v as f64 * q as f64) / p as f64).round() as u64)
827        //         .collect_vec();
828
829        //     let noise = measure_noise(&rlwe_m_k, &encoded_m_k, &ntt_op, &mod_op,
830        // s.values());     println!("Ksk noise: {noise}");
831        // }
832
833        assert_eq!(m_k_back, m_k);
834    }
835
836    /// Collect noise stats of RGSW ciphertext
837    ///
838    /// - rgsw_ct: RGSW ciphertext must be in coefficient domain
839    fn rgsw_noise_stats<T: Modulus<Element = u64> + Clone>(
840        rgsw_ct: &[Vec<u64>],
841        m: &[u64],
842        s: &[i32],
843        decomposer: &(DefaultDecomposer<u64>, DefaultDecomposer<u64>),
844        q: &T,
845    ) -> Stats<i64> {
846        let gadget_vector_a = decomposer.a().gadget_vector();
847        let gadget_vector_b = decomposer.b().gadget_vector();
848        let d_a = gadget_vector_a.len();
849        let d_b = gadget_vector_b.len();
850        let ring_size = s.len();
851        assert!(Matrix::dimension(&rgsw_ct) == (d_a * 2 + d_b * 2, ring_size));
852        assert!(m.len() == ring_size);
853
854        let mod_op = ModularOpsU64::new(q.clone());
855        let ntt_op = NttBackendU64::new(q, ring_size);
856
857        let mul_mod =
858            |a: &u64, b: &u64| ((*a as u128 * *b as u128) % q.q().unwrap() as u128) as u64;
859        let s_poly = Vec::<u64>::try_convert_from(s, q);
860        let mut neg_s = s_poly.clone();
861        mod_op.elwise_neg_mut(neg_s.as_mut());
862        let neg_sm0m1 = negacyclic_mul(&neg_s, &m, mul_mod, q.q().unwrap());
863
864        let mut stats = Stats::new();
865
866        // RLWE(\beta^j -s * m)
867        for j in 0..d_a {
868            let want_m = {
869                // RLWE(\beta^j -s * m)
870                let mut beta_neg_sm0m1 = vec![0u64; ring_size as usize];
871                mod_op.elwise_scalar_mul(beta_neg_sm0m1.as_mut(), &neg_sm0m1, &gadget_vector_a[j]);
872                beta_neg_sm0m1
873            };
874
875            let mut rlwe = vec![vec![0u64; ring_size as usize]; 2];
876            rlwe[0].copy_from_slice(rgsw_ct.get_row_slice(j));
877            rlwe[1].copy_from_slice(rgsw_ct.get_row_slice(d_a + j));
878
879            let mut got_m = vec![0; ring_size];
880            decrypt_rlwe(&rlwe, s, &mut got_m, &ntt_op, &mod_op);
881
882            let mut diff = want_m;
883            mod_op.elwise_sub_mut(diff.as_mut(), got_m.as_ref());
884            stats.add_more(&Vec::<i64>::try_convert_from(&diff, q));
885        }
886
887        // RLWE(\beta^j  m)
888        for j in 0..d_b {
889            let want_m = {
890                // RLWE(\beta^j  m)
891                let mut beta_m0m1 = vec![0u64; ring_size as usize];
892                mod_op.elwise_scalar_mul(beta_m0m1.as_mut(), &m, &gadget_vector_b[j]);
893                beta_m0m1
894            };
895
896            let mut rlwe = vec![vec![0u64; ring_size as usize]; 2];
897            rlwe[0].copy_from_slice(rgsw_ct.get_row_slice(d_a * 2 + j));
898            rlwe[1].copy_from_slice(rgsw_ct.get_row_slice(d_a * 2 + d_b + j));
899
900            let mut got_m = vec![0; ring_size];
901            decrypt_rlwe(&rlwe, s, &mut got_m, &ntt_op, &mod_op);
902
903            let mut diff = want_m;
904            mod_op.elwise_sub_mut(diff.as_mut(), got_m.as_ref());
905            stats.add_more(&Vec::<i64>::try_convert_from(&diff, q));
906        }
907
908        stats
909    }
910
911    #[test]
912    fn print_noise_stats_rgsw_x_rgsw() {
913        let logq = 60;
914        let logp = 2;
915        let ring_size = 1 << 11;
916        let q = generate_prime(logq, ring_size, 1u64 << logq).unwrap();
917        let d_rgsw = 12;
918        let logb = 5;
919
920        let s = RlweSecret::random((ring_size >> 1) as usize, ring_size as usize);
921
922        let ntt_op = NttBackendU64::new(&q, ring_size as usize);
923        let mod_op = ModularOpsU64::new(q);
924        let decomposer = (
925            DefaultDecomposer::new(q, logb, d_rgsw),
926            DefaultDecomposer::new(q, logb, d_rgsw),
927        );
928
929        let d_a = decomposer.a().decomposition_count().0;
930        let d_b = decomposer.b().decomposition_count().0;
931
932        let mul_mod = |a: &u64, b: &u64| ((*a as u128 * *b as u128) % q as u128) as u64;
933
934        let mut carry_m = vec![0u64; ring_size as usize];
935        carry_m[thread_rng().gen_range(0..ring_size) as usize] = 1 << logp;
936
937        // RGSW(carry_m)
938        let mut rgsw_carrym = {
939            let seeded_rgsw = sk_encrypt_rgsw(&carry_m, s.values(), &decomposer, &mod_op, &ntt_op);
940            let mut rgsw_eval =
941                RgswCiphertextEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(
942                    &seeded_rgsw,
943                );
944            rgsw_eval
945                .data
946                .iter_mut()
947                .for_each(|ri| ntt_op.backward(ri.as_mut()));
948            rgsw_eval.data
949        };
950
951        let mut scratch_matrix = vec![
952            vec![0u64; ring_size as usize];
953            rgsw_x_rgsw_scratch_rows(&decomposer, &decomposer)
954        ];
955
956        rgsw_noise_stats(&rgsw_carrym, &carry_m, s.values(), &decomposer, &q);
957
958        for i in 0..8 {
959            let mut m = vec![0u64; ring_size as usize];
960            m[thread_rng().gen_range(0..ring_size) as usize] = if (i & 1) == 1 { q - 1 } else { 1 };
961            let rgsw_m =
962                RgswCiphertextEvaluationDomain::<_, _, DefaultSecureRng, NttBackendU64>::from(
963                    &sk_encrypt_rgsw(&m, s.values(), &decomposer, &mod_op, &ntt_op),
964                );
965
966            rgsw_by_rgsw_inplace(
967                &mut RgswCiphertextMutRef::new(rgsw_carrym.as_mut(), d_a, d_b),
968                &RgswCiphertextRef::new(rgsw_m.data.as_ref(), d_a, d_b),
969                &decomposer,
970                &decomposer,
971                &mut RuntimeScratchMutRef::new(scratch_matrix.as_mut()),
972                &ntt_op,
973                &mod_op,
974            );
975
976            // measure noise
977            carry_m = negacyclic_mul(&carry_m, &m, mul_mod, q);
978            let stats = rgsw_noise_stats(&rgsw_carrym, &carry_m, s.values(), &decomposer, &q);
979            println!(
980                "Log2 of noise std after {i} RGSW x RGSW: {}",
981                stats.std_dev().abs().log2()
982            );
983        }
984    }
985}