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