poulpy_core/encryption/compressed/
gglwe_ct.rs

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