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