1use poulpy_core::layouts::{
  2    GGLWEAutomorphismKey, GGLWETensorKey, GLWECiphertext, GLWESecret, LWESecret,
  3    prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, PrepareAlloc},
  4};
  5use std::collections::HashMap;
  6
  7use poulpy_hal::{
  8    api::{
  9        DFT, IDFTConsume, IDFTTmpA, ScratchAvailable, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
 10        TakeScalarZnx, TakeSvpPPol, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal,
 11        VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxFillUniform, VecZnxNormalize,
 12        VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwithcDegree, VmpPMatAlloc,
 13        VmpPrepare,
 14    },
 15    layouts::{Backend, Data, DataRef, Module, Scratch},
 16    source::Source,
 17};
 18
 19use crate::tfhe::blind_rotation::{
 20    BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyPrepared,
 21};
 22
 23pub trait CircuitBootstrappingKeyEncryptSk<B: Backend> {
 24    #[allow(clippy::too_many_arguments)]
 25    fn encrypt_sk<DLwe, DGlwe>(
 26        module: &Module<B>,
 27        basek: usize,
 28        sk_lwe: &LWESecret<DLwe>,
 29        sk_glwe: &GLWESecret<DGlwe>,
 30        k_brk: usize,
 31        rows_brk: usize,
 32        k_trace: usize,
 33        rows_trace: usize,
 34        k_tsk: usize,
 35        rows_tsk: usize,
 36        source_xa: &mut Source,
 37        source_xe: &mut Source,
 38        scratch: &mut Scratch<B>,
 39    ) -> Self
 40    where
 41        DLwe: DataRef,
 42        DGlwe: DataRef;
 43}
 44
 45pub struct CircuitBootstrappingKey<D: Data, BRA: BlindRotationAlgo> {
 46    pub(crate) brk: BlindRotationKey<D, BRA>,
 47    pub(crate) tsk: GGLWETensorKey<Vec<u8>>,
 48    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>>,
 49}
 50
 51impl<BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyEncryptSk<B> for CircuitBootstrappingKey<Vec<u8>, BRA>
 52where
 53    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
 54    Module<B>: SvpApply<B>
 55        + IDFTTmpA<B>
 56        + VecZnxAddScalarInplace
 57        + VecZnxDftAllocBytes
 58        + VecZnxBigNormalize<B>
 59        + DFT<B>
 60        + SvpApplyInplace<B>
 61        + IDFTConsume<B>
 62        + VecZnxNormalizeTmpBytes
 63        + VecZnxFillUniform
 64        + VecZnxSubABInplace
 65        + VecZnxAddInplace
 66        + VecZnxNormalizeInplace<B>
 67        + VecZnxAddNormal
 68        + VecZnxNormalize<B>
 69        + VecZnxSub
 70        + SvpPrepare<B>
 71        + VecZnxSwithcDegree
 72        + SvpPPolAllocBytes
 73        + SvpPPolAlloc<B>
 74        + VecZnxAutomorphism,
 75    Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
 76{
 77    fn encrypt_sk<DLwe, DGlwe>(
 78        module: &Module<B>,
 79        basek: usize,
 80        sk_lwe: &LWESecret<DLwe>,
 81        sk_glwe: &GLWESecret<DGlwe>,
 82        k_brk: usize,
 83        rows_brk: usize,
 84        k_trace: usize,
 85        rows_trace: usize,
 86        k_tsk: usize,
 87        rows_tsk: usize,
 88        source_xa: &mut Source,
 89        source_xe: &mut Source,
 90        scratch: &mut Scratch<B>,
 91    ) -> Self
 92    where
 93        DLwe: DataRef,
 94        DGlwe: DataRef,
 95        Module<B>:,
 96    {
 97        let mut auto_keys: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>> = HashMap::new();
 98        let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(module);
 99        gal_els.iter().for_each(|gal_el| {
100            let mut key: GGLWEAutomorphismKey<Vec<u8>> =
101                GGLWEAutomorphismKey::alloc(sk_glwe.n(), basek, k_trace, rows_trace, 1, sk_glwe.rank());
102            key.encrypt_sk(module, *gal_el, sk_glwe, source_xa, source_xe, scratch);
103            auto_keys.insert(*gal_el, key);
104        });
105
106        let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch);
107
108        let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(
109            sk_glwe.n(),
110            sk_lwe.n(),
111            basek,
112            k_brk,
113            rows_brk,
114            sk_glwe.rank(),
115        );
116
117        brk.encrypt_sk(
118            module,
119            &sk_glwe_prepared,
120            sk_lwe,
121            source_xa,
122            source_xe,
123            scratch,
124        );
125
126        let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(sk_glwe.n(), basek, k_tsk, rows_tsk, 1, sk_glwe.rank());
127        tsk.encrypt_sk(module, sk_glwe, source_xa, source_xe, scratch);
128
129        Self {
130            brk,
131            atk: auto_keys,
132            tsk,
133        }
134    }
135}
136
137pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
138    pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
139    pub(crate) tsk: GGLWETensorKeyPrepared<Vec<u8>, B>,
140    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
141}
142
143impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B>>
144    for CircuitBootstrappingKey<D, BRA>
145where
146    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
147    BlindRotationKey<D, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
148    GGLWETensorKey<D>: PrepareAlloc<B, GGLWETensorKeyPrepared<Vec<u8>, B>>,
149    GGLWEAutomorphismKey<D>: PrepareAlloc<B, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
150{
151    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> {
152        let brk: BlindRotationKeyPrepared<Vec<u8>, BRA, B> = self.brk.prepare_alloc(module, scratch);
153        let tsk: GGLWETensorKeyPrepared<Vec<u8>, B> = self.tsk.prepare_alloc(module, scratch);
154        let mut atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>> = HashMap::new();
155        for (key, value) in &self.atk {
156            atk.insert(*key, value.prepare_alloc(module, scratch));
157        }
158        CircuitBootstrappingKeyPrepared { brk, tsk, atk }
159    }
160}