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 { layout_brk: self.cbt.brk_infos(), layout_atk: self.cbt.atk_infos(), layout_tsk: self.cbt.tsk_infos() }
142    }
143    fn ks_infos(&self) -> GLWEToLWEKeyLayout {
144        GLWEToLWEKeyLayout{
145            n: self.ks.n(),
146            base2k: self.ks.base2k(),
147            k: self.ks.k(),
148            rank_in: self.ks.rank_in(),
149            dnum: self.ks.dnum()
150        }
151    }
152}
153
154impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> GLWEAutomorphismKeyHelper<GLWEAutomorphismKeyPrepared<D, BE>, BE>
155    for BDDKeyPrepared<D, BRA, BE>
156{
157    fn automorphism_key_infos(&self) -> poulpy_core::layouts::GGLWELayout {
158        self.cbt.automorphism_key_infos()
159    }
160
161    fn get_automorphism_key(&self, k: i64) -> Option<&GLWEAutomorphismKeyPrepared<D, BE>> {
162        self.cbt.get_automorphism_key(k)
163    }
164}
165
166pub trait BDDKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend>
167where
168    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>,
169{
170    fn alloc_bdd_key_from_infos<A>(&self, infos: &A) -> BDDKeyPrepared<Vec<u8>, BRA, BE>
171    where
172        A: BDDKeyInfos,
173    {
174        BDDKeyPrepared {
175            cbt: CircuitBootstrappingKeyPrepared::alloc_from_infos(self, &infos.cbt_infos()),
176            ks: GLWEToLWEKeyPrepared::alloc_from_infos(self, &infos.ks_infos()),
177        }
178    }
179
180    fn prepare_bdd_key_tmp_bytes<A>(&self, infos: &A) -> usize
181    where
182        A: BDDKeyInfos,
183    {
184        self.circuit_bootstrapping_key_prepare_tmp_bytes(&infos.cbt_infos())
185            .max(self.prepare_glwe_to_lwe_key_tmp_bytes(&infos.ks_infos()))
186    }
187
188    fn prepare_bdd_key<DM, DR>(&self, res: &mut BDDKeyPrepared<DM, BRA, BE>, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
189    where
190        DM: DataMut,
191        DR: DataRef,
192        Scratch<BE>: ScratchTakeCore<BE>,
193    {
194        res.cbt.prepare(self, &other.cbt, scratch);
195        res.ks.prepare(self, &other.ks, scratch);
196    }
197}
198impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPreparedFactory<BRA, BE> for Module<BE> where
199    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>
200{
201}
202
203impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPrepared<Vec<u8>, BRA, BE> {
204    pub fn alloc_from_infos<M, A>(module: &M, infos: &A) -> Self
205    where
206        M: BDDKeyPreparedFactory<BRA, BE>,
207        A: BDDKeyInfos,
208    {
209        module.alloc_bdd_key_from_infos(infos)
210    }
211}
212
213impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> BDDKeyHelper<D, BRA, BE> for BDDKeyPrepared<D, BRA, BE> {
214    fn get_cbt_key(
215        &self,
216    ) -> (
217        &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
218        &GLWEToLWEKeyPrepared<D, BE>,
219    ) {
220        (&self.cbt, &self.ks)
221    }
222}
223
224pub trait BDDKeyHelper<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> {
225    fn get_cbt_key(
226        &self,
227    ) -> (
228        &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
229        &GLWEToLWEKeyPrepared<D, BE>,
230    );
231}
232
233pub trait FheUintPrepareDebug<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
234    fn fhe_uint_debug_prepare<DM, DR0, DR1>(
235        &self,
236        res: &mut FheUintPreparedDebug<DM, T>,
237        bits: &FheUint<DR0, T>,
238        key: &BDDKeyPrepared<DR1, BRA, BE>,
239        scratch: &mut Scratch<BE>,
240    ) where
241        DM: DataMut,
242        DR0: DataRef,
243        DR1: DataRef;
244}