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