1use poulpy_core::layouts::{
  2    AutomorphismKey, AutomorphismKeyLayout, GGLWEInfos, GGSWInfos, GLWE, GLWEInfos, GLWESecret, LWEInfos, LWESecret, TensorKey,
  3    TensorKeyLayout,
  4    prepared::{AutomorphismKeyPrepared, GLWESecretPrepared, PrepareAlloc, TensorKeyPrepared},
  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 brk_infos(&self) -> BlindRotationKeyLayout;
 27    fn atk_infos(&self) -> AutomorphismKeyLayout;
 28    fn tsk_infos(&self) -> TensorKeyLayout;
 29}
 30
 31#[derive(Debug, Clone, Copy)]
 32pub struct CircuitBootstrappingKeyLayout {
 33    pub layout_brk: BlindRotationKeyLayout,
 34    pub layout_atk: AutomorphismKeyLayout,
 35    pub layout_tsk: TensorKeyLayout,
 36}
 37
 38impl CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyLayout {
 39    fn atk_infos(&self) -> AutomorphismKeyLayout {
 40        self.layout_atk
 41    }
 42
 43    fn brk_infos(&self) -> BlindRotationKeyLayout {
 44        self.layout_brk
 45    }
 46
 47    fn tsk_infos(&self) -> TensorKeyLayout {
 48        self.layout_tsk
 49    }
 50}
 51
 52pub trait CircuitBootstrappingKeyEncryptSk<B: Backend> {
 53    #[allow(clippy::too_many_arguments)]
 54    fn encrypt_sk<DLwe, DGlwe, INFOS>(
 55        module: &Module<B>,
 56        sk_lwe: &LWESecret<DLwe>,
 57        sk_glwe: &GLWESecret<DGlwe>,
 58        cbt_infos: &INFOS,
 59        source_xa: &mut Source,
 60        source_xe: &mut Source,
 61        scratch: &mut Scratch<B>,
 62    ) -> Self
 63    where
 64        INFOS: CircuitBootstrappingKeyInfos,
 65        DLwe: DataRef,
 66        DGlwe: DataRef;
 67}
 68
 69pub struct CircuitBootstrappingKey<D: Data, BRA: BlindRotationAlgo> {
 70    pub(crate) brk: BlindRotationKey<D, BRA>,
 71    pub(crate) tsk: TensorKey<Vec<u8>>,
 72    pub(crate) atk: HashMap<i64, AutomorphismKey<Vec<u8>>>,
 73}
 74
 75impl<BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyEncryptSk<B> for CircuitBootstrappingKey<Vec<u8>, BRA>
 76where
 77    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
 78    Module<B>: SvpApplyDftToDft<B>
 79        + VecZnxIdftApplyTmpA<B>
 80        + VecZnxAddScalarInplace
 81        + VecZnxDftAllocBytes
 82        + VecZnxBigNormalize<B>
 83        + VecZnxDftApply<B>
 84        + SvpApplyDftToDftInplace<B>
 85        + VecZnxIdftApplyConsume<B>
 86        + VecZnxNormalizeTmpBytes
 87        + VecZnxFillUniform
 88        + VecZnxSubInplace
 89        + VecZnxAddInplace
 90        + VecZnxNormalizeInplace<B>
 91        + VecZnxAddNormal
 92        + VecZnxNormalize<B>
 93        + VecZnxSub
 94        + SvpPrepare<B>
 95        + VecZnxSwitchRing
 96        + SvpPPolAllocBytes
 97        + SvpPPolAlloc<B>
 98        + VecZnxAutomorphism,
 99    Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
100{
101    fn encrypt_sk<DLwe, DGlwe, INFOS>(
102        module: &Module<B>,
103        sk_lwe: &LWESecret<DLwe>,
104        sk_glwe: &GLWESecret<DGlwe>,
105        cbt_infos: &INFOS,
106        source_xa: &mut Source,
107        source_xe: &mut Source,
108        scratch: &mut Scratch<B>,
109    ) -> Self
110    where
111        INFOS: CircuitBootstrappingKeyInfos,
112        DLwe: DataRef,
113        DGlwe: DataRef,
114    {
115        assert_eq!(sk_lwe.n(), cbt_infos.brk_infos().n_lwe());
116        assert_eq!(sk_glwe.n(), cbt_infos.brk_infos().n_glwe());
117        assert_eq!(sk_glwe.n(), cbt_infos.atk_infos().n());
118        assert_eq!(sk_glwe.n(), cbt_infos.tsk_infos().n());
119
120        let atk_infos: AutomorphismKeyLayout = cbt_infos.atk_infos();
121        let brk_infos: BlindRotationKeyLayout = cbt_infos.brk_infos();
122        let trk_infos: TensorKeyLayout = cbt_infos.tsk_infos();
123
124        let mut auto_keys: HashMap<i64, AutomorphismKey<Vec<u8>>> = HashMap::new();
125        let gal_els: Vec<i64> = GLWE::trace_galois_elements(module);
126        gal_els.iter().for_each(|gal_el| {
127            let mut key: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc_from_infos(&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: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(&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: TensorKeyPrepared<Vec<u8>, B>,
158    pub(crate) atk: HashMap<i64, AutomorphismKeyPrepared<Vec<u8>, B>>,
159}
160
161impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyPrepared<D, BRA, B> {
162    fn atk_infos(&self) -> AutomorphismKeyLayout {
163        let (_, atk) = self.atk.iter().next().expect("atk is empty");
164        AutomorphismKeyLayout {
165            n: atk.n(),
166            base2k: atk.base2k(),
167            k: atk.k(),
168            dnum: atk.dnum(),
169            dsize: atk.dsize(),
170            rank: atk.rank(),
171        }
172    }
173
174    fn brk_infos(&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            dnum: self.brk.dnum(),
181            rank: self.brk.rank(),
182        }
183    }
184
185    fn tsk_infos(&self) -> TensorKeyLayout {
186        TensorKeyLayout {
187            n: self.tsk.n(),
188            base2k: self.tsk.base2k(),
189            k: self.tsk.k(),
190            dnum: self.tsk.dnum(),
191            dsize: self.tsk.dsize(),
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    TensorKey<D>: PrepareAlloc<B, TensorKeyPrepared<Vec<u8>, B>>,
203    AutomorphismKey<D>: PrepareAlloc<B, AutomorphismKeyPrepared<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: TensorKeyPrepared<Vec<u8>, B> = self.tsk.prepare_alloc(module, scratch);
208        let mut atk: HashMap<i64, AutomorphismKeyPrepared<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}