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