poulpy/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs
poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rsRaw
1use poulpy_core::{
2 GLWEAdd, GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEKeyswitch, GLWENoise, GLWEPacking, GLWERotate, GLWESub, GLWETrace,
3 LWEFromGLWE, ScratchTakeCore,
4 layouts::{
5 Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEAutomorphismKeyHelper, GLWEInfos, GLWELayout,
6 GLWEPlaintextLayout, GLWESecretPreparedToRef, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos, LWEToMut, Rank,
7 TorusPrecision,
8 },
9};
10use poulpy_hal::{
11 api::ModuleLogN,
12 layouts::{Backend, Data, DataMut, DataRef, Scratch, Stats},
13 source::Source,
14};
15use std::{collections::HashMap, marker::PhantomData};
16
17use crate::tfhe::bdd_arithmetic::{Cmux, FheUintPrepared, FromBits, GetGGSWBit, ToBits, UnsignedInteger};
18
19/// An FHE ciphertext encrypting the bits of an [UnsignedInteger].
20pub struct FheUint<D: Data, T: UnsignedInteger> {
21 pub(crate) bits: GLWE<D>,
22 pub(crate) _phantom: PhantomData<T>,
23}
24
25impl<T: UnsignedInteger> FheUint<Vec<u8>, T> {
26 pub fn alloc_from_infos<A>(infos: &A) -> Self
27 where
28 A: GLWEInfos,
29 {
30 Self::alloc(infos.n(), infos.base2k(), infos.k(), infos.rank())
31 }
32
33 pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
34 Self {
35 bits: GLWE::alloc(n, base2k, k, rank),
36 _phantom: PhantomData,
37 }
38 }
39}
40
41impl<'a, T: UnsignedInteger> FheUint<&'a mut [u8], T> {
42 pub fn from_glwe_to_mut<G>(glwe: &'a mut G) -> Self
43 where
44 G: GLWEToMut,
45 {
46 FheUint {
47 bits: glwe.to_mut(),
48 _phantom: PhantomData,
49 }
50 }
51}
52
53impl<'a, T: UnsignedInteger> FheUint<&'a [u8], T> {
54 pub fn from_glwe_to_ref<G>(glwe: &'a G) -> Self
55 where
56 G: GLWEToRef,
57 {
58 FheUint {
59 bits: glwe.to_ref(),
60 _phantom: PhantomData,
61 }
62 }
63}
64
65impl<D: DataRef, T: UnsignedInteger> LWEInfos for FheUint<D, T> {
66 fn base2k(&self) -> poulpy_core::layouts::Base2K {
67 self.bits.base2k()
68 }
69
70 fn k(&self) -> poulpy_core::layouts::TorusPrecision {
71 self.bits.k()
72 }
73
74 fn n(&self) -> poulpy_core::layouts::Degree {
75 self.bits.n()
76 }
77}
78
79impl<D: DataRef, T: UnsignedInteger> GLWEInfos for FheUint<D, T> {
80 fn rank(&self) -> poulpy_core::layouts::Rank {
81 self.bits.rank()
82 }
83}
84
85impl<D: DataMut, T: UnsignedInteger + ToBits> FheUint<D, T> {
86 pub fn encrypt_sk<S, M, BE: Backend>(
87 &mut self,
88 module: &M,
89 data: T,
90 sk: &S,
91 source_xa: &mut Source,
92 source_xe: &mut Source,
93 scratch: &mut Scratch<BE>,
94 ) where
95 S: GLWESecretPreparedToRef<BE> + GLWEInfos,
96 M: ModuleLogN + GLWEEncryptSk<BE>,
97 Scratch<BE>: ScratchTakeCore<BE>,
98 {
99 #[cfg(debug_assertions)]
100 {
101 assert!(module.n().is_multiple_of(T::BITS as usize));
102 assert_eq!(self.n(), module.n() as u32);
103 assert_eq!(sk.n(), module.n() as u32);
104 }
105
106 let mut data_bits: Vec<i64> = vec![0i64; module.n()];
107
108 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
109
110 // Interleaves bytes
111 for i in 0..T::BITS as usize {
112 data_bits[T::bit_index(i) << log_gap] = data.bit(i) as i64
113 }
114
115 let pt_infos = GLWEPlaintextLayout {
116 n: self.n(),
117 base2k: self.base2k(),
118 k: 2_usize.into(),
119 };
120
121 let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
122
123 pt.encode_vec_i64(&data_bits, TorusPrecision(2));
124 self.bits
125 .encrypt_sk(module, &pt, sk, source_xa, source_xe, scratch_1);
126 }
127}
128
129impl<D: DataRef, T: UnsignedInteger + FromBits> FheUint<D, T> {
130 pub fn noise<S, M, BE: Backend>(&self, module: &M, want: u32, sk: &S, scratch: &mut Scratch<BE>) -> Stats
131 where
132 S: GLWESecretPreparedToRef<BE> + GLWEInfos,
133 M: ModuleLogN + GLWEDecrypt<BE> + GLWENoise<BE>,
134 Scratch<BE>: ScratchTakeCore<BE>,
135 {
136 #[cfg(debug_assertions)]
137 {
138 assert!(module.n().is_multiple_of(T::BITS as usize));
139 assert_eq!(self.n(), module.n() as u32);
140 assert_eq!(sk.n(), module.n() as u32);
141 }
142
143 let (mut pt, scratch_1) = scratch.take_glwe_plaintext(self);
144 let mut data_bits = vec![0i64; module.n()];
145 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
146 for i in 0..T::BITS as usize {
147 data_bits[T::bit_index(i) << log_gap] = want.bit(i) as i64
148 }
149 pt.encode_vec_i64(&data_bits, TorusPrecision(2));
150 self.bits.noise(module, &pt, sk, scratch_1)
151 }
152
153 pub fn decrypt<S, M, BE: Backend>(&self, module: &M, sk: &S, scratch: &mut Scratch<BE>) -> T
154 where
155 S: GLWESecretPreparedToRef<BE> + GLWEInfos,
156 M: ModuleLogN + GLWEDecrypt<BE>,
157 Scratch<BE>: ScratchTakeCore<BE>,
158 {
159 #[cfg(debug_assertions)]
160 {
161 assert!(module.n().is_multiple_of(T::BITS as usize));
162 assert_eq!(self.n(), module.n() as u32);
163 assert_eq!(sk.n(), module.n() as u32);
164 }
165
166 let pt_infos = GLWEPlaintextLayout {
167 n: self.n(),
168 base2k: self.base2k(),
169 k: 1_usize.into(),
170 };
171
172 let (mut pt, scratch_1) = scratch.take_glwe_plaintext(&pt_infos);
173
174 self.bits.decrypt(module, &mut pt, sk, scratch_1);
175
176 let mut data_bits: Vec<i64> = vec![0i64; module.n()];
177 pt.decode_vec_i64(&mut data_bits, TorusPrecision(2));
178
179 let mut bits: Vec<u8> = vec![0u8; T::BITS as usize];
180
181 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
182
183 // Retrives from interleaved bytes
184 for i in 0..T::BITS as usize {
185 bits[i] = data_bits[T::bit_index(i) << log_gap] as u8
186 }
187
188 T::from_bits(&bits)
189 }
190}
191
192impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
193 /// Packs Vec<GLWE(bit[i])> into [FheUint].
194 pub fn pack<G, M, K, H, BE: Backend>(&mut self, module: &M, mut bits: Vec<G>, keys: &H, scratch: &mut Scratch<BE>)
195 where
196 G: GLWEToMut + GLWEInfos,
197 M: ModuleLogN + GLWEPacking<BE> + GLWECopy,
198 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
199 H: GLWEAutomorphismKeyHelper<K, BE>,
200 Scratch<BE>: ScratchTakeCore<BE>,
201 {
202 // Repacks the GLWE ciphertexts bits
203 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
204 let mut cts: HashMap<usize, &mut G> = HashMap::new();
205 for (i, ct) in bits.iter_mut().enumerate().take(T::BITS as usize) {
206 cts.insert(T::bit_index(i) << log_gap, ct);
207 }
208
209 module.glwe_pack(&mut self.bits, cts, log_gap, keys, scratch);
210 }
211
212 #[allow(clippy::too_many_arguments)]
213 // Self <- ((a.rotate_right(dst<<4) & 0xFFFF_0000) | (b.rotate_right(src<<4) & 0x0000_FFFF)).rotate_left(dst<<4);
214 pub fn splice_u16<A, B, H, K, M, BE: Backend>(
215 &mut self,
216 module: &M,
217 dst: usize,
218 src: usize,
219 a: &A,
220 b: &B,
221 keys: &H,
222 scratch: &mut Scratch<BE>,
223 ) where
224 A: GLWEToRef + GLWEInfos,
225 B: GLWEToRef + GLWEInfos,
226 H: GLWEAutomorphismKeyHelper<K, BE>,
227 K: GGLWEPreparedToRef<BE> + GGLWEInfos + GetGaloisElement,
228 M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWECopy,
229 Scratch<BE>: ScratchTakeBDD<T, BE>,
230 {
231 assert!(dst < (T::BITS >> 4) as usize);
232 assert!(src < (T::BITS >> 4) as usize);
233
234 let (mut tmp, scratch_1) = scratch.take_fhe_uint(self);
235 tmp.splice_u8(module, dst << 1, src << 1, a, b, keys, scratch_1);
236 self.splice_u8(
237 module,
238 (dst << 1) + 1,
239 (src << 1) + 1,
240 &tmp,
241 b,
242 keys,
243 scratch_1,
244 );
245 }
246
247 #[allow(clippy::too_many_arguments)]
248 // Self <- ((a.rotate_right(dst<<3) & 0xFFFF_FF00) | (b.rotate_right(src<<3) & 0x0000_00FF)).rotate_left(dst<<3);
249 pub fn splice_u8<A, B, H, K, M, BE: Backend>(
250 &mut self,
251 module: &M,
252 dst: usize,
253 src: usize,
254 a: &A,
255 b: &B,
256 keys: &H,
257 scratch: &mut Scratch<BE>,
258 ) where
259 A: GLWEToRef + GLWEInfos,
260 B: GLWEToRef + GLWEInfos,
261 H: GLWEAutomorphismKeyHelper<K, BE>,
262 K: GGLWEPreparedToRef<BE> + GGLWEInfos + GetGaloisElement,
263 M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWECopy,
264 Scratch<BE>: ScratchTakeBDD<T, BE>,
265 {
266 assert!(dst < (T::BITS >> 3) as usize);
267 assert!(src < (T::BITS >> 3) as usize);
268
269 // 1) Zero the byte receiver
270 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
271 let trace_start = (T::LOG_BITS - T::LOG_BYTES) as usize;
272 let rot: i64 = (T::bit_index(dst << 3) << log_gap) as i64;
273
274 module.glwe_copy(self, a);
275
276 self.zero_byte(module, dst, keys, scratch);
277
278 // Isolate the byte to transfer from a
279 let (mut tmp_fhe_uint_byte, scratch_1) = scratch.take_fhe_uint(b);
280
281 // Move a[byte_a] into a[dst]
282 module.glwe_rotate(
283 -((T::bit_index(src << 3) << log_gap) as i64),
284 &mut tmp_fhe_uint_byte,
285 b,
286 );
287
288 // Zeroes all other bytes
289 module.glwe_trace_inplace(&mut tmp_fhe_uint_byte, trace_start, keys, scratch_1);
290
291 // Moves back self[0] to self[byte_tg]
292 module.glwe_rotate_inplace(rot, &mut tmp_fhe_uint_byte, scratch_1);
293
294 // Add self[0] += a[0]
295 module.glwe_add_inplace(&mut self.bits, &tmp_fhe_uint_byte);
296 }
297}
298
299impl<D: DataMut, T: UnsignedInteger> GLWEToMut for FheUint<D, T> {
300 fn to_mut(&mut self) -> GLWE<&mut [u8]> {
301 self.bits.to_mut()
302 }
303}
304
305pub trait ScratchTakeBDD<T: UnsignedInteger, BE: Backend>
306where
307 Self: ScratchTakeCore<BE>,
308{
309 fn take_fhe_uint<A>(&mut self, infos: &A) -> (FheUint<&mut [u8], T>, &mut Self)
310 where
311 A: GLWEInfos,
312 {
313 let (glwe, scratch) = self.take_glwe(infos);
314 (
315 FheUint {
316 bits: glwe,
317 _phantom: PhantomData,
318 },
319 scratch,
320 )
321 }
322}
323
324impl<T: UnsignedInteger, BE: Backend> ScratchTakeBDD<T, BE> for Scratch<BE> where Self: ScratchTakeCore<BE> {}
325
326impl<D: DataRef, T: UnsignedInteger> FheUint<D, T> {
327 pub fn get_bit_lwe<R, KGLWE, KLWE, M, BE: Backend>(
328 &self,
329 module: &M,
330 bit: usize,
331 res: &mut R,
332 ks_glwe: Option<&KGLWE>,
333 ks_lwe: &KLWE,
334 scratch: &mut Scratch<BE>,
335 ) where
336 R: LWEToMut,
337 KGLWE: GGLWEPreparedToRef<BE> + GGLWEInfos,
338 KLWE: GGLWEPreparedToRef<BE> + GGLWEInfos,
339 M: ModuleLogN + LWEFromGLWE<BE> + GLWERotate<BE> + GLWEKeyswitch<BE>,
340 Scratch<BE>: ScratchTakeCore<BE>,
341 {
342 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
343 if let Some(ks_glwe) = ks_glwe {
344 let (mut res_tmp, scratch_1) = scratch.take_glwe(&GLWELayout {
345 n: self.n(),
346 base2k: ks_lwe.base2k(),
347 k: ks_lwe.k().min(self.k()),
348 rank: ks_lwe.rank_out(),
349 });
350 module.glwe_keyswitch(&mut res_tmp, self, ks_glwe, scratch_1);
351 res.to_mut().from_glwe(
352 module,
353 &res_tmp,
354 T::bit_index(bit) << log_gap,
355 ks_lwe,
356 scratch_1,
357 );
358 } else {
359 res.to_mut()
360 .from_glwe(module, self, T::bit_index(bit) << log_gap, ks_lwe, scratch);
361 }
362 }
363
364 pub fn get_bit_glwe<R, K, M, H, BE: Backend>(&self, module: &M, bit: usize, res: &mut R, keys: &H, scratch: &mut Scratch<BE>)
365 where
366 R: GLWEToMut,
367 K: GGLWEPreparedToRef<BE> + GGLWEInfos,
368 M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE>,
369 H: GLWEAutomorphismKeyHelper<K, BE>,
370 K: GGLWEPreparedToRef<BE> + GGLWEInfos + GetGaloisElement,
371 Scratch<BE>: ScratchTakeCore<BE>,
372 {
373 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
374 let rot = (T::bit_index(bit) << log_gap) as i64;
375 module.glwe_rotate(-rot, res, self);
376 module.glwe_trace_inplace(res, 0, keys, scratch);
377 }
378
379 pub fn get_byte<R, K, M, H, BE: Backend>(&self, module: &M, byte: usize, res: &mut R, keys: &H, scratch: &mut Scratch<BE>)
380 where
381 R: GLWEToMut,
382 K: GGLWEPreparedToRef<BE> + GGLWEInfos,
383 M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE>,
384 H: GLWEAutomorphismKeyHelper<K, BE>,
385 K: GGLWEPreparedToRef<BE> + GGLWEInfos + GetGaloisElement,
386 Scratch<BE>: ScratchTakeCore<BE>,
387 {
388 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
389 let trace_start = (T::LOG_BITS - T::LOG_BYTES) as usize;
390 let rot = (T::bit_index(byte << 3) << log_gap) as i64;
391 module.glwe_rotate(-rot, res, self);
392 module.glwe_trace_inplace(res, trace_start, keys, scratch);
393 }
394}
395
396impl<D: DataRef, T: UnsignedInteger> GLWEToRef for FheUint<D, T> {
397 fn to_ref(&self) -> GLWE<&[u8]> {
398 self.bits.to_ref()
399 }
400}
401
402impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
403 pub fn from_fhe_uint_prepared<M, DR, H, K, BE: Backend>(
404 &mut self,
405 module: &M,
406 other: &FheUintPrepared<DR, T, BE>,
407 keys: &H,
408 scratch: &mut Scratch<BE>,
409 ) where
410 DR: DataRef,
411 M: Cmux<BE> + ModuleLogN + GLWEPacking<BE> + GLWECopy,
412 Scratch<BE>: ScratchTakeCore<BE>,
413 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
414 H: GLWEAutomorphismKeyHelper<K, BE>,
415 {
416 let zero: GLWE<Vec<u8>> = GLWE::alloc_from_infos(self);
417 let mut one: GLWE<Vec<u8>> = GLWE::alloc_from_infos(self);
418 one.data_mut()
419 .encode_coeff_i64(self.base2k().into(), 0, 2, 0, 1);
420
421 let (mut out_bits, scratch_1) = scratch.take_glwe_slice(T::BITS as usize, self);
422
423 for (i, bits) in out_bits.iter_mut().enumerate().take(T::BITS as usize) {
424 module.cmux(bits, &one, &zero, &other.get_bit(i), scratch_1);
425 }
426
427 self.pack(module, out_bits, keys, scratch_1);
428 }
429
430 pub fn zero_byte<M, K, H, BE: Backend>(&mut self, module: &M, byte: usize, keys: &H, scratch: &mut Scratch<BE>)
431 where
432 H: GLWEAutomorphismKeyHelper<K, BE>,
433 K: GGLWEPreparedToRef<BE> + GGLWEInfos + GetGaloisElement,
434 M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWECopy,
435 Scratch<BE>: ScratchTakeBDD<T, BE>,
436 {
437 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
438 let trace_start = (T::LOG_BITS - T::LOG_BYTES) as usize;
439 let rot: i64 = (T::bit_index(byte << 3) << log_gap) as i64;
440
441 // Move a to self and align byte
442 module.glwe_rotate_inplace(-rot, &mut self.bits, scratch);
443
444 // Stores this byte (everything else zeroed) into tmp_trace
445 let (mut tmp_trace, scratch_1) = scratch.take_glwe(self);
446 module.glwe_trace(&mut tmp_trace, trace_start, self, keys, scratch_1);
447
448 // Subtracts to self to zero it
449 module.glwe_sub_inplace(&mut self.bits, &tmp_trace);
450
451 // Move a to self and align byte
452 module.glwe_rotate_inplace(rot, &mut self.bits, scratch);
453 }
454
455 pub fn sext<M, H, K, BE>(&mut self, module: &M, byte: usize, keys: &H, scratch: &mut Scratch<BE>)
456 where
457 M:,
458 H: GLWEAutomorphismKeyHelper<K, BE>,
459 K: GGLWEPreparedToRef<BE> + GGLWEInfos + GetGaloisElement,
460 BE: Backend,
461 M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWEAdd + GLWESub + GLWECopy,
462 Scratch<BE>: ScratchTakeCore<BE>,
463 {
464 assert!(byte < (1 << T::LOG_BYTES));
465
466 let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
467 let rot: i64 = (T::bit_index((byte << 3) + 7) << log_gap) as i64;
468
469 let (mut sext, scratch_1) = scratch.take_glwe(self);
470
471 // Extract MSB
472 module.glwe_rotate(-rot, &mut sext, &self.bits);
473 module.glwe_trace_inplace(&mut sext, 0, keys, scratch_1);
474
475 // Replicates MSB in byte
476 for i in 0..3 {
477 let (mut tmp, _) = scratch_1.take_glwe(self);
478 module.glwe_rotate(((1 << T::LOG_BYTES) << log_gap) << i, &mut tmp, &sext);
479 module.glwe_add_inplace(&mut sext, &tmp);
480 }
481
482 // Splice sext
483 let (mut tmp, scratch_2) = scratch_1.take_glwe(self);
484 for i in (byte + 1)..(1 << T::LOG_BYTES) as usize {
485 FheUint::<&mut [u8], T>::from_glwe_to_mut(&mut tmp).splice_u8(module, i, 0, &self.bits, &sext, keys, scratch_2);
486 module.glwe_copy(&mut self.bits, &tmp);
487 }
488 }
489}