1use poulpy_core::{
  2    Distribution, GGLWEToGGSWKeyEncryptSk, GLWEAutomorphismKeyEncryptSk, GetDistribution, ScratchTakeCore,
  3    layouts::{
  4        GGLWEInfos, GGLWEToGGSWKey, GGLWEToGGSWKeyLayout, GGSWInfos, GLWEAutomorphismKey, GLWEAutomorphismKeyLayout, GLWEInfos,
  5        GLWESecretPreparedFactory, GLWESecretToRef, LWEInfos, LWESecretToRef, prepared::GLWESecretPrepared,
  6    },
  7    trace_galois_elements,
  8};
  9use std::collections::HashMap;
 10
 11use poulpy_hal::{
 12    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 13    source::Source,
 14};
 15
 16use crate::tfhe::blind_rotation::{
 17    BlindRotationAlgo, BlindRotationKey, BlindRotationKeyEncryptSk, BlindRotationKeyFactory, BlindRotationKeyInfos,
 18    BlindRotationKeyLayout,
 19};
 20
 21pub trait CircuitBootstrappingKeyInfos {
 22    fn block_size(&self) -> usize;
 23    fn brk_infos(&self) -> BlindRotationKeyLayout;
 24    fn atk_infos(&self) -> GLWEAutomorphismKeyLayout;
 25    fn tsk_infos(&self) -> GGLWEToGGSWKeyLayout;
 26}
 27
 28#[derive(Debug, Clone, Copy)]
 29pub struct CircuitBootstrappingKeyLayout {
 30    pub layout_brk: BlindRotationKeyLayout,
 31    pub layout_atk: GLWEAutomorphismKeyLayout,
 32    pub layout_tsk: GGLWEToGGSWKeyLayout,
 33}
 34
 35impl CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyLayout {
 36    fn block_size(&self) -> usize {
 37        unimplemented!("unimplemented for CircuitBootstrappingKeyLayout")
 38    }
 39
 40    fn atk_infos(&self) -> GLWEAutomorphismKeyLayout {
 41        self.layout_atk
 42    }
 43
 44    fn brk_infos(&self) -> BlindRotationKeyLayout {
 45        self.layout_brk
 46    }
 47
 48    fn tsk_infos(&self) -> GGLWEToGGSWKeyLayout {
 49        self.layout_tsk
 50    }
 51}
 52
 53pub trait CircuitBootstrappingKeyEncryptSk<BRA: BlindRotationAlgo, BE: Backend> {
 54    fn circuit_bootstrapping_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
 55    where
 56        A: CircuitBootstrappingKeyInfos;
 57
 58    #[allow(clippy::too_many_arguments)]
 59    fn circuit_bootstrapping_key_encrypt_sk<D, S0, S1>(
 60        &self,
 61        res: &mut CircuitBootstrappingKey<D, BRA>,
 62        sk_lwe: &S0,
 63        sk_glwe: &S1,
 64        source_xa: &mut Source,
 65        source_xe: &mut Source,
 66        scratch: &mut Scratch<BE>,
 67    ) where
 68        D: DataMut,
 69        S0: LWESecretToRef + GetDistribution + LWEInfos,
 70        S1: GLWESecretToRef + GLWEInfos + GetDistribution;
 71}
 72
 73impl<BRA: BlindRotationAlgo> CircuitBootstrappingKey<Vec<u8>, BRA> {
 74    pub fn alloc_from_infos<A>(infos: &A) -> Self
 75    where
 76        A: CircuitBootstrappingKeyInfos,
 77        BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 78    {
 79        let atk_infos: &GLWEAutomorphismKeyLayout = &infos.atk_infos();
 80        let brk_infos: &BlindRotationKeyLayout = &infos.brk_infos();
 81        let trk_infos: &GGLWEToGGSWKeyLayout = &infos.tsk_infos();
 82        let gal_els: Vec<i64> = trace_galois_elements(atk_infos.log_n(), 2 * atk_infos.n().as_usize() as i64);
 83
 84        Self {
 85            brk: <BlindRotationKey<Vec<u8>, BRA> as BlindRotationKeyFactory<BRA>>::blind_rotation_key_alloc(brk_infos),
 86            atk: gal_els
 87                .iter()
 88                .map(|&gal_el| {
 89                    let key = GLWEAutomorphismKey::alloc_from_infos(atk_infos);
 90                    (gal_el, key)
 91                })
 92                .collect(),
 93            tsk: GGLWEToGGSWKey::alloc_from_infos(trk_infos),
 94        }
 95    }
 96}
 97
 98pub struct CircuitBootstrappingKey<D: Data, BRA: BlindRotationAlgo> {
 99    pub(crate) brk: BlindRotationKey<D, BRA>,
100    pub(crate) tsk: GGLWEToGGSWKey<Vec<u8>>,
101    pub(crate) atk: HashMap<i64, GLWEAutomorphismKey<Vec<u8>>>,
102}
103
104impl<D: DataMut, BRA: BlindRotationAlgo> CircuitBootstrappingKey<D, BRA> {
105    pub fn encrypt_sk<M, S0, S1, BE: Backend>(
106        &mut self,
107        module: &M,
108        sk_lwe: &S0,
109        sk_glwe: &S1,
110        source_xa: &mut Source,
111        source_xe: &mut Source,
112        scratch: &mut Scratch<BE>,
113    ) where
114        S0: LWESecretToRef + GetDistribution + LWEInfos,
115        S1: GLWESecretToRef + GLWEInfos + GetDistribution,
116        M: CircuitBootstrappingKeyEncryptSk<BRA, BE>,
117    {
118        module.circuit_bootstrapping_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
119    }
120}
121
122impl<BRA: BlindRotationAlgo, BE: Backend> CircuitBootstrappingKeyEncryptSk<BRA, BE> for Module<BE>
123where
124    Self: GGLWEToGGSWKeyEncryptSk<BE>
125        + BlindRotationKeyEncryptSk<BRA, BE>
126        + GLWEAutomorphismKeyEncryptSk<BE>
127        + GLWESecretPreparedFactory<BE>,
128    Scratch<BE>: ScratchTakeCore<BE>,
129{
130    fn circuit_bootstrapping_key_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
131    where
132        A: CircuitBootstrappingKeyInfos,
133    {
134        self.glwe_automorphism_key_encrypt_sk_tmp_bytes(&infos.atk_infos())
135            .max(self.blind_rotation_key_encrypt_sk_tmp_bytes(&infos.brk_infos()))
136            .max(self.gglwe_to_ggsw_key_encrypt_sk_tmp_bytes(&infos.tsk_infos()))
137    }
138
139    fn circuit_bootstrapping_key_encrypt_sk<D, S0, S1>(
140        &self,
141        res: &mut CircuitBootstrappingKey<D, BRA>,
142        sk_lwe: &S0,
143        sk_glwe: &S1,
144        source_xa: &mut Source,
145        source_xe: &mut Source,
146        scratch: &mut Scratch<BE>,
147    ) where
148        D: DataMut,
149        S0: LWESecretToRef + GetDistribution + LWEInfos,
150        S1: GLWESecretToRef + GLWEInfos + GetDistribution,
151    {
152        let brk_infos: &BlindRotationKeyLayout = &res.brk_infos();
153        let atk_infos: &GLWEAutomorphismKeyLayout = &res.atk_infos();
154        let tsk_infos: &GGLWEToGGSWKeyLayout = &res.tsk_infos();
155
156        assert_eq!(sk_lwe.n(), brk_infos.n_lwe());
157        assert_eq!(sk_glwe.n(), brk_infos.n_glwe());
158        assert_eq!(sk_glwe.n(), atk_infos.n());
159        assert_eq!(sk_glwe.n(), tsk_infos.n());
160
161        assert!(sk_glwe.dist() != &Distribution::NONE);
162
163        for (p, atk) in res.atk.iter_mut() {
164            atk.encrypt_sk(self, *p, sk_glwe, source_xa, source_xe, scratch);
165        }
166
167        let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(self, brk_infos.rank());
168        sk_glwe_prepared.prepare(self, sk_glwe);
169
170        res.brk.encrypt_sk(
171            self,
172            &sk_glwe_prepared,
173            sk_lwe,
174            source_xa,
175            source_xe,
176            scratch,
177        );
178
179        res.tsk
180            .encrypt_sk(self, sk_glwe, source_xa, source_xe, scratch);
181    }
182}
183
184impl<D: DataRef, BRA: BlindRotationAlgo> CircuitBootstrappingKeyInfos for CircuitBootstrappingKey<D, BRA> {
185    fn block_size(&self) -> usize {
186        self.brk.block_size()
187    }
188
189    fn atk_infos(&self) -> GLWEAutomorphismKeyLayout {
190        let (_, atk) = self.atk.iter().next().expect("atk is empty");
191        GLWEAutomorphismKeyLayout {
192            n: atk.n(),
193            base2k: atk.base2k(),
194            k: atk.k(),
195            dnum: atk.dnum(),
196            dsize: atk.dsize(),
197            rank: atk.rank(),
198        }
199    }
200
201    fn brk_infos(&self) -> BlindRotationKeyLayout {
202        BlindRotationKeyLayout {
203            n_glwe: self.brk.n_glwe(),
204            n_lwe: self.brk.n_lwe(),
205            base2k: self.brk.base2k(),
206            k: self.brk.k(),
207            dnum: self.brk.dnum(),
208            rank: self.brk.rank(),
209        }
210    }
211
212    fn tsk_infos(&self) -> GGLWEToGGSWKeyLayout {
213        GGLWEToGGSWKeyLayout {
214            n: self.tsk.n(),
215            base2k: self.tsk.base2k(),
216            k: self.tsk.k(),
217            dnum: self.tsk.dnum(),
218            dsize: self.tsk.dsize(),
219            rank: self.tsk.rank(),
220        }
221    }
222}