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