1use crate::tfhe::bdd_arithmetic::FheUintPreparedDebug;
  2use crate::tfhe::{
  3    bdd_arithmetic::{FheUint, FheUintPrepared, UnsignedInteger},
  4    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory},
  5    circuit_bootstrapping::{
  6        CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
  7        CircuitBootstrappingKeyPrepared, CircuitBootstrappingKeyPreparedFactory, CirtuitBootstrappingExecute,
  8    },
  9};
 10
 11use poulpy_core::{
 12    GLWEToLWESwitchingKeyEncryptSk, GetDistribution, LWEFromGLWE, ScratchTakeCore,
 13    layouts::{
 14        GGSWInfos, GGSWPreparedFactory, GLWEInfos, GLWESecretToRef, GLWEToLWEKey, GLWEToLWEKeyLayout,
 15        GLWEToLWEKeyPreparedFactory, LWE, LWEInfos, LWESecretToRef, prepared::GLWEToLWEKeyPrepared,
 16    },
 17};
 18use poulpy_hal::{
 19    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 20    source::Source,
 21};
 22
 23pub trait BDDKeyInfos {
 24    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout;
 25    fn ks_infos(&self) -> GLWEToLWEKeyLayout;
 26}
 27
 28#[derive(Debug, Clone, Copy)]
 29pub struct BDDKeyLayout {
 30    pub cbt: CircuitBootstrappingKeyLayout,
 31    pub ks: GLWEToLWEKeyLayout,
 32}
 33
 34impl BDDKeyInfos for BDDKeyLayout {
 35    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
 36        self.cbt
 37    }
 38
 39    fn ks_infos(&self) -> GLWEToLWEKeyLayout {
 40        self.ks
 41    }
 42}
 43
 44pub struct BDDKey<D, BRA>
 45where
 46    D: Data,
 47    BRA: BlindRotationAlgo,
 48{
 49    pub(crate) cbt: CircuitBootstrappingKey<D, BRA>,
 50    pub(crate) ks: GLWEToLWEKey<D>,
 51}
 52
 53impl<BRA: BlindRotationAlgo> BDDKey<Vec<u8>, BRA>
 54where
 55    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 56{
 57    pub fn alloc_from_infos<A>(infos: &A) -> Self
 58    where
 59        A: BDDKeyInfos,
 60    {
 61        Self {
 62            cbt: CircuitBootstrappingKey::alloc_from_infos(&infos.cbt_infos()),
 63            ks: GLWEToLWEKey::alloc_from_infos(&infos.ks_infos()),
 64        }
 65    }
 66}
 67
 68pub trait BDDKeyEncryptSk<BRA: BlindRotationAlgo, BE: Backend> {
 69    fn bdd_key_encrypt_sk<D, S0, S1>(
 70        &self,
 71        res: &mut BDDKey<D, BRA>,
 72        sk_lwe: &S0,
 73        sk_glwe: &S1,
 74        source_xa: &mut Source,
 75        source_xe: &mut Source,
 76        scratch: &mut Scratch<BE>,
 77    ) where
 78        D: DataMut,
 79        S0: LWESecretToRef + GetDistribution + LWEInfos,
 80        S1: GLWESecretToRef + GetDistribution + GLWEInfos;
 81}
 82
 83impl<BE: Backend, BRA: BlindRotationAlgo> BDDKeyEncryptSk<BRA, BE> for Module<BE>
 84where
 85    Self: CircuitBootstrappingKeyEncryptSk<BRA, BE> + GLWEToLWESwitchingKeyEncryptSk<BE>,
 86    Scratch<BE>: ScratchTakeCore<BE>,
 87{
 88    fn bdd_key_encrypt_sk<D, S0, S1>(
 89        &self,
 90        res: &mut BDDKey<D, BRA>,
 91        sk_lwe: &S0,
 92        sk_glwe: &S1,
 93        source_xa: &mut Source,
 94        source_xe: &mut Source,
 95        scratch: &mut Scratch<BE>,
 96    ) where
 97        D: DataMut,
 98        S0: LWESecretToRef + GetDistribution + LWEInfos,
 99        S1: GLWESecretToRef + GetDistribution + GLWEInfos,
100    {
101        res.ks
102            .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
103        res.cbt
104            .encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
105    }
106}
107
108impl<D: DataMut, BRA: BlindRotationAlgo> BDDKey<D, BRA> {
109    pub fn encrypt_sk<S0, S1, M, BE: Backend>(
110        &mut self,
111        module: &M,
112        sk_lwe: &S0,
113        sk_glwe: &S1,
114        source_xa: &mut Source,
115        source_xe: &mut Source,
116        scratch: &mut Scratch<BE>,
117    ) where
118        S0: LWESecretToRef + GetDistribution + LWEInfos,
119        S1: GLWESecretToRef + GetDistribution + GLWEInfos,
120        M: BDDKeyEncryptSk<BRA, BE>,
121        Scratch<BE>: ScratchTakeCore<BE>,
122    {
123        module.bdd_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
124    }
125}
126
127pub struct BDDKeyPrepared<D, BRA, BE>
128where
129    D: Data,
130    BRA: BlindRotationAlgo,
131    BE: Backend,
132{
133    pub(crate) cbt: CircuitBootstrappingKeyPrepared<D, BRA, BE>,
134    pub(crate) ks: GLWEToLWEKeyPrepared<D, BE>,
135}
136
137pub trait BDDKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend>
138where
139    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>,
140{
141    fn alloc_bdd_key_from_infos<A>(&self, infos: &A) -> BDDKeyPrepared<Vec<u8>, BRA, BE>
142    where
143        A: BDDKeyInfos,
144    {
145        BDDKeyPrepared {
146            cbt: CircuitBootstrappingKeyPrepared::alloc_from_infos(self, &infos.cbt_infos()),
147            ks: GLWEToLWEKeyPrepared::alloc_from_infos(self, &infos.ks_infos()),
148        }
149    }
150
151    fn prepare_bdd_key_tmp_bytes<A>(&self, infos: &A) -> usize
152    where
153        A: BDDKeyInfos,
154    {
155        self.circuit_bootstrapping_key_prepare_tmp_bytes(&infos.cbt_infos())
156            .max(self.prepare_glwe_to_lwe_key_tmp_bytes(&infos.ks_infos()))
157    }
158
159    fn prepare_bdd_key<DM, DR>(&self, res: &mut BDDKeyPrepared<DM, BRA, BE>, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
160    where
161        DM: DataMut,
162        DR: DataRef,
163        Scratch<BE>: ScratchTakeCore<BE>,
164    {
165        res.cbt.prepare(self, &other.cbt, scratch);
166        res.ks.prepare(self, &other.ks, scratch);
167    }
168}
169impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPreparedFactory<BRA, BE> for Module<BE> where
170    Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>
171{
172}
173
174impl<BRA: BlindRotationAlgo, BE: Backend> BDDKeyPrepared<Vec<u8>, BRA, BE> {
175    pub fn alloc_from_infos<M, A>(module: &M, infos: &A) -> Self
176    where
177        M: BDDKeyPreparedFactory<BRA, BE>,
178        A: BDDKeyInfos,
179    {
180        module.alloc_bdd_key_from_infos(infos)
181    }
182}
183
184impl<D: DataMut, BRA: BlindRotationAlgo, BE: Backend> BDDKeyPrepared<D, BRA, BE> {
185    pub fn prepare<DR, M>(&mut self, module: &M, other: &BDDKey<DR, BRA>, scratch: &mut Scratch<BE>)
186    where
187        DR: DataRef,
188        M: BDDKeyPreparedFactory<BRA, BE>,
189        Scratch<BE>: ScratchTakeCore<BE>,
190    {
191        module.prepare_bdd_key(self, other, scratch);
192    }
193}
194
195pub trait FheUintBlocksPrepare<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
196    fn fhe_uint_prepare_tmp_bytes<R, A>(&self, block_size: usize, extension_factor: usize, res_infos: &R, infos: &A) -> usize
197    where
198        R: GGSWInfos,
199        A: BDDKeyInfos;
200    fn fhe_uint_prepare<DM, DR0, DR1>(
201        &self,
202        res: &mut FheUintPrepared<DM, T, BE>,
203        bits: &FheUint<DR0, T>,
204        key: &BDDKeyPrepared<DR1, BRA, BE>,
205        scratch: &mut Scratch<BE>,
206    ) where
207        DM: DataMut,
208        DR0: DataRef,
209        DR1: DataRef;
210}
211
212impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintBlocksPrepare<BRA, T, BE> for Module<BE>
213where
214    Self: LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
215    Scratch<BE>: ScratchTakeCore<BE>,
216{
217    fn fhe_uint_prepare_tmp_bytes<R, A>(&self, block_size: usize, extension_factor: usize, res_infos: &R, bdd_infos: &A) -> usize
218    where
219        R: GGSWInfos,
220        A: BDDKeyInfos,
221    {
222        self.circuit_bootstrapping_execute_tmp_bytes(
223            block_size,
224            extension_factor,
225            res_infos,
226            &bdd_infos.cbt_infos(),
227        )
228    }
229
230    fn fhe_uint_prepare<DM, DR0, DR1>(
231        &self,
232        res: &mut FheUintPrepared<DM, T, BE>,
233        bits: &FheUint<DR0, T>,
234        key: &BDDKeyPrepared<DR1, BRA, BE>,
235        scratch: &mut Scratch<BE>,
236    ) where
237        DM: DataMut,
238        DR0: DataRef,
239        DR1: DataRef,
240    {
241        let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(bits); //TODO: add TakeLWE
242        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
243        for (bit, dst) in res.bits.iter_mut().enumerate() {
244            bits.get_bit(self, bit, &mut lwe, &key.ks, scratch_1);
245            key.cbt
246                .execute_to_constant(self, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
247            dst.prepare(self, &tmp_ggsw, scratch_1);
248        }
249    }
250}
251
252impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
253    pub fn prepare<BRA, M, O, K>(
254        &mut self,
255        module: &M,
256        other: &FheUint<O, T>,
257        key: &BDDKeyPrepared<K, BRA, BE>,
258        scratch: &mut Scratch<BE>,
259    ) where
260        BRA: BlindRotationAlgo,
261        O: DataRef,
262        K: DataRef,
263        M: FheUintBlocksPrepare<BRA, T, BE>,
264        Scratch<BE>: ScratchTakeCore<BE>,
265    {
266        module.fhe_uint_prepare(self, other, key, scratch);
267    }
268}
269
270pub trait FheUintBlockDebugPrepare<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
271    fn fhe_uint_debug_prepare<DM, DR0, DR1>(
272        &self,
273        res: &mut FheUintPreparedDebug<DM, T>,
274        bits: &FheUint<DR0, T>,
275        key: &BDDKeyPrepared<DR1, BRA, BE>,
276        scratch: &mut Scratch<BE>,
277    ) where
278        DM: DataMut,
279        DR0: DataRef,
280        DR1: DataRef;
281}