1use std::marker::PhantomData;
  2use std::thread;
  3
  4use poulpy_core::layouts::{
  5    Base2K, Dnum, Dsize, GGSWInfos, GGSWPreparedFactory, GLWEInfos, LWEInfos, Rank, TorusPrecision, prepared::GGSWPrepared,
  6};
  7use poulpy_core::layouts::{
  8    GGLWEInfos, GGLWEPreparedToRef, GGSW, GGSWLayout, GGSWPreparedToMut, GGSWPreparedToRef, GLWEAutomorphismKeyHelper,
  9    GetGaloisElement, LWE,
 10};
 11use poulpy_core::{GLWECopy, GLWEDecrypt, GLWEPacking, LWEFromGLWE};
 12
 13use poulpy_core::{GGSWEncryptSk, ScratchTakeCore, layouts::GLWESecretPreparedToRef};
 14use poulpy_hal::api::{ModuleLogN, ScratchAvailable, ScratchFromBytes};
 15use poulpy_hal::layouts::{Backend, Data, DataRef, Module};
 16
 17use poulpy_hal::{
 18    api::ModuleN,
 19    layouts::{DataMut, Scratch},
 20    source::Source,
 21};
 22
 23use crate::tfhe::bdd_arithmetic::{
 24    BDDKey, BDDKeyHelper, BDDKeyInfos, BDDKeyPrepared, BDDKeyPreparedFactory, BitSize, FheUint, ToBits,
 25};
 26use crate::tfhe::bdd_arithmetic::{Cmux, FromBits, ScratchTakeBDD, UnsignedInteger};
 27use crate::tfhe::blind_rotation::BlindRotationAlgo;
 28use crate::tfhe::circuit_bootstrapping::{CircuitBootstrappingKeyInfos, CirtuitBootstrappingExecute};
 29
 30/// A prepared FHE ciphertext encrypting the bits of an [UnsignedInteger].
 31pub struct FheUintPrepared<D: Data, T: UnsignedInteger, B: Backend> {
 32    pub(crate) bits: Vec<GGSWPrepared<D, B>>,
 33    pub(crate) _phantom: PhantomData<T>,
 34}
 35
 36impl<T: UnsignedInteger, BE: Backend> FheUintPreparedFactory<T, BE> for Module<BE> where Self: Sized + GGSWPreparedFactory<BE> {}
 37
 38pub trait GetGGSWBit<BE: Backend>: Sync {
 39    fn get_bit(&self, bit: usize) -> GGSWPrepared<&[u8], BE>;
 40}
 41
 42impl<D: DataRef, T: UnsignedInteger, BE: Backend> GetGGSWBit<BE> for FheUintPrepared<D, T, BE> {
 43    fn get_bit(&self, bit: usize) -> GGSWPrepared<&[u8], BE> {
 44        assert!(bit <= self.bits.len());
 45        self.bits[bit].to_ref()
 46    }
 47}
 48
 49pub trait GetGGSWBitMut<T: UnsignedInteger, BE: Backend> {
 50    fn get_bit(&mut self, bit: usize) -> GGSWPrepared<&mut [u8], BE>;
 51}
 52
 53impl<D: DataMut, T: UnsignedInteger, BE: Backend> GetGGSWBitMut<T, BE> for FheUintPrepared<D, T, BE> {
 54    fn get_bit(&mut self, bit: usize) -> GGSWPrepared<&mut [u8], BE> {
 55        assert!(bit <= self.bits.len());
 56        self.bits[bit].to_mut()
 57    }
 58}
 59
 60impl<D: Data, T: UnsignedInteger, BE: Backend> BitSize for FheUintPrepared<D, T, BE> {
 61    fn bit_size(&self) -> usize {
 62        T::BITS as usize
 63    }
 64}
 65
 66pub trait FheUintPreparedFactory<T: UnsignedInteger, BE: Backend>
 67where
 68    Self: Sized + GGSWPreparedFactory<BE>,
 69{
 70    fn alloc_fhe_uint_prepared(
 71        &self,
 72        base2k: Base2K,
 73        k: TorusPrecision,
 74        dnum: Dnum,
 75        dsize: Dsize,
 76        rank: Rank,
 77    ) -> FheUintPrepared<Vec<u8>, T, BE> {
 78        FheUintPrepared {
 79            bits: (0..T::BITS)
 80                .map(|_| GGSWPrepared::alloc(self, base2k, k, dnum, dsize, rank))
 81                .collect(),
 82            _phantom: PhantomData,
 83        }
 84    }
 85
 86    fn alloc_fhe_uint_prepared_from_infos<A>(&self, infos: &A) -> FheUintPrepared<Vec<u8>, T, BE>
 87    where
 88        A: GGSWInfos,
 89    {
 90        self.alloc_fhe_uint_prepared(
 91            infos.base2k(),
 92            infos.k(),
 93            infos.dnum(),
 94            infos.dsize(),
 95            infos.rank(),
 96        )
 97    }
 98}
 99
