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