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