1use crate::tfhe::bdd_arithmetic::FheUintBlocksPreparedDebug;
  2use crate::tfhe::{
  3    bdd_arithmetic::{FheUintBlocks, FheUintBlocksPrepared, 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, GLWEToLWEKeyLayout, GLWEToLWESwitchingKey,
 14        GLWEToLWESwitchingKeyPreparedFactory, LWE, LWEInfos, LWESecretToRef, prepared::GLWEToLWESwitchingKeyPrepared,
 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: GLWEToLWESwitchingKey<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: GLWEToLWESwitchingKey::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: GLWEToLWESwitchingKeyPrepared<D, BE>,
134}
135
136pub trait BDDKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend>
137where
138    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWESwitchingKeyPreparedFactory<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: GLWEToLWESwitchingKeyPrepared::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_switching_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> + GLWEToLWESwitchingKeyPreparedFactory<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_blocks_prepare_tmp_bytes<R, A>(
196        &self,
197        block_size: usize,
198        extension_factor: usize,
199        res_infos: &R,
200        infos: &A,
201    ) -> usize
202    where
203        R: GGSWInfos,
204        A: BDDKeyInfos;
205    fn fhe_uint_blocks_prepare<DM, DR0, DR1>(
206        &self,
207        res: &mut FheUintBlocksPrepared<DM, T, BE>,
208        bits: &FheUintBlocks<DR0, T>,
209        key: &BDDKeyPrepared<DR1, BRA, BE>,
210        scratch: &mut Scratch<BE>,
211    ) where
212        DM: DataMut,
213        DR0: DataRef,
214        DR1: DataRef;
215}
216
217impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintBlocksPrepare<BRA, T, BE> for Module<BE>
218where
219    Self: LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
220    Scratch<BE>: ScratchTakeCore<BE>,
221{
222    fn fhe_uint_blocks_prepare_tmp_bytes<R, A>(
223        &self,
224        block_size: usize,
225        extension_factor: usize,
226        res_infos: &R,
227        bdd_infos: &A,
228    ) -> usize
229    where
230        R: GGSWInfos,
231        A: BDDKeyInfos,
232    {
233        self.circuit_bootstrapping_execute_tmp_bytes(
234            block_size,
235            extension_factor,
236            res_infos,
237            &bdd_infos.cbt_infos(),
238        )
239    }
240
241    fn fhe_uint_blocks_prepare<DM, DR0, DR1>(
242        &self,
243        res: &mut FheUintBlocksPrepared<DM, T, BE>,
244        bits: &FheUintBlocks<DR0, T>,
245        key: &BDDKeyPrepared<DR1, BRA, BE>,
246        scratch: &mut Scratch<BE>,
247    ) where
248        DM: DataMut,
249        DR0: DataRef,
250        DR1: DataRef,
251    {
252        assert_eq!(res.blocks.len(), bits.blocks.len());
253
254        let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&bits.blocks[0]); //TODO: add TakeLWE
255        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
256        for (dst, src) in res.blocks.iter_mut().zip(bits.blocks.iter()) {
257            lwe.from_glwe(self, src, &key.ks, scratch_1);
258            key.cbt
259                .execute_to_constant(self, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
260            dst.prepare(self, &tmp_ggsw, scratch_1);
261        }
262    }
263}
264
265impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintBlocksPrepared<D, T, BE> {
266    pub fn prepare<BRA, M, O, K>(
267        &mut self,
268        module: &M,
269        other: &FheUintBlocks<O, T>,
270        key: &BDDKeyPrepared<K, BRA, BE>,
271        scratch: &mut Scratch<BE>,
272    ) where
273        BRA: BlindRotationAlgo,
274        O: DataRef,
275        K: DataRef,
276        M: FheUintBlocksPrepare<BRA, T, BE>,
277        Scratch<BE>: ScratchTakeCore<BE>,
278    {
279        module.fhe_uint_blocks_prepare(self, other, key, scratch);
280    }
281}
282
283pub trait FheUintBlockDebugPrepare<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
284    fn fhe_uint_block_debug_prepare<DM, DR0, DR1>(
285        &self,
286        res: &mut FheUintBlocksPreparedDebug<DM, T>,
287        bits: &FheUintBlocks<DR0, T>,
288        key: &BDDKeyPrepared<DR1, BRA, BE>,
289        scratch: &mut Scratch<BE>,
290    ) where
291        DM: DataMut,
292        DR0: DataRef,
293        DR1: DataRef;
294}