poulpy/core/src/lwe/encryption.rs
core/src/lwe/encryption.rsRaw
1use backend::hal::{
2 api::{
3 ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft, VecZnxAddNormal,
4 VecZnxAddScalarInplace, VecZnxAutomorphismInplace, VecZnxFillUniform, VecZnxNormalizeInplace, VecZnxSwithcDegree,
5 ZnxView, ZnxViewMut, ZnxZero,
6 },
7 layouts::{Backend, DataMut, DataRef, Module, Scratch, ScratchOwned, VecZnx},
8 oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl},
9};
10use sampling::source::Source;
11
12use crate::{
13 GGLWEEncryptSkFamily, GLWESecret, GLWEToLWESwitchingKey, Infos, LWECiphertext, LWESecret, LWESwitchingKey,
14 LWEToGLWESwitchingKey, SIX_SIGMA, TakeGLWESecret, TakeGLWESecretExec, lwe::LWEPlaintext,
15};
16
17impl<DataSelf: DataMut> LWECiphertext<DataSelf> {
18 pub fn encrypt_sk<DataPt, DataSk, B: Backend>(
19 &mut self,
20 module: &Module<B>,
21 pt: &LWEPlaintext<DataPt>,
22 sk: &LWESecret<DataSk>,
23 source_xa: &mut Source,
24 source_xe: &mut Source,
25 sigma: f64,
26 ) where
27 DataPt: DataRef,
28 DataSk: DataRef,
29 Module<B>: VecZnxFillUniform + VecZnxAddNormal + VecZnxNormalizeInplace<B>,
30 B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
31 {
32 #[cfg(debug_assertions)]
33 {
34 assert_eq!(self.n(), sk.n())
35 }
36
37 let basek: usize = self.basek();
38 let k: usize = self.k();
39
40 module.vec_znx_fill_uniform(basek, &mut self.data, 0, k, source_xa);
41
42 let mut tmp_znx: VecZnx<Vec<u8>> = VecZnx::alloc(1, 1, self.size());
43
44 let min_size = self.size().min(pt.size());
45
46 (0..min_size).for_each(|i| {
47 tmp_znx.at_mut(0, i)[0] = pt.data.at(0, i)[0]
48 - self.data.at(0, i)[1..]
49 .iter()
50 .zip(sk.data.at(0, 0))
51 .map(|(x, y)| x * y)
52 .sum::<i64>();
53 });
54
55 (min_size..self.size()).for_each(|i| {
56 tmp_znx.at_mut(0, i)[0] -= self.data.at(0, i)[1..]
57 .iter()
58 .zip(sk.data.at(0, 0))
59 .map(|(x, y)| x * y)
60 .sum::<i64>();
61 });
62
63 module.vec_znx_add_normal(
64 basek,
65 &mut self.data,
66 0,
67 k,
68 source_xe,
69 sigma,
70 sigma * SIX_SIGMA,
71 );
72
73 module.vec_znx_normalize_inplace(
74 basek,
75 &mut tmp_znx,
76 0,
77 ScratchOwned::alloc(size_of::<i64>()).borrow(),
78 );
79
80 (0..self.size()).for_each(|i| {
81 self.data.at_mut(0, i)[0] = tmp_znx.at(0, i)[0];
82 });
83 }
84}
85
86impl<D: DataMut> GLWEToLWESwitchingKey<D> {
87 pub fn encrypt_sk<DLwe, DGlwe, B: Backend>(
88 &mut self,
89 module: &Module<B>,
90 sk_lwe: &LWESecret<DLwe>,
91 sk_glwe: &GLWESecret<DGlwe>,
92 source_xa: &mut Source,
93 source_xe: &mut Source,
94 sigma: f64,
95 scratch: &mut Scratch<B>,
96 ) where
97 DLwe: DataRef,
98 DGlwe: DataRef,
99 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAutomorphismInplace + VecZnxSwithcDegree + VecZnxAddScalarInplace,
100 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx,
101 {
102 #[cfg(debug_assertions)]
103 {
104 assert!(sk_lwe.n() <= module.n());
105 }
106
107 let (mut sk_lwe_as_glwe, scratch1) = scratch.take_glwe_secret(sk_glwe.n(), 1);
108 sk_lwe_as_glwe.data.zero();
109 sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n()].copy_from_slice(sk_lwe.data.at(0, 0));
110 module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0);
111
112 self.0.encrypt_sk(
113 module,
114 sk_glwe,
115 &sk_lwe_as_glwe,
116 source_xa,
117 source_xe,
118 sigma,
119 scratch1,
120 );
121 }
122}
123
124impl<D: DataMut> LWEToGLWESwitchingKey<D> {
125 pub fn encrypt_sk<DLwe, DGlwe, B: Backend>(
126 &mut self,
127 module: &Module<B>,
128 sk_lwe: &LWESecret<DLwe>,
129 sk_glwe: &GLWESecret<DGlwe>,
130 source_xa: &mut Source,
131 source_xe: &mut Source,
132 sigma: f64,
133 scratch: &mut Scratch<B>,
134 ) where
135 DLwe: DataRef,
136 DGlwe: DataRef,
137 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAutomorphismInplace + VecZnxSwithcDegree + VecZnxAddScalarInplace,
138 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx,
139 {
140 #[cfg(debug_assertions)]
141 {
142 assert!(sk_lwe.n() <= module.n());
143 }
144
145 let (mut sk_lwe_as_glwe, scratch1) = scratch.take_glwe_secret(sk_glwe.n(), 1);
146 sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n()].copy_from_slice(sk_lwe.data.at(0, 0));
147 sk_lwe_as_glwe.data.at_mut(0, 0)[sk_lwe.n()..].fill(0);
148 module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data.as_vec_znx_mut(), 0);
149
150 self.0.encrypt_sk(
151 module,
152 &sk_lwe_as_glwe,
153 &sk_glwe,
154 source_xa,
155 source_xe,
156 sigma,
157 scratch1,
158 );
159 }
160}
161
162impl<D: DataMut> LWESwitchingKey<D> {
163 pub fn encrypt_sk<DIn, DOut, B: Backend>(
164 &mut self,
165 module: &Module<B>,
166 sk_lwe_in: &LWESecret<DIn>,
167 sk_lwe_out: &LWESecret<DOut>,
168 source_xa: &mut Source,
169 source_xe: &mut Source,
170 sigma: f64,
171 scratch: &mut Scratch<B>,
172 ) where
173 DIn: DataRef,
174 DOut: DataRef,
175 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAutomorphismInplace + VecZnxSwithcDegree + VecZnxAddScalarInplace,
176 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx,
177 {
178 #[cfg(debug_assertions)]
179 {
180 assert!(sk_lwe_in.n() <= self.n());
181 assert!(sk_lwe_out.n() <= self.n());
182 assert!(self.n() <= module.n());
183 }
184
185 let (mut sk_in_glwe, scratch1) = scratch.take_glwe_secret(self.n(), 1);
186 let (mut sk_out_glwe, scratch2) = scratch1.take_glwe_secret(self.n(), 1);
187
188 sk_out_glwe.data.at_mut(0, 0)[..sk_lwe_out.n()].copy_from_slice(sk_lwe_out.data.at(0, 0));
189 sk_out_glwe.data.at_mut(0, 0)[sk_lwe_out.n()..].fill(0);
190 module.vec_znx_automorphism_inplace(-1, &mut sk_out_glwe.data.as_vec_znx_mut(), 0);
191
192 sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_in.n()].copy_from_slice(sk_lwe_in.data.at(0, 0));
193 sk_in_glwe.data.at_mut(0, 0)[sk_lwe_in.n()..].fill(0);
194 module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data.as_vec_znx_mut(), 0);
195
196 self.0.encrypt_sk(
197 module,
198 &sk_in_glwe,
199 &sk_out_glwe,
200 source_xa,
201 source_xe,
202 sigma,
203 scratch2,
204 );
205 }
206}