1use crate::tfhe::bdd_arithmetic::FheUintPreparedDebug;
  2use crate::tfhe::circuit_bootstrapping::CircuitBootstrappingKeyInfos;
  3use crate::tfhe::{
  4    bdd_arithmetic::{FheUint, UnsignedInteger},
  5    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory},
  6    circuit_bootstrapping::{
  7        CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
  8        CircuitBootstrappingKeyPrepared, CircuitBootstrappingKeyPreparedFactory,
  9    },
 10};
 11
 12use poulpy_core::GLWESwitchingKeyEncryptSk;
 13use poulpy_core::layouts::{
 14    GGLWEInfos, GLWEAutomorphismKeyHelper, GLWEAutomorphismKeyPrepared, GLWESecret, GLWESwitchingKey, GLWESwitchingKeyLayout,
 15    GLWESwitchingKeyPrepared,
 16};
 17use poulpy_core::{
 18    GLWEToLWESwitchingKeyEncryptSk, GetDistribution, ScratchTakeCore,
 19    layouts::{
 20        GLWEInfos, GLWESecretToRef, GLWEToLWEKey, GLWEToLWEKeyLayout, GLWEToLWEKeyPreparedFactory, LWEInfos, LWESecretToRef,
 21        prepared::GLWEToLWEKeyPrepared,
 22    },
 23};
 24use poulpy_hal::{
 25    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 26    source::Source,
 27};
 28
 29pub trait BDDKeyInfos {
 30    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout;
 31    fn ks_lwe_infos(&self) -> GLWEToLWEKeyLayout;
 32    fn ks_glwe_infos(&self) -> Option<GLWESwitchingKeyLayout>;
 33}
 34
 35#[derive(Debug, Clone, Copy)]
 36pub struct BDDKeyLayout {
 37    pub cbt: CircuitBootstrappingKeyLayout,
 38    pub ks_glwe: Option<GLWESwitchingKeyLayout>,
 39    pub ks_lwe: GLWEToLWEKeyLayout,
 40}
 41
 42impl BDDKeyInfos for BDDKeyLayout {
 43    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
 44        self.cbt
 45    }
 46
 47    fn ks_glwe_infos(&self) -> Option<GLWESwitchingKeyLayout> {
 48        self.ks_glwe
 49    }
 50
 51    fn ks_lwe_infos(&self) -> GLWEToLWEKeyLayout {
 52        self.ks_lwe
 53    }
 54}
 55
 56pub struct BDDKey<D, BRA>
 57where
 58    D: Data,
 59    BRA: BlindRotationAlgo,
 60{
 61    pub(crate) cbt: CircuitBootstrappingKey<D, BRA>,
 62    pub(crate) ks_glwe: Option<GLWESwitchingKey<D>>,
 63    pub(crate) ks_lwe: GLWEToLWEKey<D>,
 64}
 65
 66impl<BRA: BlindRotationAlgo> BDDKey<Vec<u8>, BRA>
 67where
 68    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 69{
 70    pub fn alloc_from_infos<A>(infos: &A) -> Self
 71    where
 72        A: BDDKeyInfos,
 73    {
 74        let ks_glwe: Option<GLWESwitchingKey<Vec<u8>>> = if let Some(ks_infos) = &infos.ks_glwe_infos() {
 75            Some(GLWESwitchingKey::alloc_from_infos(ks_infos))
 76        } else {
 77            None
 78        };
 79
 80        Self {
 81            cbt: CircuitBootstrappingKey::alloc_from_infos(&infos.cbt_infos()),
 82            ks_glwe: ks_glwe,
 83            ks_lwe: GLWEToLWEKey::alloc_from_infos(&infos.ks_lwe_infos()),
 84        }
 85    }
 86}
 87
 88pub trait BDDKeyEncryptSk<BRA: BlindRotationAlgo, BE: Backend> {
 89    fn bdd_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 90    where
 91        A: BDDKeyInfos;
 92
 93    fn bdd_key_encrypt_sk<D, S0, S1>(
 94        &self,
 95        res: &mut BDDKey<D, BRA>,
 96        sk_lwe: &S0,
 97        sk_glwe: &S1,
 98        source_xa: &mut Source,
 99        source_xe: &mut Source,
100        scratch: &mut Scratch<BE>,
101    ) where
102        D: DataMut,
103        S0: LWESecretToRef + GetDistribution + LWEInfos,
104        S1: GLWESecretToRef + GetDistribution + GLWEInfos;
105}
106
107impl<BE: Backend, BRA: BlindRotationAlgo> BDDKeyEncryptSk<BRA, BE> for Module<BE>
108where
109    Self: CircuitBootstrappingKeyEncryptSk<BRA, BE> + GLWEToLWESwitchingKeyEncryptSk<BE> + GLWESwitchingKeyEncryptSk<BE>,
110    Scratch<BE>: ScratchTakeCore<BE>,
111{
112    fn bdd_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
113    where
114        A: BDDKeyInfos,
115    {
116        self.circuit_bootstrapping_key_encrypt_sk_tmp_bytes(&infos.cbt_infos())
117            .max(self.glwe_to_lwe_key_encrypt_sk_tmp_bytes(&infos.ks_lwe_infos()))
118    }
119
120    fn bdd_key_encrypt_sk<D, S0, S1>(
121        &self,
122        res: &mut BDDKey<D, BRA>,
123        sk_lwe: &S0,
124        sk_glwe: &S1,
125        source_xa: &mut Source,
126        source_xe: &mut Source,
127        scratch: &mut Scratch<BE>,
128    ) where
129        D: DataMut,
130        S0: LWESecretToRef + GetDistribution + LWEInfos,
131        S1: GLWESecretToRef + GetDistribution + GLWEInfos,
132    {
133        if let Some(key) = &mut res.ks_glwe {
134            let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(sk_glwe.n(), key.rank_out());
135            sk_out.fill_ternary_prob(0.5, source_xe);
136            key.encrypt_sk(self, sk_glwe, &sk_out, source_xa, source_xe, scratch);
137            res.ks_lwe
138                .encrypt_sk(self, sk_lwe, &sk_out, source_xa, source_xe, scratch);
139        } else {
140            res.ks_lwe
141                .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
142        }
143
144        res.cbt
145            .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
146    }
147}
148
149impl<D: DataMut, BRA: BlindRotationAlgo> BDDKey<D, BRA> {
150    pub fn encrypt_sk<S0, S1, M, BE: Backend>(
151        &mut self,
152        module: &M,
153        sk_lwe: &S0,
154        sk_glwe: &S1,
155        source_xa: &mut Source,
156        source_xe: &mut Source,
157        scratch: &mut Scratch<BE>,
158    ) where
159        S0: LWESecretToRef + GetDistribution + LWEInfos,
160        S1: GLWESecretToRef + GetDistribution + GLWEInfos,
161        M: BDDKeyEncryptSk<BRA, BE>,
162        Scratch<BE>: ScratchTakeCore<BE>,
163    {
164        module.bdd_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
165    }
166}
167
168pub struct BDDKeyPrepared<D, BRA, BE>
169where
170    D: Data,
171    BRA: BlindRotationAlgo,
172    BE: Backend,
173{
174    pub(crate) cbt: CircuitBootstrappingKeyPrepared<D, BRA, BE>,
175    pub(crate) ks_glwe: Option<GLWESwitchingKeyPrepared<D, BE>>,
176    pub(crate) ks_lwe: GLWEToLWEKeyPrepared<D, BE>,
177}
178
179impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> BDDKeyInfos for BDDKeyPrepared<D, BRA, BE> {
180    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
181        CircuitBootstrappingKeyLayout {
182            layout_brk: self.cbt.brk_infos(),
183            layout_atk: self.cbt.atk_infos(),
184            layout_tsk: self.cbt.tsk_infos(),
185        }
186    }
187    fn ks_glwe_infos(&self) -> Option<GLWESwitchingKeyLayout> {
188        if let Some(ks_glwe) = &self.ks_glwe {
189            Some(GLWESwitchingKeyLayout {
190                n: ks_glwe.n(),
191                base2k: ks_glwe.base2k(),
192                k: ks_glwe.k(),
193                rank_in: ks_glwe.rank_in(),
194                rank_out: ks_glwe.rank_out(),
195                dnum: ks_glwe.dnum(),
196                dsize: ks_glwe.dsize(),
197            })
198        } else {
199            None
200        }
201    }
202    fn ks_lwe_infos(&self) -> GLWEToLWEKeyLayout {
203        GLWEToLWEKeyLayout {
204            n: self.ks_lwe.n(),
205            base2k: self.ks_lwe.base2k(),
206            k: self.ks_lwe.k(),
207            rank_in: self.ks_lwe.rank_in(),
208            dnum: self.ks_lwe.dnum(),
209        }
210    }
211}
212
213impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> GLWEAutomorphismKeyHelper<GLWEAutomorphismKeyPrepared<D, BE>, BE>
214    for BDDKeyPrepared<D, BRA, BE>
215{
216    fn automorphism_key_infos(&self) -> poulpy_core::layouts::GGLWELayout {
217        self.cbt.automorphism_key_infos()
218    }
219
220    fn get_automorphism_key(&self, k: i64) -> Option<&GLWEAutomorphismKeyPrepared<D, BE>> {
221        self.cbt.get_automorphism_key(k)
222    }
223}
224
225pub trait BDDKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend>
226where
227    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>,
228{
229    fn alloc_bdd_key_from_infos<A>(&self, infos: &A) -> BDDKeyPrepared<Vec<u8>, BRA, BE>
230    where
231        A: BDDKeyInfos,
232    {
233        let ks_glwe = if let Some(ks_glwe_infos) = &infos.ks_glwe_infos() {
234            Some(GLWESwitchingKeyPrepared::alloc_from_infos(
235                self,
236                ks_glwe_infos,
237            ))
238        } else {
239            None
240        };
241
242        BDDKeyPrepared {
243            cbt: CircuitBootstrappingKeyPrepared::alloc_from_infos(self, &infos.cbt_infos()),
244            ks_glwe,
245            ks_lwe: GLWEToLWEKeyPrepared::alloc_from_infos(self, &infos.ks_lwe_infos()),
246        }
247    }
248
249    fn prepare_bdd_key_tmp_bytes<A>(&self, infos: &A) -> usize
250    where
251        A: BDDKeyInfos,
252    {
253        self.circuit_bootstrapping_key_prepare_tmp_bytes(&infos.cbt_infos())
254            .max(self.prepare_glwe_to_lwe_key_tmp_bytes(&infos.ks_lwe_infos()))
255    }
256
257    fn prepare_bdd_key<DM, DR>(&self, res: &mut BDDKeyPrepared<DM, BRA, BE>, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
258    where
259        DM: DataMut,
260        DR: DataRef,
261        Scratch<BE>: ScratchTakeCore<BE>,
262    {
263        res.cbt.prepare(self, &other.cbt, scratch);
264
265        if let Some(key_prep) = &mut res.ks_glwe {
266            if let Some(other) = &other.ks_glwe {
267                key_prep.prepare(self, other, scratch);
268            } else {
269                panic!("incompatible keys: res has Some(ks_glwe) but other has none")
270            }
271        }
272
273        res.ks_lwe.prepare(self, &other.ks_lwe, scratch);
274    }
275}
276impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPreparedFactory<BRA, BE> for Module<BE> where
277    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>
278{
279}
280
281impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPrepared<Vec<u8>, BRA, BE> {
282    pub fn alloc_from_infos<M, A>(module: &M, infos: &A) -> Self
283    where
284        M: BDDKeyPreparedFactory<BRA, BE>,
285        A: BDDKeyInfos,
286    {
287        module.alloc_bdd_key_from_infos(infos)
288    }
289}
290
291impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> BDDKeyHelper<D, BRA, BE> for BDDKeyPrepared<D, BRA, BE> {
292    fn get_cbt_key(
293        &self,
294    ) -> (
295        &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
296        Option<&GLWESwitchingKeyPrepared<D, BE>>,
297        &GLWEToLWEKeyPrepared<D, BE>,
298    ) {
299        (&self.cbt, self.ks_glwe.as_ref(), &self.ks_lwe)
300    }
301}
302
303pub trait BDDKeyHelper<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> {
304    fn get_cbt_key(
305        &self,
306    ) -> (
307        &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
308        Option<&GLWESwitchingKeyPrepared<D, BE>>,
309        &GLWEToLWEKeyPrepared<D, BE>,
310    );
311}
312
313pub trait FheUintPrepareDebug<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
314    fn fhe_uint_debug_prepare<DM, DR0, DR1>(
315        &self,
316        res: &mut FheUintPreparedDebug<DM, T>,
317        bits: &FheUint<DR0, T>,
318        key: &BDDKeyPrepared<DR1, BRA, BE>,
319        scratch: &mut Scratch<BE>,
320    ) where
321        DM: DataMut,
322        DR0: DataRef,
323        DR1: DataRef;
324}