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