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