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_tmp_bytes<A>(&self, infos: &A) -> usize
 72    where
 73        A: BDDKeyInfos;
 74
 75    fn bdd_key_encrypt_sk<D, S0, S1>(
 76        &self,
 77        res: &mut BDDKey<D, BRA>,
 78        sk_lwe: &S0,
 79        sk_glwe: &S1,
 80        source_xa: &mut Source,
 81        source_xe: &mut Source,
 82        scratch: &mut Scratch<BE>,
 83    ) where
 84        D: DataMut,
 85        S0: LWESecretToRef + GetDistribution + LWEInfos,
 86        S1: GLWESecretToRef + GetDistribution + GLWEInfos;
 87}
 88
 89impl<BE: Backend, BRA: BlindRotationAlgo> BDDKeyEncryptSk<BRA, BE> for Module<BE>
 90where
 91    Self: CircuitBootstrappingKeyEncryptSk<BRA, BE> + GLWEToLWESwitchingKeyEncryptSk<BE>,
 92    Scratch<BE>: ScratchTakeCore<BE>,
 93{
 94    fn bdd_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 95    where
 96        A: BDDKeyInfos,
 97    {
 98        self.circuit_bootstrapping_key_encrypt_sk_tmp_bytes(&infos.cbt_infos())
 99            .max(self.glwe_to_lwe_key_encrypt_sk_tmp_bytes(&infos.ks_infos()))
100    }
101
102    fn bdd_key_encrypt_sk<D, S0, S1>(
103        &self,
104        res: &mut BDDKey<D, BRA>,
105        sk_lwe: &S0,
106        sk_glwe: &S1,
107        source_xa: &mut Source,
108        source_xe: &mut Source,
109        scratch: &mut Scratch<BE>,
110    ) where
111        D: DataMut,
112        S0: LWESecretToRef + GetDistribution + LWEInfos,
113        S1: GLWESecretToRef + GetDistribution + GLWEInfos,
114    {
115        res.ks
116            .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
117        res.cbt
118            .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
119    }
120}
121
122impl<D: DataMut, BRA: BlindRotationAlgo> BDDKey<D, BRA> {
123    pub fn encrypt_sk<S0, S1, M, BE: Backend>(
124        &mut self,
125        module: &M,
126        sk_lwe: &S0,
127        sk_glwe: &S1,
128        source_xa: &mut Source,
129        source_xe: &mut Source,
130        scratch: &mut Scratch<BE>,
131    ) where
132        S0: LWESecretToRef + GetDistribution + LWEInfos,
133        S1: GLWESecretToRef + GetDistribution + GLWEInfos,
134        M: BDDKeyEncryptSk<BRA, BE>,
135        Scratch<BE>: ScratchTakeCore<BE>,
136    {
137        module.bdd_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
138    }
139}
140
141pub struct BDDKeyPrepared<D, BRA, BE>
142where
143    D: Data,
144    BRA: BlindRotationAlgo,
145    BE: Backend,
146{
147    pub(crate) cbt: CircuitBootstrappingKeyPrepared<D, BRA, BE>,
148    pub(crate) ks: GLWEToLWEKeyPrepared<D, BE>,
149}
150
151impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> BDDKeyInfos for BDDKeyPrepared<D, BRA, BE> {
152    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
153        CircuitBootstrappingKeyLayout {
154            layout_brk: self.cbt.brk_infos(),
155            layout_atk: self.cbt.atk_infos(),
156            layout_tsk: self.cbt.tsk_infos(),
157        }
158    }
159    fn ks_infos(&self) -> GLWEToLWEKeyLayout {
160        GLWEToLWEKeyLayout {
161            n: self.ks.n(),
162            base2k: self.ks.base2k(),
163            k: self.ks.k(),
164            rank_in: self.ks.rank_in(),
165            dnum: self.ks.dnum(),
166        }
167    }
168}
169
170impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> GLWEAutomorphismKeyHelper<GLWEAutomorphismKeyPrepared<D, BE>, BE>
171    for BDDKeyPrepared<D, BRA, BE>
172{
173    fn automorphism_key_infos(&self) -> poulpy_core::layouts::GGLWELayout {
174        self.cbt.automorphism_key_infos()
175    }
176
177    fn get_automorphism_key(&self, k: i64) -> Option<&GLWEAutomorphismKeyPrepared<D, BE>> {
178        self.cbt.get_automorphism_key(k)
179    }
180}
181
182pub trait BDDKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend>
183where
184    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>,
185{
186    fn alloc_bdd_key_from_infos<A>(&self, infos: &A) -> BDDKeyPrepared<Vec<u8>, BRA, BE>
187    where
188        A: BDDKeyInfos,
189    {
190        BDDKeyPrepared {
191            cbt: CircuitBootstrappingKeyPrepared::alloc_from_infos(self, &infos.cbt_infos()),
192            ks: GLWEToLWEKeyPrepared::alloc_from_infos(self, &infos.ks_infos()),
193        }
194    }
195
196    fn prepare_bdd_key_tmp_bytes<A>(&self, infos: &A) -> usize
197    where
198        A: BDDKeyInfos,
199    {
200        self.circuit_bootstrapping_key_prepare_tmp_bytes(&infos.cbt_infos())
201            .max(self.prepare_glwe_to_lwe_key_tmp_bytes(&infos.ks_infos()))
202    }
203
204    fn prepare_bdd_key<DM, DR>(&self, res: &mut BDDKeyPrepared<DM, BRA, BE>, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
205    where
206        DM: DataMut,
207        DR: DataRef,
208        Scratch<BE>: ScratchTakeCore<BE>,
209    {
210        res.cbt.prepare(self, &other.cbt, scratch);
211        res.ks.prepare(self, &other.ks, scratch);
212    }
213}
214impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPreparedFactory<BRA, BE> for Module<BE> where
215    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>
216{
217}
218
219impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPrepared<Vec<u8>, BRA, BE> {
220    pub fn alloc_from_infos<M, A>(module: &M, infos: &A) -> Self
221    where
222        M: BDDKeyPreparedFactory<BRA, BE>,
223        A: BDDKeyInfos,
224    {
225        module.alloc_bdd_key_from_infos(infos)
226    }
227}
228
229impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> BDDKeyHelper<D, BRA, BE> for BDDKeyPrepared<D, BRA, BE> {
230    fn get_cbt_key(
231        &self,
232    ) -> (
233        &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
234        &GLWEToLWEKeyPrepared<D, BE>,
235    ) {
236        (&self.cbt, &self.ks)
237    }
238}
239
240pub trait BDDKeyHelper<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> {
241    fn get_cbt_key(
242        &self,
243    ) -> (
244        &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
245        &GLWEToLWEKeyPrepared<D, BE>,
246    );
247}
248
249pub trait FheUintPrepareDebug<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
250    fn fhe_uint_debug_prepare<DM, DR0, DR1>(
251        &self,
252        res: &mut FheUintPreparedDebug<DM, T>,
253        bits: &FheUint<DR0, T>,
254        key: &BDDKeyPrepared<DR1, BRA, BE>,
255        scratch: &mut Scratch<BE>,
256    ) where
257        DM: DataMut,
258        DR0: DataRef,
259        DR1: DataRef;
260}