1use crate::tfhe::bdd_arithmetic::FheUintPreparedDebug;
  2use crate::tfhe::{
  3    bdd_arithmetic::{FheUint, FheUintPrepared, UnsignedInteger},
  4    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory},
  5    circuit_bootstrapping::{
  6        CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
  7        CircuitBootstrappingKeyPrepared, CircuitBootstrappingKeyPreparedFactory, CirtuitBootstrappingExecute,
  8    },
  9};
 10use poulpy_core::{
 11    GLWEToLWESwitchingKeyEncryptSk, GetDistribution, LWEFromGLWE, ScratchTakeCore,
 12    layouts::{
 13        GGSWInfos, GGSWPreparedFactory, GLWEInfos, GLWESecretToRef, GLWEToLWEKey, GLWEToLWEKeyLayout,
 14        GLWEToLWEKeyPreparedFactory, LWE, LWEInfos, LWESecretToRef, prepared::GLWEToLWEKeyPrepared,
 15    },
 16};
 17use poulpy_hal::{
 18    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 19    source::Source,
 20};
 21
 22pub trait BDDKeyInfos {
 23    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout;
 24    fn ks_infos(&self) -> GLWEToLWEKeyLayout;
 25}
 26
 27#[derive(Debug, Clone, Copy)]
 28pub struct BDDKeyLayout {
 29    pub cbt: CircuitBootstrappingKeyLayout,
 30    pub ks: GLWEToLWEKeyLayout,
 31}
 32
 33impl BDDKeyInfos for BDDKeyLayout {
 34    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
 35        self.cbt
 36    }
 37
 38    fn ks_infos(&self) -> GLWEToLWEKeyLayout {
 39        self.ks
 40    }
 41}
 42
 43pub struct BDDKey<D, BRA>
 44where
 45    D: Data,
 46    BRA: BlindRotationAlgo,
 47{
 48    cbt: CircuitBootstrappingKey<D, BRA>,
 49    ks: GLWEToLWEKey<D>,
 50}
 51
 52impl<BRA: BlindRotationAlgo> BDDKey<Vec<u8>, BRA>
 53where
 54    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 55{
 56    pub fn alloc_from_infos<A>(infos: &A) -> Self
 57    where
 58        A: BDDKeyInfos,
 59    {
 60        Self {
 61            cbt: CircuitBootstrappingKey::alloc_from_infos(&infos.cbt_infos()),
 62            ks: GLWEToLWEKey::alloc_from_infos(&infos.ks_infos()),
 63        }
 64    }
 65}
 66
 67pub trait BDDKeyEncryptSk<BRA: BlindRotationAlgo, BE: Backend> {
 68    fn bdd_key_encrypt_sk<D, S0, S1>(
 69        &self,
 70        res: &mut BDDKey<D, BRA>,
 71        sk_lwe: &S0,
 72        sk_glwe: &S1,
 73        source_xa: &mut Source,
 74        source_xe: &mut Source,
 75        scratch: &mut Scratch<BE>,
 76    ) where
 77        D: DataMut,
 78        S0: LWESecretToRef + GetDistribution + LWEInfos,
 79        S1: GLWESecretToRef + GetDistribution + GLWEInfos;
 80}
 81
 82impl<BE: Backend, BRA: BlindRotationAlgo> BDDKeyEncryptSk<BRA, BE> for Module<BE>
 83where
 84    Self: CircuitBootstrappingKeyEncryptSk<BRA, BE> + GLWEToLWESwitchingKeyEncryptSk<BE>,
 85    Scratch<BE>: ScratchTakeCore<BE>,
 86{
 87    fn bdd_key_encrypt_sk<D, S0, S1>(
 88        &self,
 89        res: &mut BDDKey<D, BRA>,
 90        sk_lwe: &S0,
 91        sk_glwe: &S1,
 92        source_xa: &mut Source,
 93        source_xe: &mut Source,
 94        scratch: &mut Scratch<BE>,
 95    ) where
 96        D: DataMut,
 97        S0: LWESecretToRef + GetDistribution + LWEInfos,
 98        S1: GLWESecretToRef + GetDistribution + GLWEInfos,
 99    {
100        res.ks
101            .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
102        res.cbt
103            .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
104    }
105}
106
107impl<D: DataMut, BRA: BlindRotationAlgo> BDDKey<D, BRA> {
108    pub fn encrypt_sk<S0, S1, M, BE: Backend>(
109        &mut self,
110        module: &M,
111        sk_lwe: &S0,
112        sk_glwe: &S1,
113        source_xa: &mut Source,
114        source_xe: &mut Source,
115        scratch: &mut Scratch<BE>,
116    ) where
117        S0: LWESecretToRef + GetDistribution + LWEInfos,
118        S1: GLWESecretToRef + GetDistribution + GLWEInfos,
119        M: BDDKeyEncryptSk<BRA, BE>,
120        Scratch<BE>: ScratchTakeCore<BE>,
121    {
122        module.bdd_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
123    }
124}
125
126pub struct BddKeyPrepared<D, BRA, BE>
127where
128    D: Data,
129    BRA: BlindRotationAlgo,
130    BE: Backend,
131{
132    pub(crate) cbt: CircuitBootstrappingKeyPrepared<D, BRA, BE>,
133    pub(crate) ks: GLWEToLWEKeyPrepared<D, BE>,
134}
135
136pub trait BDDKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend>
137where
138    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>,
139{
140    fn alloc_bdd_key_from_infos<A>(&self, infos: &A) -> BddKeyPrepared<Vec<u8>, BRA, BE>
141    where
142        A: BDDKeyInfos,
143    {
144        BddKeyPrepared {
145            cbt: CircuitBootstrappingKeyPrepared::alloc_from_infos(self, &infos.cbt_infos()),
146            ks: GLWEToLWEKeyPrepared::alloc_from_infos(self, &infos.ks_infos()),
147        }
148    }
149
150    fn prepare_bdd_key_tmp_bytes<A>(&self, infos: &A) -> usize
151    where
152        A: BDDKeyInfos,
153    {
154        self.circuit_bootstrapping_key_prepare_tmp_bytes(&infos.cbt_infos())
155            .max(self.prepare_glwe_to_lwe_key_tmp_bytes(&infos.ks_infos()))
156    }
157
158    fn prepare_bdd_key<DM, DR>(&self, res: &mut BddKeyPrepared<DM, BRA, BE>, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
159    where
160        DM: DataMut,
161        DR: DataRef,
162        Scratch<BE>: ScratchTakeCore<BE>,
163    {
164        res.cbt.prepare(self, &other.cbt, scratch);
165        res.ks.prepare(self, &other.ks, scratch);
166    }
167}
168impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPreparedFactory<BRA, BE> for Module<BE> where
169    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>
170{
171}
172
173impl<BRA: BlindRotationAlgo, BE: Backend> BddKeyPrepared<Vec<u8>, BRA, BE> {
174    pub fn alloc_from_infos<M, A>(module: &M, infos: &A) -> Self
175    where
176        M: BDDKeyPreparedFactory<BRA, BE>,
177        A: BDDKeyInfos,
178    {
179        module.alloc_bdd_key_from_infos(infos)
180    }
181}
182
183impl<D: DataMut, BRA: BlindRotationAlgo, BE: Backend> BddKeyPrepared<D, BRA, BE> {
184    pub fn prepare<DR, M>(&mut self, module: &M, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
185    where
186        DR: DataRef,
187        M: BDDKeyPreparedFactory<BRA, BE>,
188        Scratch<BE>: ScratchTakeCore<BE>,
189    {
190        module.prepare_bdd_key(self, other, scratch);
191    }
192}
193
194pub trait FheUintBlocksPrepare<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
195    fn fhe_uint_prepare_tmp_bytes<R, A>(&self, block_size: usize, extension_factor: usize, res_infos: &R, infos: &A) -> usize
196    where
197        R: GGSWInfos,
198        A: BDDKeyInfos;
199    fn fhe_uint_prepare<DM, DR0, DR1>(
200        &self,
201        res: &mut FheUintPrepared<DM, T, BE>,
202        bits: &FheUint<DR0, T>,
203        key: &BddKeyPrepared<DR1, BRA, BE>,
204        scratch: &mut Scratch<BE>,
205    ) where
206        DM: DataMut,
207        DR0: DataRef,
208        DR1: DataRef;
209}
210
211impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintBlocksPrepare<BRA, T, BE> for Module<BE>
212where
213    Self: LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
214    Scratch<BE>: ScratchTakeCore<BE>,
215{
216    fn fhe_uint_prepare_tmp_bytes<R, A>(&self, block_size: usize, extension_factor: usize, res_infos: &R, bdd_infos: &A) -> usize
217    where
218        R: GGSWInfos,
219        A: BDDKeyInfos,
220    {
221        self.circuit_bootstrapping_execute_tmp_bytes(
222            block_size,
223            extension_factor,
224            res_infos,
225            &bdd_infos.cbt_infos(),
226        )
227    }
228
229    fn fhe_uint_prepare<DM, DR0, DR1>(
230        &self,
231        res: &mut FheUintPrepared<DM, T, BE>,
232        bits: &FheUint<DR0, T>,
233        key: &BddKeyPrepared<DR1, BRA, BE>,
234        scratch: &mut Scratch<BE>,
235    ) where
236        DM: DataMut,
237        DR0: DataRef,
238        DR1: DataRef,
239    {
240        assert_eq!(res.bits.len(), bits.bits.len());
241
242        let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&bits.bits[0]); //TODO: add TakeLWE
243        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
244        for (dst, src) in res.bits.iter_mut().zip(bits.bits.iter()) {
245            lwe.from_glwe(self, src, &key.ks, scratch_1);
246            key.cbt
247                .execute_to_constant(self, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
248            dst.prepare(self, &tmp_ggsw, scratch_1);
249        }
250    }
251}
252
253impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
254    pub fn prepare<BRA, M, O, K>(
255        &mut self,
256        module: &M,
257        other: &FheUint<O, T>,
258        key: &BddKeyPrepared<K, BRA, BE>,
259        scratch: &mut Scratch<BE>,
260    ) where
261        BRA: BlindRotationAlgo,
262        O: DataRef,
263        K: DataRef,
264        M: FheUintBlocksPrepare<BRA, T, BE>,
265        Scratch<BE>: ScratchTakeCore<BE>,
266    {
267        module.fhe_uint_prepare(self, other, key, scratch);
268    }
269}
270
271pub trait FheUintBlockDebugPrepare<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
272    fn fhe_uint_debug_prepare<DM, DR0, DR1>(
273        &self,
274        res: &mut FheUintPreparedDebug<DM, T>,
275        bits: &FheUint<DR0, T>,
276        key: &BddKeyPrepared<DR1, BRA, BE>,
277        scratch: &mut Scratch<BE>,
278    ) where
279        DM: DataMut,
280        DR0: DataRef,
281        DR1: DataRef;
282}