1use std::marker::PhantomData;
  2
  3use poulpy_core::layouts::{
  4    Base2K, Dnum, Dsize, GGSWInfos, GGSWPreparedFactory, GLWEInfos, LWEInfos, Rank, TorusPrecision, prepared::GGSWPrepared,
  5};
  6use poulpy_core::layouts::{
  7    GGLWEInfos, GGLWEPreparedToRef, GGSWPreparedToMut, GGSWPreparedToRef, GLWEAutomorphismKeyHelper, GetGaloisElement, LWE,
  8};
  9use poulpy_core::{GLWECopy, GLWEDecrypt, GLWEPacking, LWEFromGLWE};
 10
 11use poulpy_core::{GGSWEncryptSk, ScratchTakeCore, layouts::GLWESecretPreparedToRef};
 12use poulpy_hal::api::ModuleLogN;
 13use poulpy_hal::layouts::{Backend, Data, DataRef, Module};
 14
 15use poulpy_hal::{
 16    api::ModuleN,
 17    layouts::{DataMut, Scratch},
 18    source::Source,
 19};
 20
 21use crate::tfhe::bdd_arithmetic::{BDDKey, BDDKeyHelper, BDDKeyInfos, BDDKeyPrepared, BDDKeyPreparedFactory, FheUint, ToBits};
 22use crate::tfhe::bdd_arithmetic::{Cmux, FromBits, ScratchTakeBDD, UnsignedInteger};
 23use crate::tfhe::blind_rotation::BlindRotationAlgo;
 24use crate::tfhe::circuit_bootstrapping::CirtuitBootstrappingExecute;
 25
 26/// A prepared FHE ciphertext encrypting the bits of an [UnsignedInteger].
 27pub struct FheUintPrepared<D: Data, T: UnsignedInteger, B: Backend> {
 28    pub(crate) bits: Vec<GGSWPrepared<D, B>>,
 29    pub(crate) _phantom: PhantomData<T>,
 30}
 31
 32impl<T: UnsignedInteger, BE: Backend> FheUintPreparedFactory<T, BE> for Module<BE> where Self: Sized + GGSWPreparedFactory<BE> {}
 33
 34pub trait GetGGSWBit<BE: Backend> {
 35    fn get_bit(&self, bit: usize) -> GGSWPrepared<&[u8], BE>;
 36}
 37
 38impl<D: DataRef, T: UnsignedInteger, BE: Backend> GetGGSWBit<BE> for FheUintPrepared<D, T, BE> {
 39    fn get_bit(&self, bit: usize) -> GGSWPrepared<&[u8], BE> {
 40        assert!(bit <= self.bits.len());
 41        self.bits[bit].to_ref()
 42    }
 43}
 44
 45pub trait GetGGSWBitMut<T: UnsignedInteger, BE: Backend> {
 46    fn get_bit(&mut self, bit: usize) -> GGSWPrepared<&mut [u8], BE>;
 47}
 48
 49impl<D: DataMut, T: UnsignedInteger, BE: Backend> GetGGSWBitMut<T, BE> for FheUintPrepared<D, T, BE> {
 50    fn get_bit(&mut self, bit: usize) -> GGSWPrepared<&mut [u8], BE> {
 51        assert!(bit <= self.bits.len());
 52        self.bits[bit].to_mut()
 53    }
 54}
 55
 56pub trait FheUintPreparedFactory<T: UnsignedInteger, BE: Backend>
 57where
 58    Self: Sized + GGSWPreparedFactory<BE>,
 59{
 60    fn alloc_fhe_uint_prepared(
 61        &self,
 62        base2k: Base2K,
 63        k: TorusPrecision,
 64        dnum: Dnum,
 65        dsize: Dsize,
 66        rank: Rank,
 67    ) -> FheUintPrepared<Vec<u8>, T, BE> {
 68        FheUintPrepared {
 69            bits: (0..T::BITS)
 70                .map(|_| GGSWPrepared::alloc(self, base2k, k, dnum, dsize, rank))
 71                .collect(),
 72            _phantom: PhantomData,
 73        }
 74    }
 75
 76    fn alloc_fhe_uint_prepared_from_infos<A>(&self, infos: &A) -> FheUintPrepared<Vec<u8>, T, BE>
 77    where
 78        A: GGSWInfos,
 79    {
 80        self.alloc_fhe_uint_prepared(
 81            infos.base2k(),
 82            infos.k(),
 83            infos.dnum(),
 84            infos.dsize(),
 85            infos.rank(),
 86        )
 87    }
 88}
 89
 90impl<T: UnsignedInteger, BE: Backend> FheUintPrepared<Vec<u8>, T, BE> {
 91    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
 92    where
 93        A: GGSWInfos,
 94        M: FheUintPreparedFactory<T, BE>,
 95    {
 96        module.alloc_fhe_uint_prepared_from_infos(infos)
 97    }
 98
 99    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
100    where
101        M: FheUintPreparedFactory<T, BE>,
102    {
103        module.alloc_fhe_uint_prepared(base2k, k, dnum, dsize, rank)
104    }
105}
106
107impl<T: UnsignedInteger + ToBits, BE: Backend> FheUintPreparedEncryptSk<T, BE> for Module<BE> where
108    Self: Sized + ModuleN + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>
109{
110}
111
112pub trait FheUintPreparedEncryptSk<T: UnsignedInteger + ToBits, BE: Backend>
113where
114    Self: Sized + ModuleN + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>,
115{
116    fn fhe_uint_prepared_encrypt_sk<DM, S>(
117        &self,
118        res: &mut FheUintPrepared<DM, T, BE>,
119        value: T,
120        sk: &S,
121        source_xa: &mut Source,
122        source_xe: &mut Source,
123        scratch: &mut Scratch<BE>,
124    ) where
125        DM: DataMut,
126        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
127        Scratch<BE>: ScratchTakeCore<BE>,
128    {
129        use poulpy_hal::{api::ScratchTakeBasic, layouts::ZnxZero};
130
131        assert!(self.n().is_multiple_of(T::BITS as usize));
132        assert_eq!(res.n(), self.n() as u32);
133        assert_eq!(sk.n(), self.n() as u32);
134
135        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
136        let (mut pt, scratch_2) = scratch_1.take_scalar_znx(self.n(), 1);
137        pt.zero();
138
139        for i in 0..T::BITS as usize {
140            use poulpy_hal::layouts::ZnxViewMut;
141            pt.at_mut(0, 0)[0] = value.bit(i) as i64;
142            tmp_ggsw.encrypt_sk(self, &pt, sk, source_xa, source_xe, scratch_2);
143            res.bits[i].prepare(self, &tmp_ggsw, scratch_2);
144        }
145    }
146}
147
148impl<D: DataMut, T: UnsignedInteger + ToBits, BE: Backend> FheUintPrepared<D, T, BE> {
149    pub fn encrypt_sk<M, S>(
150        &mut self,
151        module: &M,
152        value: T,
153        sk: &S,
154        source_xa: &mut Source,
155        source_xe: &mut Source,
156        scratch: &mut Scratch<BE>,
157    ) where
158        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
159        M: FheUintPreparedEncryptSk<T, BE>,
160        Scratch<BE>: ScratchTakeCore<BE>,
161    {
162        module.fhe_uint_prepared_encrypt_sk(self, value, sk, source_xa, source_xe, scratch);
163    }
164}
165
166impl<D: DataRef, T: UnsignedInteger + FromBits, BE: Backend> FheUintPrepared<D, T, BE> {
167    pub fn decrypt<M, S, H, K>(&self, module: &M, sk: &S, keys: &H, scratch: &mut Scratch<BE>) -> T
168    where
169        M: ModuleLogN + GLWEDecrypt<BE> + Cmux<BE> + GLWEPacking<BE> + GLWECopy,
170        S: GLWESecretPreparedToRef<BE> + GLWEInfos,
171        Scratch<BE>: ScratchTakeCore<BE>,
172        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
173        H: GLWEAutomorphismKeyHelper<K, BE>,
174    {
175        let (mut tmp, scratch_1) = scratch.take_fhe_uint(self);
176        tmp.from_fhe_uint_prepared(module, self, keys, scratch_1);
177        tmp.decrypt(module, sk, scratch_1)
178    }
179}
180
181impl<D: DataRef, T: UnsignedInteger, B: Backend> LWEInfos for FheUintPrepared<D, T, B> {
182    fn base2k(&self) -> poulpy_core::layouts::Base2K {
183        self.bits[0].base2k()
184    }
185
186    fn k(&self) -> poulpy_core::layouts::TorusPrecision {
187        self.bits[0].k()
188    }
189
190    fn n(&self) -> poulpy_core::layouts::Degree {
191        self.bits[0].n()
192    }
193}
194
195impl<D: DataRef, T: UnsignedInteger, B: Backend> GLWEInfos for FheUintPrepared<D, T, B> {
196    fn rank(&self) -> poulpy_core::layouts::Rank {
197        self.bits[0].rank()
198    }
199}
200
201impl<D: DataRef, T: UnsignedInteger, B: Backend> GGSWInfos for FheUintPrepared<D, T, B> {
202    fn dsize(&self) -> poulpy_core::layouts::Dsize {
203        self.bits[0].dsize()
204    }
205
206    fn dnum(&self) -> poulpy_core::layouts::Dnum {
207        self.bits[0].dnum()
208    }
209}
210
211impl<D: DataMut, BRA: BlindRotationAlgo, BE: Backend> BDDKeyPrepared<D, BRA, BE> {
212    pub fn prepare<DR, M>(&mut self, module: &M, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
213    where
214        DR: DataRef,
215        M: BDDKeyPreparedFactory<BRA, BE>,
216        Scratch<BE>: ScratchTakeCore<BE>,
217    {
218        module.prepare_bdd_key(self, other, scratch);
219    }
220}
221
222pub trait FheUintPrepare<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
223    fn fhe_uint_prepare_tmp_bytes<R, A>(&self, block_size: usize, extension_factor: usize, res_infos: &R, infos: &A) -> usize
224    where
225        R: GGSWInfos,
226        A: BDDKeyInfos;
227    fn fhe_uint_prepare<DM, DB, DK, K>(
228        &self,
229        res: &mut FheUintPrepared<DM, T, BE>,
230        bits: &FheUint<DB, T>,
231        key: &K,
232        scratch: &mut Scratch<BE>,
233    ) where
234        DM: DataMut,
235        DB: DataRef,
236        DK: DataRef,
237        K: BDDKeyHelper<DK, BRA, BE>;
238    fn fhe_uint_prepare_custom<DM, DB, DK, K>(
239        &self,
240        res: &mut FheUintPrepared<DM, T, BE>,
241        bits: &FheUint<DB, T>,
242        bit_start: usize,
243        bit_end: usize,
244        key: &K,
245        scratch: &mut Scratch<BE>,
246    ) where
247        DM: DataMut,
248        DB: DataRef,
249        DK: DataRef,
250        K: BDDKeyHelper<DK, BRA, BE>;
251}
252
253impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintPrepare<BRA, T, BE> for Module<BE>
254where
255    Self: LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
256    Scratch<BE>: ScratchTakeCore<BE>,
257{
258    fn fhe_uint_prepare_tmp_bytes<R, A>(&self, block_size: usize, extension_factor: usize, res_infos: &R, bdd_infos: &A) -> usize
259    where
260        R: GGSWInfos,
261        A: BDDKeyInfos,
262    {
263        self.circuit_bootstrapping_execute_tmp_bytes(
264            block_size,
265            extension_factor,
266            res_infos,
267            &bdd_infos.cbt_infos(),
268        )
269    }
270
271    fn fhe_uint_prepare<DM, DB, DK, K>(
272        &self,
273        res: &mut FheUintPrepared<DM, T, BE>,
274        bits: &FheUint<DB, T>,
275        key: &K,
276        scratch: &mut Scratch<BE>,
277    ) where
278        DM: DataMut,
279        DB: DataRef,
280        DK: DataRef,
281        K: BDDKeyHelper<DK, BRA, BE>,
282    {
283        self.fhe_uint_prepare_custom(res, bits, 0, T::BITS as usize, key, scratch);
284    }
285
286    fn fhe_uint_prepare_custom<DM, DB, DK, K>(
287        &self,
288        res: &mut FheUintPrepared<DM, T, BE>,
289        bits: &FheUint<DB, T>,
290        bit_start: usize,
291        bit_end: 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>,
299    {
300        let (cbt, ks) = key.get_cbt_key();
301
302        let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(bits); //TODO: add TakeLWE
303        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
304        for (bit, dst) in res.bits[bit_start..bit_end].iter_mut().enumerate() {
305            // TODO: set the rest of the bits to a prepared zero GGSW
306            bits.get_bit_lwe(self, bit, &mut lwe, ks, scratch_1);
307            cbt.execute_to_constant(self, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
308            dst.prepare(self, &tmp_ggsw, scratch_1);
309        }
310
311        for i in 0..bit_start {
312            res.bits[i].zero(self);
313        }
314
315        for i in bit_end..T::BITS as usize {
316            res.bits[i].zero(self);
317        }
318    }
319}
320
321impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
322    pub fn prepare<BRA, M, O, K, DK>(&mut self, module: &M, other: &FheUint<O, T>, key: &K, scratch: &mut Scratch<BE>)
323    where
324        BRA: BlindRotationAlgo,
325        O: DataRef,
326        DK: DataRef,
327        K: BDDKeyHelper<DK, BRA, BE>,
328        M: FheUintPrepare<BRA, T, BE>,
329        Scratch<BE>: ScratchTakeCore<BE>,
330    {
331        module.fhe_uint_prepare(self, other, key, scratch);
332    }
333    pub fn prepare_custom<BRA, M, O, K, DK>(
334        &mut self,
335        module: &M,
336        other: &FheUint<O, T>,
337        bit_start: usize,
338        bit_end: usize,
339        key: &K,
340        scratch: &mut Scratch<BE>,
341    ) where
342        BRA: BlindRotationAlgo,
343        O: DataRef,
344        DK: DataRef,
345        K: BDDKeyHelper<DK, BRA, BE>,
346        M: FheUintPrepare<BRA, T, BE>,
347        Scratch<BE>: ScratchTakeCore<BE>,
348    {
349        module.fhe_uint_prepare_custom(self, other, bit_start, bit_end, key, scratch);
350    }
351}