1use core::layouts::{
  2    GGLWEAutomorphismKey, GGLWETensorKey, GLWECiphertext, GLWESecret, LWESecret,
  3    prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, PrepareAlloc},
  4};
  5use std::{collections::HashMap, usize};
  6
  7use backend::hal::{
  8    api::{
  9        ScratchAvailable, TakeScalarZnx, TakeSvpPPol, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddScalarInplace,
 10        VecZnxAutomorphism, VecZnxSwithcDegree, VmpPMatAlloc, VmpPMatPrepare,
 11    },
 12    layouts::{Backend, Data, DataRef, Module, Scratch},
 13};
 14use sampling::source::Source;
 15
 16use core::trait_families::{
 17    GGLWEAutomorphismKeyEncryptSkFamily, GGLWETensorKeyEncryptSkFamily, GGSWEncryptSkFamily, GLWESecretPreparedModuleFamily,
 18};
 19
 20use crate::tfhe::blind_rotation::{
 21    BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyPrepared,
 22};
 23
 24pub trait CircuitBootstrappingKeyEncryptSk<B: Backend> {
 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>: GGSWEncryptSkFamily<B>
 56        + GLWESecretPreparedModuleFamily<B>
 57        + VecZnxAddScalarInplace
 58        + GGLWEAutomorphismKeyEncryptSkFamily<B>
 59        + VecZnxAutomorphism
 60        + VecZnxSwithcDegree
 61        + GGLWETensorKeyEncryptSkFamily<B>,
 62    Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
 63{
 64    fn encrypt_sk<DLwe, DGlwe>(
 65        module: &Module<B>,
 66        basek: usize,
 67        sk_lwe: &LWESecret<DLwe>,
 68        sk_glwe: &GLWESecret<DGlwe>,
 69        k_brk: usize,
 70        rows_brk: usize,
 71        k_trace: usize,
 72        rows_trace: usize,
 73        k_tsk: usize,
 74        rows_tsk: usize,
 75        source_xa: &mut Source,
 76        source_xe: &mut Source,
 77        sigma: f64,
 78        scratch: &mut Scratch<B>,
 79    ) -> Self
 80    where
 81        DLwe: DataRef,
 82        DGlwe: DataRef,
 83    {
 84        let mut auto_keys: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>> = HashMap::new();
 85        let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(&module);
 86        gal_els.iter().for_each(|gal_el| {
 87            let mut key: GGLWEAutomorphismKey<Vec<u8>> =
 88                GGLWEAutomorphismKey::alloc(sk_glwe.n(), basek, k_trace, rows_trace, 1, sk_glwe.rank());
 89            key.encrypt_sk(
 90                &module, *gal_el, &sk_glwe, source_xa, source_xe, sigma, scratch,
 91            );
 92            auto_keys.insert(*gal_el, key);
 93        });
 94
 95        let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch);
 96
 97        let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(
 98            sk_glwe.n(),
 99            sk_lwe.n(),
100            basek,
101            k_brk,
102            rows_brk,
103            sk_glwe.rank(),
104        );
105
106        brk.encrypt_sk(
107            module,
108            &sk_glwe_prepared,
109            sk_lwe,
110            source_xa,
111            source_xe,
112            sigma,
113            scratch,
114        );
115
116        let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(sk_glwe.n(), basek, k_tsk, rows_tsk, 1, sk_glwe.rank());
117        tsk.encrypt_sk(module, &sk_glwe, source_xa, source_xe, sigma, scratch);
118
119        Self {
120            brk,
121            atk: auto_keys,
122            tsk,
123        }
124    }
125}
126
127pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
128    pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
129    pub(crate) tsk: GGLWETensorKeyPrepared<Vec<u8>, B>,
130    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
131}
132
133impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B>>
134    for CircuitBootstrappingKey<D, BRA>
135where
136    Module<B>: VmpPMatAlloc<B> + VmpPMatPrepare<B>,
137    BlindRotationKey<D, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
138    GGLWETensorKey<D>: PrepareAlloc<B, GGLWETensorKeyPrepared<Vec<u8>, B>>,
139    GGLWEAutomorphismKey<D>: PrepareAlloc<B, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
140{
141    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> {
142        let brk: BlindRotationKeyPrepared<Vec<u8>, BRA, B> = self.brk.prepare_alloc(module, scratch);
143        let tsk: GGLWETensorKeyPrepared<Vec<u8>, B> = self.tsk.prepare_alloc(module, scratch);
144        let mut atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>> = HashMap::new();
145        for (key, value) in &self.atk {
146            atk.insert(*key, value.prepare_alloc(module, scratch));
147        }
148        CircuitBootstrappingKeyPrepared { brk, tsk, atk }
149    }
150}