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, 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};
 17use sampling::source::Source;
 18
 19use crate::tfhe::blind_rotation::{
 20    BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyPrepared,
 21};
 22
 23pub trait CircuitBootstrappingKeyEncryptSk<B: Backend> {
 24    fn encrypt_sk<DLwe, DGlwe>(
 25        module: &Module<B>,
 26        basek: usize,
 27        sk_lwe: &LWESecret<DLwe>,
 28        sk_glwe: &GLWESecret<DGlwe>,
 29        k_brk: usize,
 30        rows_brk: usize,
 31        k_trace: usize,
 32        rows_trace: usize,
 33        k_tsk: usize,
 34        rows_tsk: usize,
 35        source_xa: &mut Source,
 36        source_xe: &mut Source,
 37        sigma: f64,
 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        + VecZnxDftToVecZnxBigTmpA<B>
 56        + VecZnxAddScalarInplace
 57        + VecZnxDftAllocBytes
 58        + VecZnxBigNormalize<B>
 59        + VecZnxDftFromVecZnx<B>
 60        + SvpApplyInplace<B>
 61        + VecZnxDftToVecZnxBigConsume<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        sigma: f64,
 91        scratch: &mut Scratch<B>,
 92    ) -> Self
 93    where
 94        DLwe: DataRef,
 95        DGlwe: DataRef,
 96        Module<B>:,
 97    {
 98        let mut auto_keys: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>> = HashMap::new();
 99        let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(&module);
100        gal_els.iter().for_each(|gal_el| {
101            let mut key: GGLWEAutomorphismKey<Vec<u8>> =
102                GGLWEAutomorphismKey::alloc(sk_glwe.n(), basek, k_trace, rows_trace, 1, sk_glwe.rank());
103            key.encrypt_sk(
104                &module, *gal_el, &sk_glwe, source_xa, source_xe, sigma, scratch,
105            );
106            auto_keys.insert(*gal_el, key);
107        });
108
109        let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch);
110
111        let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(
112            sk_glwe.n(),
113            sk_lwe.n(),
114            basek,
115            k_brk,
116            rows_brk,
117            sk_glwe.rank(),
118        );
119
120        brk.encrypt_sk(
121            module,
122            &sk_glwe_prepared,
123            sk_lwe,
124            source_xa,
125            source_xe,
126            sigma,
127            scratch,
128        );
129
130        let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(sk_glwe.n(), basek, k_tsk, rows_tsk, 1, sk_glwe.rank());
131        tsk.encrypt_sk(module, &sk_glwe, source_xa, source_xe, sigma, scratch);
132
133        Self {
134            brk,
135            atk: auto_keys,
136            tsk,
137        }
138    }
139}
140
141pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
142    pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
143    pub(crate) tsk: GGLWETensorKeyPrepared<Vec<u8>, B>,
144    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
145}
146
147impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B>>
148    for CircuitBootstrappingKey<D, BRA>
149where
150    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
151    BlindRotationKey<D, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
152    GGLWETensorKey<D>: PrepareAlloc<B, GGLWETensorKeyPrepared<Vec<u8>, B>>,
153    GGLWEAutomorphismKey<D>: PrepareAlloc<B, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
154{
155    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> {
156        let brk: BlindRotationKeyPrepared<Vec<u8>, BRA, B> = self.brk.prepare_alloc(module, scratch);
157        let tsk: GGLWETensorKeyPrepared<Vec<u8>, B> = self.tsk.prepare_alloc(module, scratch);
158        let mut atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>> = HashMap::new();
159        for (key, value) in &self.atk {
160            atk.insert(*key, value.prepare_alloc(module, scratch));
161        }
162        CircuitBootstrappingKeyPrepared { brk, tsk, atk }
163    }
164}