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    #[allow(clippy::too_many_arguments)]
 55    fn circuit_bootstrapping_key_encrypt_sk<D, S0, S1>(
 56        &self,
 57        res: &mut CircuitBootstrappingKey<D, BRA>,
 58        sk_lwe: &S0,
 59        sk_glwe: &S1,
 60        source_xa: &mut Source,
 61        source_xe: &mut Source,
 62        scratch: &mut Scratch<BE>,
 63    ) where
 64        D: DataMut,
 65        S0: LWESecretToRef + GetDistribution + LWEInfos,
 66        S1: GLWESecretToRef + GLWEInfos + GetDistribution;
 67}
 68
 69impl<BRA: BlindRotationAlgo> CircuitBootstrappingKey<Vec<u8>, BRA> {
 70    pub fn alloc_from_infos<A>(infos: &A) -> Self
 71    where
 72        A: CircuitBootstrappingKeyInfos,
 73        BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 74    {
 75        let atk_infos: &GLWEAutomorphismKeyLayout = &infos.atk_infos();
 76        let brk_infos: &BlindRotationKeyLayout = &infos.brk_infos();
 77        let trk_infos: &GGLWEToGGSWKeyLayout = &infos.tsk_infos();
 78        let gal_els: Vec<i64> = trace_galois_elements(atk_infos.log_n(), 2 * atk_infos.n().as_usize() as i64);
 79
 80        Self {
 81            brk: <BlindRotationKey<Vec<u8>, BRA> as BlindRotationKeyFactory<BRA>>::blind_rotation_key_alloc(brk_infos),
 82            atk: gal_els
 83                .iter()
 84                .map(|&gal_el| {
 85                    let key = GLWEAutomorphismKey::alloc_from_infos(atk_infos);
 86                    (gal_el, key)
 87                })
 88                .collect(),
 89            tsk: GGLWEToGGSWKey::alloc_from_infos(trk_infos),
 90        }
 91    }
 92}
 93
 94pub struct CircuitBootstrappingKey<D: Data, BRA: BlindRotationAlgo> {
 95    pub(crate) brk: BlindRotationKey<D, BRA>,
 96    pub(crate) tsk: GGLWEToGGSWKey<Vec<u8>>,
 97    pub(crate) atk: HashMap<i64, GLWEAutomorphismKey<Vec<u8>>>,
 98}
 99
100impl<D: DataMut, BRA: BlindRotationAlgo> CircuitBootstrappingKey<D, BRA> {
101    pub fn encrypt_sk<M, S0, S1, BE: Backend>(
102        &mut self,
103        module: &M,
104        sk_lwe: &S0,
105        sk_glwe: &S1,
106        source_xa: &mut Source,
107        source_xe: &mut Source,
108        scratch: &mut Scratch<BE>,
109    ) where
110        S0: LWESecretToRef + GetDistribution + LWEInfos,
111        S1: GLWESecretToRef + GLWEInfos + GetDistribution,
112        M: CircuitBootstrappingKeyEncryptSk<BRA, BE>,
113    {
114        module.circuit_bootstrapping_key_encrypt_sk(self, sk_lwe, sk_glwe, source_xa, source_xe, scratch);
115    }
116}
117
118impl<BRA: BlindRotationAlgo, BE: Backend> CircuitBootstrappingKeyEncryptSk<BRA, BE> for Module<BE>
119where
120    Self: GGLWEToGGSWKeyEncryptSk<BE>
121        + BlindRotationKeyEncryptSk<BRA, BE>
122        + GLWEAutomorphismKeyEncryptSk<BE>
123        + GLWESecretPreparedFactory<BE>,
124    Scratch<BE>: ScratchTakeCore<BE>,
125{
126    fn circuit_bootstrapping_key_encrypt_sk<D, S0, S1>(
127        &self,
128        res: &mut CircuitBootstrappingKey<D, BRA>,
129        sk_lwe: &S0,
130        sk_glwe: &S1,
131        source_xa: &mut Source,
132        source_xe: &mut Source,
133        scratch: &mut Scratch<BE>,
134    ) where
135        D: DataMut,
136        S0: LWESecretToRef + GetDistribution + LWEInfos,
137        S1: GLWESecretToRef + GLWEInfos + GetDistribution,
138    {
139        let brk_infos: &BlindRotationKeyLayout = &res.brk_infos();
140        let atk_infos: &GLWEAutomorphismKeyLayout = &res.atk_infos();
141        let tsk_infos: &GGLWEToGGSWKeyLayout = &res.tsk_infos();
142
143        assert_eq!(sk_lwe.n(), brk_infos.n_lwe());
144        assert_eq!(sk_glwe.n(), brk_infos.n_glwe());
145        assert_eq!(sk_glwe.n(), atk_infos.n());
146        assert_eq!(sk_glwe.n(), tsk_infos.n());
147
148        assert!(sk_glwe.dist() != &Distribution::NONE);
149
150        for (p, atk) in res.atk.iter_mut() {
151            atk.encrypt_sk(self, *p, sk_glwe, source_xa, source_xe, scratch);
152        }
153
154        let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(self, brk_infos.rank());
155        sk_glwe_prepared.prepare(self, sk_glwe);
156
157        res.brk.encrypt_sk(
158            self,
159            &sk_glwe_prepared,
160            sk_lwe,
161            source_xa,
162            source_xe,
163            scratch,
164        );
165
166        res.tsk
167            .encrypt_sk(self, sk_glwe, source_xa, source_xe, scratch);
168    }
169}
170
171impl<D: DataRef, BRA: BlindRotationAlgo> CircuitBootstrappingKeyInfos for CircuitBootstrappingKey<D, BRA> {
172    fn block_size(&self) -> usize {
173        self.brk.block_size()
174    }
175
176    fn atk_infos(&self) -> GLWEAutomorphismKeyLayout {
177        let (_, atk) = self.atk.iter().next().expect("atk is empty");
178        GLWEAutomorphismKeyLayout {
179            n: atk.n(),
180            base2k: atk.base2k(),
181            k: atk.k(),
182            dnum: atk.dnum(),
183            dsize: atk.dsize(),
184            rank: atk.rank(),
185        }
186    }
187
188    fn brk_infos(&self) -> BlindRotationKeyLayout {
189        BlindRotationKeyLayout {
190            n_glwe: self.brk.n_glwe(),
191            n_lwe: self.brk.n_lwe(),
192            base2k: self.brk.base2k(),
193            k: self.brk.k(),
194            dnum: self.brk.dnum(),
195            rank: self.brk.rank(),
196        }
197    }
198
199    fn tsk_infos(&self) -> GGLWEToGGSWKeyLayout {
200        GGLWEToGGSWKeyLayout {
201            n: self.tsk.n(),
202            base2k: self.tsk.base2k(),
203            k: self.tsk.k(),
204            dnum: self.tsk.dnum(),
205            dsize: self.tsk.dsize(),
206            rank: self.tsk.rank(),
207        }
208    }
209}