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> {
 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>(&self, block_size: usize, extension_factor: usize, res_infos: &R, bits_infos: &A, bdd_infos: &B) -> usize
226    where
227        R: GGSWInfos,
228        A: GLWEInfos,
229        B: BDDKeyInfos;
230    fn fhe_uint_prepare<DM, DB, DK, K, T: UnsignedInteger>(
231        &self,
232        res: &mut FheUintPrepared<DM, T, BE>,
233        bits: &FheUint<DB, T>,
234        key: &K,
235        scratch: &mut Scratch<BE>,
236    ) where
237        DM: DataMut,
238        DB: DataRef,
239        DK: DataRef,
240        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
241        Scratch<BE>: ScratchFromBytes<BE>,
242    {
243        self.fhe_uint_prepare_custom(res, bits, 0, T::BITS as usize, key, scratch);
244    }
245    fn fhe_uint_prepare_custom<DM, DB, DK, K, T: UnsignedInteger>(
246        &self,
247        res: &mut FheUintPrepared<DM, T, BE>,
248        bits: &FheUint<DB, T>,
249        bit_start: usize,
250        bit_count: usize,
251        key: &K,
252        scratch: &mut Scratch<BE>,
253    ) where
254        DM: DataMut,
255        DB: DataRef,
256        DK: DataRef,
257        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
258    {
259        self.fhe_uint_prepare_custom_multi_thread(1, res, bits, bit_start, bit_count, key, scratch)
260    }
261    fn fhe_uint_prepare_custom_multi_thread<DM, DB, DK, K, T: UnsignedInteger>(
262        &self,
263        threads: usize,
264        res: &mut FheUintPrepared<DM, T, BE>,
265        bits: &FheUint<DB, T>,
266        bit_start: usize,
267        bit_count: usize,
268        key: &K,
269        scratch: &mut Scratch<BE>,
270    ) where
271        DM: DataMut,
272        DB: DataRef,
273        DK: DataRef,
274        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos;
275}
276
277impl<BRA: BlindRotationAlgo, BE: Backend> FheUintPrepare<BRA, BE> for Module<BE>
278where
279    Self: LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
280    Scratch<BE>: ScratchTakeCore<BE>,
281{
282    fn fhe_uint_prepare_tmp_bytes<R, A, B>(&self, block_size: usize, extension_factor: usize, res_infos: &R, bits_infos: &A, bdd_infos: &B) -> usize
283    where
284        R: GGSWInfos,
285        A: GLWEInfos,
286        B: BDDKeyInfos,
287    {
288        self.circuit_bootstrapping_execute_tmp_bytes(
289            block_size,
290            extension_factor,
291            res_infos,
292            &bdd_infos.cbt_infos(),
293        ) + GGSW::bytes_of_from_infos(res_infos)
294            + LWE::bytes_of_from_infos(bits_infos)
295    }
296
297    fn fhe_uint_prepare_custom_multi_thread<DM, DB, DK, K, T: UnsignedInteger>(
298        &self,
299        threads: usize,
300        res: &mut FheUintPrepared<DM, T, BE>,
301        bits: &FheUint<DB, T>,
302        bit_start: usize,
303        bit_count: usize,
304        key: &K,
305        mut scratch: &mut Scratch<BE>,
306    ) where
307        DM: DataMut,
308        DB: DataRef,
309        DK: DataRef,
310        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
311    {
312        let bit_end = bit_start + bit_count;
313        let (cbt, ks) = key.get_cbt_key();
314
315        assert!(bit_end <= T::BITS as usize);
316
317        let scratch_thread_size = self.fhe_uint_prepare_tmp_bytes(cbt.block_size(), 1, res, bits, key);
318
319        assert!(scratch.available() >= threads * scratch_thread_size);
320
321        // How many bits we need to process
322        let chunk_size: usize = bit_count.div_ceil(threads); // ceil division
323
324        let mut scratches = Vec::new();
325        for _ in 0..(threads - 1) {
326            let (tmp, scratch_new) = scratch.split_at_mut(scratch_thread_size);
327            scratch = scratch_new;
328            scratches.push(tmp);
329        }
330        scratches.push(scratch);
331
332        let ggsw_infos: &GGSWLayout = &res.ggsw_layout();
333
334        thread::scope(|scope| {
335            for (thread_index, (scratch_thread, res_bits_chunk)) in scratches
336                .iter_mut()
337                .zip(res.bits[bit_start..bit_end].chunks_mut(chunk_size))
338                .enumerate()
339            {
340                let start: usize = bit_start + thread_index * chunk_size;
341
342                scope.spawn(move || {
343                    let (mut tmp_ggsw, scratch_1) = scratch_thread.take_ggsw(ggsw_infos);
344                    let (mut tmp_lwe, scratch_2) = scratch_1.take_lwe(bits);
345                    for (local_bit, dst) in res_bits_chunk.iter_mut().enumerate() {
346                        bits.get_bit_lwe(self, start + local_bit, &mut tmp_lwe, ks, scratch_2);
347                        cbt.execute_to_constant(self, &mut tmp_ggsw, &tmp_lwe, 1, 1, scratch_2);
348                        dst.prepare(self, &tmp_ggsw, scratch_2);
349                    }
350                });
351            }
352        });
353
354        for i in 0..bit_start {
355            res.bits[i].zero(self);
356        }
357
358        for i in bit_end..T::BITS as usize {
359            res.bits[i].zero(self);
360        }
361    }
362}
363
364impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
365    pub fn prepare<BRA, M, O, K, DK>(&mut self, module: &M, other: &FheUint<O, T>, key: &K, scratch: &mut Scratch<BE>)
366    where
367        BRA: BlindRotationAlgo,
368        O: DataRef,
369        DK: DataRef,
370        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
371        M: FheUintPrepare<BRA, BE>,
372        Scratch<BE>: ScratchTakeCore<BE>,
373    {
374        module.fhe_uint_prepare(self, other, key, scratch);
375    }
376    pub fn prepare_custom<BRA, M, O, K, DK>(
377        &mut self,
378        module: &M,
379        other: &FheUint<O, T>,
380        bit_start: usize,
381        bit_end: usize,
382        key: &K,
383        scratch: &mut Scratch<BE>,
384    ) where
385        BRA: BlindRotationAlgo,
386        O: DataRef,
387        DK: DataRef,
388        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
389        M: FheUintPrepare<BRA, BE>,
390        Scratch<BE>: ScratchTakeCore<BE>,
391    {
392        module.fhe_uint_prepare_custom(self, other, bit_start, bit_end, key, scratch);
393    }
394
395    pub fn prepare_custom_multi_thread<BRA, M, O, K, DK>(
396        &mut self,
397        threads: usize,
398        module: &M,
399        other: &FheUint<O, T>,
400        bit_start: usize,
401        bit_end: usize,
402        key: &K,
403        scratch: &mut Scratch<BE>,
404    ) where
405        BRA: BlindRotationAlgo,
406        O: DataRef,
407        DK: DataRef,
408        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
409        M: FheUintPrepare<BRA, BE>,
410        Scratch<BE>: ScratchTakeCore<BE>,
411    {
412        module.fhe_uint_prepare_custom_multi_thread(threads, self, other, bit_start, bit_end, key, scratch);
413    }
414}