1use poulpy_core::layouts::{
  2    AutomorphismKey, AutomorphismKeyLayout, GGLWEInfos, GGSWInfos, GLWE, GLWEInfos, GLWESecret, LWEInfos, LWESecret, TensorKey,
  3    TensorKeyLayout,
  4    prepared::{GLWEAutomorphismKeyPrepared, GLWESecretPrepared, TensorKeyPrepared},
  5};
  6use std::collections::HashMap;
  7
  8use poulpy_hal::{
  9    api::{
 10        ScratchAvailable, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace,
 11        VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform, VecZnxIdftApplyConsume,
 12        VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
 13        VecZnxSwitchRing, VmpPMatAlloc, VmpPrepare,
 14    },
 15    layouts::{Backend, Data, DataRef, Module, Scratch},
 16    source::Source,
 17};
 18
 19use crate::tfhe::blind_rotation::{
 20    BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyInfos,
 21    BlindRotationKeyLayout, BlindRotationKeyPrepared,
 22};
 23
 24pub trait CircuitBootstrappingKeyInfos {
 25    fn brk_infos(&self) -> BlindRotationKeyLayout;
 26    fn atk_infos(&self) -> AutomorphismKeyLayout;
 27    fn tsk_infos(&self) -> TensorKeyLayout;
 28}
 29
 30#[derive(Debug, Clone, Copy)]
 31pub struct CircuitBootstrappingKeyLayout {
 32    pub layout_brk: BlindRotationKeyLayout,
 33    pub layout_atk: AutomorphismKeyLayout,
 34    pub layout_tsk: TensorKeyLayout,
 35}
 36
 37impl CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyLayout {
 38    fn atk_infos(&self) -> AutomorphismKeyLayout {
 39        self.layout_atk
 40    }
 41
 42    fn brk_infos(&self) -> BlindRotationKeyLayout {
 43        self.layout_brk
 44    }
 45
 46    fn tsk_infos(&self) -> TensorKeyLayout {
 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: TensorKey<Vec<u8>>,
 71    pub(crate) atk: HashMap<i64, AutomorphismKey<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        + VecZnxDftBytesOf
 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        + SvpPPolBytesOf
 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    {
114        assert_eq!(sk_lwe.n(), cbt_infos.brk_infos().n_lwe());
115        assert_eq!(sk_glwe.n(), cbt_infos.brk_infos().n_glwe());
116        assert_eq!(sk_glwe.n(), cbt_infos.atk_infos().n());
117        assert_eq!(sk_glwe.n(), cbt_infos.tsk_infos().n());
118
119        let atk_infos: AutomorphismKeyLayout = cbt_infos.atk_infos();
120        let brk_infos: BlindRotationKeyLayout = cbt_infos.brk_infos();
121        let trk_infos: TensorKeyLayout = cbt_infos.tsk_infos();
122
123        let mut auto_keys: HashMap<i64, AutomorphismKey<Vec<u8>>> = HashMap::new();
124        let gal_els: Vec<i64> = GLWE::trace_galois_elements(module);
125        gal_els.iter().for_each(|gal_el| {
126            let mut key: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc_from_infos(&atk_infos);
127            key.encrypt_sk(module, *gal_el, sk_glwe, source_xa, source_xe, scratch);
128            auto_keys.insert(*gal_el, key);
129        });
130
131        let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch);
132
133        let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(&brk_infos);
134        brk.encrypt_sk(
135            module,
136            &sk_glwe_prepared,
137            sk_lwe,
138            source_xa,
139            source_xe,
140            scratch,
141        );
142
143        let mut tsk: TensorKey<Vec<u8>> = TensorKey::alloc_from_infos(&trk_infos);
144        tsk.encrypt_sk(module, sk_glwe, source_xa, source_xe, scratch);
145
146        Self {
147            brk,
148            atk: auto_keys,
149            tsk,
150        }
151    }
152}
153
154pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
155    pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
156    pub(crate) tsk: TensorKeyPrepared<Vec<u8>, B>,
157    pub(crate) atk: HashMap<i64, GLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
158}
159
160impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyPrepared<D, BRA, B> {
161    fn atk_infos(&self) -> AutomorphismKeyLayout {
162        let (_, atk) = self.atk.iter().next().expect("atk is empty");
163        AutomorphismKeyLayout {
164            n: atk.n(),
165            base2k: atk.base2k(),
166            k: atk.k(),
167            dnum: atk.dnum(),
168            dsize: atk.dsize(),
169            rank: atk.rank(),
170        }
171    }
172
173    fn brk_infos(&self) -> BlindRotationKeyLayout {
174        BlindRotationKeyLayout {
175            n_glwe: self.brk.n_glwe(),
176            n_lwe: self.brk.n_lwe(),
177            base2k: self.brk.base2k(),
178            k: self.brk.k(),
179            dnum: self.brk.dnum(),
180            rank: self.brk.rank(),
181        }
182    }
183
184    fn tsk_infos(&self) -> TensorKeyLayout {
185        TensorKeyLayout {
186            n: self.tsk.n(),
187            base2k: self.tsk.base2k(),
188            k: self.tsk.k(),
189            dnum: self.tsk.dnum(),
190            dsize: self.tsk.dsize(),
191            rank: self.tsk.rank(),
192        }
193    }
194}
195
196impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B>>
197    for CircuitBootstrappingKey<D, BRA>
198where
199    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
200    BlindRotationKey<D, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
201    TensorKey<D>: PrepareAlloc<B, TensorKeyPrepared<Vec<u8>, B>>,
202    AutomorphismKey<D>: PrepareAlloc<B, GLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
203{
204    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> {
205        let brk: BlindRotationKeyPrepared<Vec<u8>, BRA, B> = self.brk.prepare_alloc(module, scratch);
206        let tsk: TensorKeyPrepared<Vec<u8>, B> = self.tsk.prepare_alloc(module, scratch);
207        let mut atk: HashMap<i64, GLWEAutomorphismKeyPrepared<Vec<u8>, B>> = HashMap::new();
208        for (key, value) in &self.atk {
209            atk.insert(*key, value.prepare_alloc(module, scratch));
210        }
211        CircuitBootstrappingKeyPrepared { brk, tsk, atk }
212    }
213}