100impl<T: UnsignedInteger, BE: Backend> FheUintPrepared<Vec<u8>, T, BE> {
101    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
102    where
103        A: GGSWInfos,
104        M: FheUintPreparedFactory<T, BE>,
105    {
106        module.alloc_fhe_uint_prepared_from_infos(infos)
107    }
108
109    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
110    where
111        M: FheUintPreparedFactory<T, BE>,
112    {
113        module.alloc_fhe_uint_prepared(base2k, k, dnum, dsize, rank)
114    }
115}
116
117impl<T: UnsignedInteger + ToBits, BE: Backend> FheUintPreparedEncryptSk<T, BE> for Module<BE> where
118    Self: Sized + ModuleN + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>
119{
120}
121
122pub trait FheUintPreparedEncryptSk<T: UnsignedInteger + ToBits, BE: Backend>
123where
124    Self: Sized + ModuleN + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>,
125{
126    fn fhe_uint_prepared_encrypt_sk<DM, S>(
127        &self,
128        res: &mut FheUintPrepared<DM, T, BE>,
129        value: T,
130        sk: &S,
131        source_xa: &mut Source,
132        source_xe: &mut Source,
133        scratch: &mut Scratch<BE>,
134    ) where
135        DM: DataMut,
136        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
137        Scratch<BE>: ScratchTakeCore<BE>,
138    {
139        use poulpy_hal::{api::ScratchTakeBasic, layouts::ZnxZero};
140
141        assert!(self.n().is_multiple_of(T::BITS as usize));
142        assert_eq!(res.n(), self.n() as u32);
143        assert_eq!(sk.n(), self.n() as u32);
144
145        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
146        let (mut pt, scratch_2) = scratch_1.take_scalar_znx(self.n(), 1);
147        pt.zero();
148
149        for i in 0..T::BITS as usize {
150            use poulpy_hal::layouts::ZnxViewMut;
151            pt.at_mut(0, 0)[0] = value.bit(i) as i64;
152            tmp_ggsw.encrypt_sk(self, &pt, sk, source_xa, source_xe, scratch_2);
153            res.bits[i].prepare(self, &tmp_ggsw, scratch_2);
154        }
155    }
156}
157
158impl<D: DataMut, T: UnsignedInteger + ToBits, BE: Backend> FheUintPrepared<D, T, BE> {
159    pub fn encrypt_sk<M, S>(
160        &mut self,
161        module: &M,
162        value: T,
163        sk: &S,
164        source_xa: &mut Source,
165        source_xe: &mut Source,
166        scratch: &mut Scratch<BE>,
167    ) where
168        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
169        M: FheUintPreparedEncryptSk<T, BE>,
170        Scratch<BE>: ScratchTakeCore<BE>,
171    {
172        module.fhe_uint_prepared_encrypt_sk(self, value, sk, source_xa, source_xe, scratch);
173    }
174}
175
176impl<D: DataRef, T: UnsignedInteger + FromBits, BE: Backend> FheUintPrepared<D, T, BE> {
177    pub fn decrypt<M, S, H, K>(&self, module: &M, sk: &S, keys: &H, scratch: &mut Scratch<BE>) -> T
178    where
179        M: ModuleLogN + GLWEDecrypt<BE> + Cmux<BE> + GLWEPacking<BE> + GLWECopy,
180        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
181        Scratch<BE>: ScratchTakeCore<BE>,
182        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
183        H: GLWEAutomorphismKeyHelper<K, BE>,
184    {
185        let (mut tmp, scratch_1) = scratch.take_fhe_uint(self);
186        tmp.from_fhe_uint_prepared(module, self, keys, scratch_1);
187        tmp.decrypt(module, sk, scratch_1)
188    }
189}
190
191impl<D: DataRef, T: UnsignedInteger, B: Backend> LWEInfos for FheUintPrepared<D, T, B> {
192    fn base2k(&self) -> poulpy_core::layouts::Base2K {
193        self.bits[0].base2k()
194    }
195
196    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
197        self.bits[0].k()
198    }
199
200    fn n(&self) -> poulpy_core::layouts::Degree {
201        self.bits[0].n()
202    }
203}
204
205impl<D: DataRef, T: UnsignedInteger, B: Backend> GLWEInfos for FheUintPrepared<D, T, B> {
206    fn rank(&self) -> poulpy_core::layouts::Rank {
207        self.bits[0].rank()
208    }
209}
210
211impl<D: DataRef, T: UnsignedInteger, B: Backend> GGSWInfos for FheUintPrepared<D, T, B> {
212    fn dsize(&self) -> poulpy_core::layouts::Dsize {
213        self.bits[0].dsize()
214    }
215
216    fn dnum(&self) -> poulpy_core::layouts::Dnum {
217        self.bits[0].dnum()
218    }
219}
220
221impl<D: DataMut, BRA: BlindRotationAlgo, BE: Backend> BDDKeyPrepared<D, BRA, BE> {
222    pub fn prepare<DR, M>(&mut self, module: &M, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
223    where
224        DR: DataRef,
225        M: BDDKeyPreparedFactory<BRA, BE>,
226        Scratch<BE>: ScratchTakeCore<BE>,
227    {
228        module.prepare_bdd_key(self, other, scratch);
229    }
230}
231
232pub trait FheUintPrepare<BRA: BlindRotationAlgo, BE: Backend> {
233    fn fhe_uint_prepare_tmp_bytes<R, A, B>(
234        &self,
235        block_size: usize,
236        extension_factor: usize,
237        res_infos: &R,
238        bits_infos: &A,
239        bdd_infos: &B,
240    ) -> usize
241    where
242        R: GGSWInfos,
243        A: GLWEInfos,
244        B: BDDKeyInfos;
245    fn fhe_uint_prepare<DM, DB, DK, K, T: UnsignedInteger>(
246        &self,
247        res: &mut FheUintPrepared<DM, T, BE>,
248        bits: &FheUint<DB, T>,
249        key: &K,
250        scratch: &mut Scratch<BE>,
251    ) where
252        DM: DataMut,
253        DB: DataRef,
254        DK: DataRef,
255        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
256        Scratch<BE>: ScratchFromBytes<BE>,
257    {
258        self.fhe_uint_prepare_custom(res, bits, 0, T::BITS as usize, key, scratch);
259    }
260    fn fhe_uint_prepare_custom<DM, DB, DK, K, T: UnsignedInteger>(
261        &self,
262        res: &mut FheUintPrepared<DM, T, BE>,
263        bits: &FheUint<DB, T>,
264        bit_start: usize,
265        bit_count: usize,
266        key: &K,
267        scratch: &mut Scratch<BE>,
268    ) where
269        DM: DataMut,
270        DB: DataRef,
271        DK: DataRef,
272        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
273    {
274        self.fhe_uint_prepare_custom_multi_thread(1, res, bits, bit_start, bit_count, key, scratch)
275    }
276    #[allow(clippy::too_many_arguments)]
277    fn fhe_uint_prepare_custom_multi_thread<DM, DB, DK, K, T: UnsignedInteger>(
278        &self,
279        threads: usize,
280        res: &mut FheUintPrepared<DM, T, BE>,
281        bits: &FheUint<DB, T>,
282        bit_start: usize,
283        bit_count: usize,
284        key: &K,
285        scratch: &mut Scratch<BE>,
286    ) where
287        DM: DataMut,
288        DB: DataRef,
289        DK: DataRef,
290        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos;
291}
292
293impl<BRA: BlindRotationAlgo, BE: Backend> FheUintPrepare<BRA, BE> for Module<BE>
294where
295    Self: LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
296    Scratch<BE>: ScratchTakeCore<BE>,
297{
298    fn fhe_uint_prepare_tmp_bytes<R, A, B>(
299        &self,
300        block_size: usize,
301        extension_factor: usize,
302        res_infos: &R,
303        bits_infos: &A,
304        bdd_infos: &B,
305    ) -> usize
306    where
307        R: GGSWInfos,
308        A: GLWEInfos,
309        B: BDDKeyInfos,
310    {
311        self.circuit_bootstrapping_execute_tmp_bytes(
312            block_size,
313            extension_factor,
314            res_infos,
315            &bdd_infos.cbt_infos(),
316        ) + GGSW::bytes_of_from_infos(res_infos)
317            + LWE::bytes_of_from_infos(bits_infos)
318    }
319
320    fn fhe_uint_prepare_custom_multi_thread<DM, DB, DK, K, T: UnsignedInteger>(
321        &self,
322        threads: usize,
323        res: &mut FheUintPrepared<DM, T, BE>,
324        bits: &FheUint<DB, T>,
325        bit_start: usize,
326        bit_count: usize,
327        key: &K,
328        scratch: &mut Scratch<BE>,
329    ) where
330        DM: DataMut,
331        DB: DataRef,
332        DK: DataRef,
333        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
334    {
335        let bit_end = bit_start + bit_count;
336        let (cbt, ks_glwe, ks_lwe) = key.get_cbt_key();
337
338        assert!(bit_end <= T::BITS as usize);
339
340        let scratch_thread_size = self.fhe_uint_prepare_tmp_bytes(cbt.block_size(), 1, res, bits, key);
341
342        assert!(
343            scratch.available() >= threads * scratch_thread_size,
344            "scratch.available():{} < threads:{threads} * scratch_thread_size:{scratch_thread_size}",
345            scratch.available()
346        );
347
348        let chunk_size: usize = bit_count.div_ceil(threads);
349
350        let (mut scratches, _) = scratch.split_mut(threads, scratch_thread_size);
351
352        let ggsw_infos: &GGSWLayout = &res.ggsw_layout();
353
354        thread::scope(|scope| {
355            for (thread_index, (scratch_thread, res_bits_chunk)) in scratches
356                .iter_mut()
357                .zip(res.bits[bit_start..bit_end].chunks_mut(chunk_size))
358                .enumerate()
359            {
360                let start: usize = bit_start + thread_index * chunk_size;
361
362                scope.spawn(move || {
363                    let (mut tmp_ggsw, scratch_1) = scratch_thread.take_ggsw(ggsw_infos);
364                    let (mut tmp_lwe, scratch_2) = scratch_1.take_lwe(bits);
365                    for (local_bit, dst) in res_bits_chunk.iter_mut().enumerate() {
366                        bits.get_bit_lwe(
367                            self,
368                            start + local_bit,
369                            &mut tmp_lwe,
370                            ks_glwe,
371                            ks_lwe,
372                            scratch_2,
373                        );
374                        cbt.execute_to_constant(self, &mut tmp_ggsw, &tmp_lwe, 1, 1, scratch_2);
375                        dst.prepare(self, &tmp_ggsw, scratch_2);
376                    }
377                });
378            }
379        });
380
381        for i in 0..bit_start {
382            res.bits[i].zero(self);
383        }
384
385        for i in bit_end..T::BITS as usize {
386            res.bits[i].zero(self);
387        }
388    }
389}
390
391impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
392    pub fn prepare<BRA, M, O, K, DK>(&mut self, module: &M, other: &FheUint<O, T>, key: &K, scratch: &mut Scratch<BE>)
393    where
394        BRA: BlindRotationAlgo,
395        O: DataRef,
396        DK: DataRef,
397        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
398        M: FheUintPrepare<BRA, BE>,
399        Scratch<BE>: ScratchTakeCore<BE>,
400    {
401        module.fhe_uint_prepare(self, other, key, scratch);
402    }
403    pub fn prepare_custom<BRA, M, O, K, DK>(
404        &mut self,
405        module: &M,
406        other: &FheUint<O, T>,
407        bit_start: usize,
408        bit_end: usize,
409        key: &K,
410        scratch: &mut Scratch<BE>,
411    ) where
412        BRA: BlindRotationAlgo,
413        O: DataRef,
414        DK: DataRef,
415        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
416        M: FheUintPrepare<BRA, BE>,
417        Scratch<BE>: ScratchTakeCore<BE>,
418    {
419        module.fhe_uint_prepare_custom(self, other, bit_start, bit_end, key, scratch);
420    }
421
422    #[allow(clippy::too_many_arguments)]
423    pub fn prepare_custom_multi_thread<BRA, M, O, K, DK>(
424        &mut self,
425        threads: usize,
426        module: &M,
427        other: &FheUint<O, T>,
428        bit_start: usize,
429        bit_end: usize,
430        key: &K,
431        scratch: &mut Scratch<BE>,
432    ) where
433        BRA: BlindRotationAlgo,
434        O: DataRef,
435        DK: DataRef,
436        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
437        M: FheUintPrepare<BRA, BE>,
438        Scratch<BE>: ScratchTakeCore<BE>,
439    {
440        module.fhe_uint_prepare_custom_multi_thread(threads, self, other, bit_start, bit_end, key, scratch);
441    }
442}