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