1use poulpy_core::layouts::{
  2    GGLWEAutomorphismKey, GGLWEAutomorphismKeyLayout, GGLWELayoutInfos, GGLWETensorKey, GGLWETensorKeyLayout, GGSWInfos,
  3    GLWECiphertext, GLWEInfos, GLWESecret, LWEInfos, LWESecret,
  4    prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, PrepareAlloc},
  5};
  6use std::collections::HashMap;
  7
  8use poulpy_hal::{
  9    api::{
 10        ScratchAvailable, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx,
 11        TakeSvpPPol, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
 12        VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume,
 13        VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
 14        VecZnxSwitchRing, VmpPMatAlloc, VmpPrepare,
 15    },
 16    layouts::{Backend, Data, DataRef, Module, Scratch},
 17    source::Source,
 18};
 19
 20use crate::tfhe::blind_rotation::{
 21    BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyInfos,
 22    BlindRotationKeyLayout, BlindRotationKeyPrepared,
 23};
 24
 25pub trait CircuitBootstrappingKeyInfos {
 26    fn layout_brk(&self) -> BlindRotationKeyLayout;
 27    fn layout_atk(&self) -> GGLWEAutomorphismKeyLayout;
 28    fn layout_tsk(&self) -> GGLWETensorKeyLayout;
 29}
 30
 31pub struct CircuitBootstrappingKeyLayout {
 32    pub layout_brk: BlindRotationKeyLayout,
 33    pub layout_atk: GGLWEAutomorphismKeyLayout,
 34    pub layout_tsk: GGLWETensorKeyLayout,
 35}
 36
 37impl CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyLayout {
 38    fn layout_atk(&self) -> GGLWEAutomorphismKeyLayout {
 39        self.layout_atk
 40    }
 41
 42    fn layout_brk(&self) -> BlindRotationKeyLayout {
 43        self.layout_brk
 44    }
 45
 46    fn layout_tsk(&self) -> GGLWETensorKeyLayout {
 47        self.layout_tsk
 48    }
 49}
 50
 51pub trait CircuitBootstrappingKeyEncryptSk<B: Backend> {
 52    #[allow(clippy::too_many_arguments)]
 53    fn encrypt_sk<DLwe, DGlwe, INFOS>(
 54        module: &Module<B>,
 55        sk_lwe: &LWESecret<DLwe>,
 56        sk_glwe: &GLWESecret<DGlwe>,
 57        cbt_infos: &INFOS,
 58        source_xa: &mut Source,
 59        source_xe: &mut Source,
 60        scratch: &mut Scratch<B>,
 61    ) -> Self
 62    where
 63        INFOS: CircuitBootstrappingKeyInfos,
 64        DLwe: DataRef,
 65        DGlwe: DataRef;
 66}
 67
 68pub struct CircuitBootstrappingKey<D: Data, BRA: BlindRotationAlgo> {
 69    pub(crate) brk: BlindRotationKey<D, BRA>,
 70    pub(crate) tsk: GGLWETensorKey<Vec<u8>>,
 71    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>>,
 72}
 73
 74impl<BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyEncryptSk<B> for CircuitBootstrappingKey<Vec<u8>, BRA>
 75where
 76    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
 77    Module<B>: SvpApplyDftToDft<B>
 78        + VecZnxIdftApplyTmpA<B>
 79        + VecZnxAddScalarInplace
 80        + VecZnxDftAllocBytes
 81        + VecZnxBigNormalize<B>
 82        + VecZnxDftApply<B>
 83        + SvpApplyDftToDftInplace<B>
 84        + VecZnxIdftApplyConsume<B>
 85        + VecZnxNormalizeTmpBytes
 86        + VecZnxFillUniform
 87        + VecZnxSubInplace
 88        + VecZnxAddInplace
 89        + VecZnxNormalizeInplace<B>
 90        + VecZnxAddNormal
 91        + VecZnxNormalize<B>
 92        + VecZnxSub
 93        + SvpPrepare<B>
 94        + VecZnxSwitchRing
 95        + SvpPPolAllocBytes
 96        + SvpPPolAlloc<B>
 97        + VecZnxAutomorphism,
 98    Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
 99{
100    fn encrypt_sk<DLwe, DGlwe, INFOS>(
101        module: &Module<B>,
102        sk_lwe: &LWESecret<DLwe>,
103        sk_glwe: &GLWESecret<DGlwe>,
104        cbt_infos: &INFOS,
105        source_xa: &mut Source,
106        source_xe: &mut Source,
107        scratch: &mut Scratch<B>,
108    ) -> Self
109    where
110        INFOS: CircuitBootstrappingKeyInfos,
111        DLwe: DataRef,
112        DGlwe: DataRef,
113        Module<B>:,
114    {
115        assert_eq!(sk_lwe.n(), cbt_infos.layout_brk().n_lwe());
116        assert_eq!(sk_glwe.n(), cbt_infos.layout_brk().n_glwe());
117        assert_eq!(sk_glwe.n(), cbt_infos.layout_atk().n());
118        assert_eq!(sk_glwe.n(), cbt_infos.layout_tsk().n());
119
120        let atk_infos: GGLWEAutomorphismKeyLayout = cbt_infos.layout_atk();
121        let brk_infos: BlindRotationKeyLayout = cbt_infos.layout_brk();
122        let trk_infos: GGLWETensorKeyLayout = cbt_infos.layout_tsk();
123
124        let mut auto_keys: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>> = HashMap::new();
125        let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(module);
126        gal_els.iter().for_each(|gal_el| {
127            let mut key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(&atk_infos);
128            key.encrypt_sk(module, *gal_el, sk_glwe, source_xa, source_xe, scratch);
129            auto_keys.insert(*gal_el, key);
130        });
131
132        let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch);
133
134        let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(&brk_infos);
135        brk.encrypt_sk(
136            module,
137            &sk_glwe_prepared,
138            sk_lwe,
139            source_xa,
140            source_xe,
141            scratch,
142        );
143
144        let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(&trk_infos);
145        tsk.encrypt_sk(module, sk_glwe, source_xa, source_xe, scratch);
146
147        Self {
148            brk,
149            atk: auto_keys,
150            tsk,
151        }
152    }
153}
154
155pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
156    pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
157    pub(crate) tsk: GGLWETensorKeyPrepared<Vec<u8>, B>,
158    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
159}
160
161impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyPrepared<D, BRA, B> {
162    fn layout_atk(&self) -> GGLWEAutomorphismKeyLayout {
163        let (_, atk) = self.atk.iter().next().expect("atk is empty");
164        GGLWEAutomorphismKeyLayout {
165            n: atk.n(),
166            base2k: atk.base2k(),
167            k: atk.k(),
168            rows: atk.rows(),
169            digits: atk.digits(),
170            rank: atk.rank(),
171        }
172    }
173
174    fn layout_brk(&self) -> BlindRotationKeyLayout {
175        BlindRotationKeyLayout {
176            n_glwe: self.brk.n_glwe(),
177            n_lwe: self.brk.n_lwe(),
178            base2k: self.brk.base2k(),
179            k: self.brk.k(),
180            rows: self.brk.rows(),
181            rank: self.brk.rank(),
182        }
183    }
184
185    fn layout_tsk(&self) -> GGLWETensorKeyLayout {
186        GGLWETensorKeyLayout {
187            n: self.tsk.n(),
188            base2k: self.tsk.base2k(),
189            k: self.tsk.k(),
190            rows: self.tsk.rows(),
191            digits: self.tsk.digits(),
192            rank: self.tsk.rank(),
193        }
194    }
195}
196
197impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B>>
198    for CircuitBootstrappingKey<D, BRA>
199where
200    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
201    BlindRotationKey<D, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
202    GGLWETensorKey<D>: PrepareAlloc<B, GGLWETensorKeyPrepared<Vec<u8>, B>>,
203    GGLWEAutomorphismKey<D>: PrepareAlloc<B, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
204{
205    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> {
206        let brk: BlindRotationKeyPrepared<Vec<u8>, BRA, B> = self.brk.prepare_alloc(module, scratch);
207        let tsk: GGLWETensorKeyPrepared<Vec<u8>, B> = self.tsk.prepare_alloc(module, scratch);
208        let mut atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>> = HashMap::new();
209        for (key, value) in &self.atk {
210            atk.insert(*key, value.prepare_alloc(module, scratch));
211        }
212        CircuitBootstrappingKeyPrepared { brk, tsk, atk }
213    }
214}