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::layouts::{GGLWEInfos, GLWEAutomorphismKeyHelper, GLWEAutomorphismKeyPrepared};
 13use poulpy_core::{
 14    GLWEToLWESwitchingKeyEncryptSk, GetDistribution, ScratchTakeCore,
 15    layouts::{
 16        GLWEInfos, GLWESecretToRef, GLWEToLWEKey, GLWEToLWEKeyLayout, GLWEToLWEKeyPreparedFactory, LWEInfos, LWESecretToRef,
 17        prepared::GLWEToLWEKeyPrepared,
 18    },
 19};
 20use poulpy_hal::{
 21    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 22    source::Source,
 23};
 24
 25pub trait BDDKeyInfos {
 26    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout;
 27    fn ks_infos(&self) -> GLWEToLWEKeyLayout;
 28}
 29
 30#[derive(Debug, Clone, Copy)]
 31pub struct BDDKeyLayout {
 32    pub cbt: CircuitBootstrappingKeyLayout,
 33    pub ks: GLWEToLWEKeyLayout,
 34}
 35
 36impl BDDKeyInfos for BDDKeyLayout {
 37    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
 38        self.cbt
 39    }
 40
 41    fn ks_infos(&self) -> GLWEToLWEKeyLayout {
 42        self.ks
 43    }
 44}
 45
 46pub struct BDDKey<D, BRA>
 47where
 48    D: Data,
 49    BRA: BlindRotationAlgo,
 50{
 51    pub(crate) cbt: CircuitBootstrappingKey<D, BRA>,
 52    pub(crate) ks: GLWEToLWEKey<D>,
 53}
 54
 55impl<BRA: BlindRotationAlgo> BDDKey<Vec<u8>, BRA>
 56where
 57    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 58{
 59    pub fn alloc_from_infos<A>(infos: &A) -> Self
 60    where
 61        A: BDDKeyInfos,
 62    {
 63        Self {
 64            cbt: CircuitBootstrappingKey::alloc_from_infos(&infos.cbt_infos()),
 65            ks: GLWEToLWEKey::alloc_from_infos(&infos.ks_infos()),
 66        }
 67    }
 68}
 69
 70pub trait BDDKeyEncryptSk<BRA: BlindRotationAlgo, BE: Backend> {
 71    fn bdd_key_encrypt_sk<D, S0, S1>(
 72        &self,
 73        res: &mut BDDKey<D, BRA>,
 74        sk_lwe: &S0,
 75        sk_glwe: &S1,
 76        source_xa: &mut Source,
 77        source_xe: &mut Source,
 78        scratch: &mut Scratch<BE>,
 79    ) where
 80        D: DataMut,
 81        S0: LWESecretToRef + GetDistribution + LWEInfos,
 82        S1: GLWESecretToRef + GetDistribution + GLWEInfos;
 83}
 84
 85impl<BE: Backend, BRA: BlindRotationAlgo> BDDKeyEncryptSk<BRA, BE> for Module<BE>
 86where
 87    Self: CircuitBootstrappingKeyEncryptSk<BRA, BE> + GLWEToLWESwitchingKeyEncryptSk<BE>,
 88    Scratch<BE>: ScratchTakeCore<BE>,
 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        res.ks
