poulpy_core/encryption/compressed/
gglwe_ct.rs

1use poulpy_hal::{
2    api::{
3        ScratchAvailable, SvpApplyInplace, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
4        VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxFillUniform,
5        VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, ZnxZero,
6    },
7    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
8    source::Source,
9};
10
11use crate::{
12    TakeGLWEPt,
13    encryption::glwe_encrypt_sk_internal,
14    layouts::{GGLWECiphertext, Infos, compressed::GGLWECiphertextCompressed, prepared::GLWESecretPrepared},
15};
16
17impl GGLWECiphertextCompressed<Vec<u8>> {
18    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
19    where
20        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes,
21    {
22        GGLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k)
23    }
24}
25
26impl<D: DataMut> GGLWECiphertextCompressed<D> {
27    #[allow(clippy::too_many_arguments)]
28    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
29        &mut self,
30        module: &Module<B>,
31        pt: &ScalarZnx<DataPt>,
32        sk: &GLWESecretPrepared<DataSk, B>,
33        seed: [u8; 32],
34        source_xe: &mut Source,
35        sigma: f64,
36        scratch: &mut Scratch<B>,
37    ) where
38        Module<B>: VecZnxAddScalarInplace
39            + VecZnxDftAllocBytes
40            + VecZnxBigNormalize<B>
41            + VecZnxDftFromVecZnx<B>
42            + SvpApplyInplace<B>
43            + VecZnxDftToVecZnxBigConsume<B>
44            + VecZnxNormalizeTmpBytes
45            + VecZnxFillUniform
46            + VecZnxSubABInplace
47            + VecZnxAddInplace
48            + VecZnxNormalizeInplace<B>
49            + VecZnxAddNormal
50            + VecZnxNormalize<B>
51            + VecZnxSub,
52        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
53    {
54        #[cfg(debug_assertions)]
55        {
56            use poulpy_hal::api::ZnxInfos;
57
58            assert_eq!(
59                self.rank_in(),
60                pt.cols(),
61                "self.rank_in(): {} != pt.cols(): {}",
62                self.rank_in(),
63                pt.cols()
64            );
65            assert_eq!(
66                self.rank_out(),
67                sk.rank(),
68                "self.rank_out(): {} != sk.rank(): {}",
69                self.rank_out(),
70                sk.rank()
71            );
72            assert_eq!(self.n(), sk.n());
73            assert_eq!(pt.n(), sk.n());
74            assert!(
75                scratch.available()
76                    >= GGLWECiphertextCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k()),
77                "scratch.available: {} < GGLWECiphertext::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
78                scratch.available(),
79                self.rank(),
80                self.size(),
81                GGLWECiphertextCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k())
82            );
83            assert!(
84                self.rows() * self.digits() * self.basek() <= self.k(),
85                "self.rows() : {} * self.digits() : {} * self.basek() : {} = {} >= self.k() = {}",
86                self.rows(),
87                self.digits(),
88                self.basek(),
89                self.rows() * self.digits() * self.basek(),
90                self.k()
91            );
92        }
93
94        let rows: usize = self.rows();
95        let digits: usize = self.digits();
96        let basek: usize = self.basek();
97        let k: usize = self.k();
98        let rank_in: usize = self.rank_in();
99        let cols: usize = self.rank_out() + 1;
100
101        let mut source_xa = Source::new(seed);
102
103        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(sk.n(), basek, k);
104        (0..rank_in).for_each(|col_i| {
105            (0..rows).for_each(|row_i| {
106                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
107                tmp_pt.data.zero(); // zeroes for next iteration
108                module.vec_znx_add_scalar_inplace(
109                    &mut tmp_pt.data,
110                    0,
111                    (digits - 1) + row_i * digits,
112                    pt,
113                    col_i,
114                );
115                module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scrach_1);
116
117                let (seed, mut source_xa_tmp) = source_xa.branch();
118                self.seed[col_i * rows + row_i] = seed;
119
120                glwe_encrypt_sk_internal(
121                    module,
122                    self.basek(),
123                    self.k(),
124                    &mut self.at_mut(row_i, col_i).data,
125                    cols,
126                    true,
127                    Some((&tmp_pt, 0)),
128                    sk,
129                    &mut source_xa_tmp,
130                    source_xe,
131                    sigma,
132                    scrach_1,
133                );
134            });
135        });
136    }
137}