poulpy/bin_fhe/blind_rotation/key.rs
bin_fhe/blind_rotation/key.rsRaw
1use backend::hal::{
2 api::{
3 FillUniform, Reset, ScratchAvailable, SvpPPolAlloc, SvpPrepare, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace,
4 VmpPMatAlloc, VmpPMatPrepare, ZnxInfos, ZnxView, ZnxViewMut,
5 },
6 layouts::{Backend, Data, DataMut, DataRef, Module, ReaderFrom, ScalarZnx, ScalarZnxToRef, Scratch, SvpPPol, WriterTo},
7};
8use sampling::source::Source;
9
10use crate::{
11 Distribution, Infos,
12 layouts::{
13 GGSWCiphertext, LWESecret,
14 prepared::{GGSWCiphertextExec, GLWESecretExec},
15 },
16};
17use std::fmt;
18
19use crate::trait_families::GGSWEncryptSkFamily;
20
21#[derive(Clone)]
22pub struct BlindRotationKeyCGGI<D: Data> {
23 pub(crate) keys: Vec<GGSWCiphertext<D>>,
24 pub(crate) dist: Distribution,
25}
26
27impl<D: DataRef> fmt::Debug for BlindRotationKeyCGGI<D> {
28 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
29 write!(f, "{}", self)
30 }
31}
32
33impl<D: Data> PartialEq for BlindRotationKeyCGGI<D> {
34 fn eq(&self, other: &Self) -> bool {
35 if self.keys.len() != other.keys.len() {
36 return false;
37 }
38 for (a, b) in self.keys.iter().zip(other.keys.iter()) {
39 if a != b {
40 return false;
41 }
42 }
43 self.dist == other.dist
44 }
45}
46
47impl<D: Data> Eq for BlindRotationKeyCGGI<D> {}
48
49impl<D: DataRef> fmt::Display for BlindRotationKeyCGGI<D> {
50 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51 for (i, key) in self.keys.iter().enumerate() {
52 write!(f, "key[{}]: {}", i, key)?;
53 }
54 writeln!(f, "{:?}", self.dist)
55 }
56}
57
58impl<D: DataMut> Reset for BlindRotationKeyCGGI<D> {
59 fn reset(&mut self) {
60 self.keys.iter_mut().for_each(|key| key.reset());
61 self.dist = Distribution::NONE;
62 }
63}
64
65impl<D: DataMut> FillUniform for BlindRotationKeyCGGI<D> {
66 fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
67 self.keys
68 .iter_mut()
69 .for_each(|key| key.fill_uniform(source));
70 }
71}
72
73use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
74
75impl<D: DataMut> ReaderFrom for BlindRotationKeyCGGI<D> {
76 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
77 match Distribution::read_from(reader) {
78 Ok(dist) => self.dist = dist,
79 Err(e) => return Err(e),
80 }
81 let len: usize = reader.read_u64::<LittleEndian>()? as usize;
82 if self.keys.len() != len {
83 return Err(std::io::Error::new(
84 std::io::ErrorKind::InvalidData,
85 format!("self.keys.len()={} != read len={}", self.keys.len(), len),
86 ));
87 }
88 for key in &mut self.keys {
89 key.read_from(reader)?;
90 }
91 Ok(())
92 }
93}
94
95impl<D: DataRef> WriterTo for BlindRotationKeyCGGI<D> {
96 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
97 match self.dist.write_to(writer) {
98 Ok(()) => {}
99 Err(e) => return Err(e),
100 }
101 writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
102 for key in &self.keys {
103 key.write_to(writer)?;
104 }
105 Ok(())
106 }
107}
108
109impl BlindRotationKeyCGGI<Vec<u8>> {
110 pub fn alloc(n_gglwe: usize, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
111 let mut data: Vec<GGSWCiphertext<Vec<u8>>> = Vec::with_capacity(n_lwe);
112 (0..n_lwe).for_each(|_| data.push(GGSWCiphertext::alloc(n_gglwe, basek, k, rows, 1, rank)));
113 Self {
114 keys: data,
115 dist: Distribution::NONE,
116 }
117 }
118
119 pub fn generate_from_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
120 where
121 Module<B>: GGSWEncryptSkFamily<B>,
122 {
123 GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k, rank)
124 }
125}
126
127impl<D: DataRef> BlindRotationKeyCGGI<D> {
128 #[allow(dead_code)]
129 pub(crate) fn n(&self) -> usize {
130 self.keys[0].n()
131 }
132
133 #[allow(dead_code)]
134 pub(crate) fn rows(&self) -> usize {
135 self.keys[0].rows()
136 }
137
138 #[allow(dead_code)]
139 pub(crate) fn k(&self) -> usize {
140 self.keys[0].k()
141 }
142
143 #[allow(dead_code)]
144 pub(crate) fn size(&self) -> usize {
145 self.keys[0].size()
146 }
147
148 #[allow(dead_code)]
149 pub(crate) fn rank(&self) -> usize {
150 self.keys[0].rank()
151 }
152
153 pub(crate) fn basek(&self) -> usize {
154 self.keys[0].basek()
155 }
156
157 #[allow(dead_code)]
158 pub(crate) fn block_size(&self) -> usize {
159 match self.dist {
160 Distribution::BinaryBlock(value) => value,
161 _ => 1,
162 }
163 }
164}
165
166impl<D: DataMut> BlindRotationKeyCGGI<D> {
167 pub fn generate_from_sk<DataSkGLWE, DataSkLWE, B: Backend>(
168 &mut self,
169 module: &Module<B>,
170 sk_glwe: &GLWESecretExec<DataSkGLWE, B>,
171 sk_lwe: &LWESecret<DataSkLWE>,
172 source_xa: &mut Source,
173 source_xe: &mut Source,
174 sigma: f64,
175 scratch: &mut Scratch<B>,
176 ) where
177 DataSkGLWE: DataRef,
178 DataSkLWE: DataRef,
179 Module<B>: GGSWEncryptSkFamily<B> + VecZnxAddScalarInplace,
180 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
181 {
182 #[cfg(debug_assertions)]
183 {
184 assert_eq!(self.keys.len(), sk_lwe.n());
185 assert!(sk_glwe.n() <= module.n());
186 assert_eq!(sk_glwe.rank(), self.keys[0].rank());
187 match sk_lwe.dist {
188 Distribution::BinaryBlock(_)
189 | Distribution::BinaryFixed(_)
190 | Distribution::BinaryProb(_)
191 | Distribution::ZERO => {}
192 _ => panic!(
193 "invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
194 ),
195 }
196 }
197
198 self.dist = sk_lwe.dist;
199
200 let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n(), 1);
201 let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data.to_ref();
202
203 self.keys.iter_mut().enumerate().for_each(|(i, ggsw)| {
204 pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
205 ggsw.encrypt_sk(module, &pt, sk_glwe, source_xa, source_xe, sigma, scratch);
206 });
207 }
208}
209
210#[derive(PartialEq, Eq)]
211pub struct BlindRotationKeyCGGIExec<D: Data, B: Backend> {
212 pub(crate) data: Vec<GGSWCiphertextExec<D, B>>,
213 pub(crate) dist: Distribution,
214 pub(crate) x_pow_a: Option<Vec<SvpPPol<Vec<u8>, B>>>,
215}
216
217impl<D: Data, B: Backend> BlindRotationKeyCGGIExec<D, B> {
218 #[allow(dead_code)]
219 pub(crate) fn n(&self) -> usize {
220 self.data[0].n()
221 }
222
223 #[allow(dead_code)]
224 pub(crate) fn rows(&self) -> usize {
225 self.data[0].rows()
226 }
227
228 #[allow(dead_code)]
229 pub(crate) fn k(&self) -> usize {
230 self.data[0].k()
231 }
232
233 #[allow(dead_code)]
234 pub(crate) fn size(&self) -> usize {
235 self.data[0].size()
236 }
237
238 #[allow(dead_code)]
239 pub(crate) fn rank(&self) -> usize {
240 self.data[0].rank()
241 }
242
243 pub(crate) fn basek(&self) -> usize {
244 self.data[0].basek()
245 }
246
247 pub(crate) fn block_size(&self) -> usize {
248 match self.dist {
249 Distribution::BinaryBlock(value) => value,
250 _ => 1,
251 }
252 }
253}
254
255pub trait BlindRotationKeyCGGIExecLayoutFamily<B: Backend> = SvpPPolAlloc<B> + SvpPrepare<B>;
256
257impl<B: Backend> BlindRotationKeyCGGIExec<Vec<u8>, B> {
258 pub fn alloc(module: &Module<B>, n_glwe: usize, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self
259 where
260 Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B> + VmpPMatAlloc<B> + VmpPMatPrepare<B>,
261 {
262 let mut data: Vec<GGSWCiphertextExec<Vec<u8>, B>> = Vec::with_capacity(n_lwe);
263 (0..n_lwe).for_each(|_| {
264 data.push(GGSWCiphertextExec::alloc(
265 module, n_glwe, basek, k, rows, 1, rank,
266 ))
267 });
268 Self {
269 data,
270 dist: Distribution::NONE,
271 x_pow_a: None,
272 }
273 }
274
275 pub fn from<DataOther>(module: &Module<B>, other: &BlindRotationKeyCGGI<DataOther>, scratch: &mut Scratch<B>) -> Self
276 where
277 DataOther: DataRef,
278 Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B> + VmpPMatAlloc<B> + VmpPMatPrepare<B>,
279 {
280 let mut brk: BlindRotationKeyCGGIExec<Vec<u8>, B> = Self::alloc(
281 module,
282 other.n(),
283 other.keys.len(),
284 other.basek(),
285 other.k(),
286 other.rows(),
287 other.rank(),
288 );
289 brk.prepare(module, other, scratch);
290 brk
291 }
292}
293
294impl<D: DataMut, B: Backend> BlindRotationKeyCGGIExec<D, B> {
295 pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &BlindRotationKeyCGGI<DataOther>, scratch: &mut Scratch<B>)
296 where
297 DataOther: DataRef,
298 Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B> + VmpPMatAlloc<B> + VmpPMatPrepare<B>,
299 {
300 #[cfg(debug_assertions)]
301 {
302 assert_eq!(self.data.len(), other.keys.len());
303 }
304
305 let n: usize = other.n();
306
307 self.data
308 .iter_mut()
309 .zip(other.keys.iter())
310 .for_each(|(ggsw_exec, other)| {
311 ggsw_exec.prepare(module, other, scratch);
312 });
313
314 self.dist = other.dist;
315
316 match other.dist {
317 Distribution::BinaryBlock(_) => {
318 let mut x_pow_a: Vec<SvpPPol<Vec<u8>, B>> = Vec::with_capacity(n << 1);
319 let mut buf: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
320 (0..n << 1).for_each(|i| {
321 let mut res: SvpPPol<Vec<u8>, B> = module.svp_ppol_alloc(n, 1);
322 set_xai_plus_y(module, i, 0, &mut res, &mut buf);
323 x_pow_a.push(res);
324 });
325 self.x_pow_a = Some(x_pow_a);
326 }
327 _ => {}
328 }
329 }
330}
331
332pub fn set_xai_plus_y<A, C, B: Backend>(module: &Module<B>, ai: usize, y: i64, res: &mut SvpPPol<A, B>, buf: &mut ScalarZnx<C>)
333where
334 A: DataMut,
335 C: DataMut,
336 Module<B>: SvpPrepare<B>,
337{
338 let n: usize = res.n();
339
340 {
341 let raw: &mut [i64] = buf.at_mut(0, 0);
342 if ai < n {
343 raw[ai] = 1;
344 } else {
345 raw[(ai - n) & (n - 1)] = -1;
346 }
347 raw[0] += y;
348 }
349
350 module.svp_prepare(res, 0, buf, 0);
351
352 {
353 let raw: &mut [i64] = buf.at_mut(0, 0);
354
355 if ai < n {
356 raw[ai] = 0;
357 } else {
358 raw[(ai - n) & (n - 1)] = 0;
359 }
360 raw[0] = 0;
361 }
362}