104            .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
105        res.cbt
106            .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
107    }
108}
109
110impl<D: DataMut, BRA: BlindRotationAlgo> BDDKey<D, BRA> {
111    pub fn encrypt_sk<S0, S1, M, BE: Backend>(
112        &mut self,
113        module: &M,
114        sk_lwe: &S0,
115        sk_glwe: &S1,
116        source_xa: &mut Source,
117        source_xe: &mut Source,
118        scratch: &mut Scratch<BE>,
119    ) where
120        S0: LWESecretToRef + GetDistribution + LWEInfos,
121        S1: GLWESecretToRef + GetDistribution + GLWEInfos,
122        M: BDDKeyEncryptSk<BRA, BE>,
123        Scratch<BE>: ScratchTakeCore<BE>,
124    {
125        module.bdd_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
126    }
127}
128
129pub struct BDDKeyPrepared<D, BRA, BE>
130where
131    D: Data,
132    BRA: BlindRotationAlgo,
133    BE: Backend,
134{
135    pub(crate) cbt: CircuitBootstrappingKeyPrepared<D, BRA, BE>,
136    pub(crate) ks: GLWEToLWEKeyPrepared<D, BE>,
137}
138
139impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> BDDKeyInfos for BDDKeyPrepared<D, BRA, BE> {
140    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
141        CircuitBootstrappingKeyLayout {
142            layout_brk: self.cbt.brk_infos(),
143            layout_atk: self.cbt.atk_infos(),
144            layout_tsk: self.cbt.tsk_infos(),
145        }
146    }
147    fn ks_infos(&self) -> GLWEToLWEKeyLayout {
148        GLWEToLWEKeyLayout {
149            n: self.ks.n(),
150            base2k: self.ks.base2k(),
151            k: self.ks.k(),
152            rank_in: self.ks.rank_in(),
153            dnum: self.ks.dnum(),
154        }
155    }
156}
157
158impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> GLWEAutomorphismKeyHelper<GLWEAutomorphismKeyPrepared<D, BE>, BE>
159    for BDDKeyPrepared<D, BRA, BE>
160{
161    fn automorphism_key_infos(&self) -> poulpy_core::layouts::GGLWELayout {
162        self.cbt.automorphism_key_infos()
163    }
164
165    fn get_automorphism_key(&self, k: i64) -> Option<&GLWEAutomorphismKeyPrepared<D, BE>> {
166        self.cbt.get_automorphism_key(k)
167    }
168}
169
170pub trait BDDKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend>
171where
172    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>,
173{
174    fn alloc_bdd_key_from_infos<A>(&self, infos: &A) -> BDDKeyPrepared<Vec<u8>, BRA, BE>
175    where
176        A: BDDKeyInfos,
177    {
178        BDDKeyPrepared {
179            cbt: CircuitBootstrappingKeyPrepared::alloc_from_infos(self, &infos.cbt_infos()),
180            ks: GLWEToLWEKeyPrepared::alloc_from_infos(self, &infos.ks_infos()),
181        }
182    }
183
184    fn prepare_bdd_key_tmp_bytes<A>(&self, infos: &A) -> usize
185    where
186        A: BDDKeyInfos,
187    {
188        self.circuit_bootstrapping_key_prepare_tmp_bytes(&infos.cbt_infos())
189            .max(self.prepare_glwe_to_lwe_key_tmp_bytes(&infos.ks_infos()))
190    }
191
192    fn prepare_bdd_key<DM, DR>(&self, res: &mut BDDKeyPrepared<DM, BRA, BE>, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
193    where
194        DM: DataMut,
195        DR: DataRef,
196        Scratch<BE>: ScratchTakeCore<BE>,
197    {
198        res.cbt.prepare(self, &other.cbt, scratch);
199        res.ks.prepare(self, &other.ks, scratch);
200    }
201}
202impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPreparedFactory<BRA, BE> for Module<BE> where
203    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>
204{
205}
206
207impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPrepared<Vec<u8>, BRA, BE> {
208    pub fn alloc_from_infos<M, A>(module: &M, infos: &A) -> Self
209    where
210        M: BDDKeyPreparedFactory<BRA, BE>,
211        A: BDDKeyInfos,
212    {
213        module.alloc_bdd_key_from_infos(infos)
214    }
215}
216
217impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> BDDKeyHelper<D, BRA, BE> for BDDKeyPrepared<D, BRA, BE> {
218    fn get_cbt_key(
219        &self,
220    ) -> (
221        &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
222        &GLWEToLWEKeyPrepared<D, BE>,
223    ) {
224        (&self.cbt, &self.ks)
225    }
226}
227
228pub trait BDDKeyHelper<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> {
229    fn get_cbt_key(
230        &self,
231    ) -> (
232        &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
233        &GLWEToLWEKeyPrepared<D, BE>,
234    );
235}
236
237pub trait FheUintPrepareDebug<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
238    fn fhe_uint_debug_prepare<DM, DR0, DR1>(
239        &self,
240        res: &mut FheUintPreparedDebug<DM, T>,
241        bits: &FheUint<DR0, T>,
242        key: &BDDKeyPrepared<DR1, BRA, BE>,
243        scratch: &mut Scratch<BE>,
244    ) where
245        DM: DataMut,
246        DR0: DataRef,
247        DR1: DataRef;
248}