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        ScratchAvailable, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeSvpPPol,
 10        TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism,
 11        VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpA,
 12        VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace,
 13        VecZnxSwithcDegree, 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, 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        sigma: f64,
 39        scratch: &mut Scratch<B>,
 40    ) -> Self
 41    where
 42        DLwe: DataRef,
 43        DGlwe: DataRef;
 44}
 45
 46pub struct CircuitBootstrappingKey<D: Data, BRA: BlindRotationAlgo> {
 47    pub(crate) brk: BlindRotationKey<D, BRA>,
 48    pub(crate) tsk: GGLWETensorKey<Vec<u8>>,
 49    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>>,
 50}
 51
 52impl<BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyEncryptSk<B> for CircuitBootstrappingKey<Vec<u8>, BRA>
 53where
 54    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
 55    Module<B>: SvpApply<B>
 56        + VecZnxDftToVecZnxBigTmpA<B>
 57        + VecZnxAddScalarInplace
 58        + VecZnxDftAllocBytes
 59        + VecZnxBigNormalize<B>
 60        + VecZnxDftFromVecZnx<B>
 61        + SvpApplyInplace<B>
 62        + VecZnxDftToVecZnxBigConsume<B>
 63        + VecZnxNormalizeTmpBytes
 64        + VecZnxFillUniform
 65        + VecZnxSubABInplace
 66        + VecZnxAddInplace
 67        + VecZnxNormalizeInplace<B>
 68        + VecZnxAddNormal
 69        + VecZnxNormalize<B>
 70        + VecZnxSub
 71        + SvpPrepare<B>
 72        + VecZnxSwithcDegree
 73        + SvpPPolAllocBytes
 74        + SvpPPolAlloc<B>
 75        + VecZnxAutomorphism,
 76    Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
 77{
 78    fn encrypt_sk<DLwe, DGlwe>(
 79        module: &Module<B>,
 80        basek: usize,
 81        sk_lwe: &LWESecret<DLwe>,
 82        sk_glwe: &GLWESecret<DGlwe>,
 83        k_brk: usize,
 84        rows_brk: usize,
 85        k_trace: usize,
 86        rows_trace: usize,
 87        k_tsk: usize,
 88        rows_tsk: usize,
 89        source_xa: &mut Source,
 90        source_xe: &mut Source,
 91        sigma: f64,
 92        scratch: &mut Scratch<B>,
 93    ) -> Self
 94    where
 95        DLwe: DataRef,
 96        DGlwe: DataRef,
 97        Module<B>:,
 98    {
 99        let mut auto_keys: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>> = HashMap::new();
100        let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(module);
101        gal_els.iter().for_each(|gal_el| {
102            let mut key: GGLWEAutomorphismKey<Vec<u8>> =
103                GGLWEAutomorphismKey::alloc(sk_glwe.n(), basek, k_trace, rows_trace, 1, sk_glwe.rank());
104            key.encrypt_sk(
105                module, *gal_el, sk_glwe, source_xa, source_xe, sigma, scratch,
106            );
107            auto_keys.insert(*gal_el, key);
108        });
109
110        let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch);
111
112        let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(
113            sk_glwe.n(),
114            sk_lwe.n(),
115            basek,
116            k_brk,
117            rows_brk,
118            sk_glwe.rank(),
119        );
120
121        brk.encrypt_sk(
122            module,
123            &sk_glwe_prepared,
124            sk_lwe,
125            source_xa,
126            source_xe,
127            sigma,
128            scratch,
129        );
130
131        let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(sk_glwe.n(), basek, k_tsk, rows_tsk, 1, sk_glwe.rank());
132        tsk.encrypt_sk(module, sk_glwe, source_xa, source_xe, sigma, scratch);
133
134        Self {
135            brk,
136            atk: auto_keys,
137            tsk,
138        }
139    }
140}
141
142pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
143    pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
144    pub(crate) tsk: GGLWETensorKeyPrepared<Vec<u8>, B>,
145    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
146}
147
148impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B>>
149    for CircuitBootstrappingKey<D, BRA>
150where
151    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
152    BlindRotationKey<D, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
153    GGLWETensorKey<D>: PrepareAlloc<B, GGLWETensorKeyPrepared<Vec<u8>, B>>,
154    GGLWEAutomorphismKey<D>: PrepareAlloc<B, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
155{
156    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> {
157        let brk: BlindRotationKeyPrepared<Vec<u8>, BRA, B> = self.brk.prepare_alloc(module, scratch);
158        let tsk: GGLWETensorKeyPrepared<Vec<u8>, B> = self.tsk.prepare_alloc(module, scratch);
159        let mut atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>> = HashMap::new();
160        for (key, value) in &self.atk {
161            atk.insert(*key, value.prepare_alloc(module, scratch));
162        }
163        CircuitBootstrappingKeyPrepared { brk, tsk, atk }
164    }
165}