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