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