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