1use poulpy_core::{
  2    layouts::{
  3        GGLWEInfos, GGLWEToGGSWKeyLayout, GGLWEToGGSWKeyPrepared, GGLWEToGGSWKeyPreparedFactory, GGSWInfos,
  4        GLWEAutomorphismKeyHelper, GLWEAutomorphismKeyLayout, GLWEAutomorphismKeyPreparedFactory, GLWEInfos,
  5        GLWETensorKeyPreparedFactory, LWEInfos, prepared::GLWEAutomorphismKeyPrepared,
  6    },
  7    trace_galois_elements,
  8};
  9use std::collections::HashMap;
 10
 11use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
 12
 13use crate::tfhe::{
 14    blind_rotation::{
 15        BlindRotationAlgo, BlindRotationKeyInfos, BlindRotationKeyLayout, BlindRotationKeyPrepared,
 16        BlindRotationKeyPreparedFactory,
 17    },
 18    circuit_bootstrapping::{CircuitBootstrappingKey, CircuitBootstrappingKeyInfos},
 19};
 20
 21impl<BRA: BlindRotationAlgo, BE: Backend> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> {
 22    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE>
 23    where
 24        A: CircuitBootstrappingKeyInfos,
 25        M: CircuitBootstrappingKeyPreparedFactory<BRA, BE>,
 26    {
 27        module.circuit_bootstrapping_key_prepared_alloc_from_infos(infos)
 28    }
 29}
 30
 31impl<D: DataMut, BRA: BlindRotationAlgo, BE: Backend> CircuitBootstrappingKeyPrepared<D, BRA, BE> {
 32    pub fn prepare<DR, M>(&mut self, module: &M, other: &CircuitBootstrappingKey<DR, BRA>, scratch: &mut Scratch<BE>)
 33    where
 34        DR: DataRef,
 35        M: CircuitBootstrappingKeyPreparedFactory<BRA, BE>,
 36    {
 37        module.circuit_bootstrapping_key_prepare(self, other, scratch);
 38    }
 39}
 40
 41impl<BE: Backend, BRA: BlindRotationAlgo> CircuitBootstrappingKeyPreparedFactory<BRA, BE> for Module<BE> where
 42    Self: Sized
 43        + BlindRotationKeyPreparedFactory<BRA, BE>
 44        + GLWETensorKeyPreparedFactory<BE>
 45        + GLWEAutomorphismKeyPreparedFactory<BE>
 46{
 47}
 48
 49pub trait CircuitBootstrappingKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend>
 50where
 51    Self: Sized
 52        + BlindRotationKeyPreparedFactory<BRA, BE>
 53        + GGLWEToGGSWKeyPreparedFactory<BE>
 54        + GLWEAutomorphismKeyPreparedFactory<BE>,
 55{
 56    fn circuit_bootstrapping_key_prepared_alloc_from_infos<A>(
 57        &self,
 58        infos: &A,
 59    ) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE>
 60    where
 61        A: CircuitBootstrappingKeyInfos,
 62    {
 63        let atk_infos: &GLWEAutomorphismKeyLayout = &infos.atk_infos();
 64        let gal_els: Vec<i64> = trace_galois_elements(atk_infos.log_n(), 2 * atk_infos.n().as_usize() as i64);
 65
 66        CircuitBootstrappingKeyPrepared {
 67            brk: BlindRotationKeyPrepared::alloc(self, &infos.brk_infos()),
 68            tsk: GGLWEToGGSWKeyPrepared::alloc_from_infos(self, &infos.tsk_infos()),
 69            atk: gal_els
 70                .iter()
 71                .map(|&gal_el| {
 72                    let key = GLWEAutomorphismKeyPrepared::alloc_from_infos(self, atk_infos);
 73                    (gal_el, key)
 74                })
 75                .collect(),
 76        }
 77    }
 78
 79    fn circuit_bootstrapping_key_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
 80    where
 81        A: CircuitBootstrappingKeyInfos,
 82    {
 83        self.blind_rotation_key_prepare_tmp_bytes(&infos.brk_infos())
 84            .max(self.prepare_gglwe_to_ggsw_key_tmp_bytes(&infos.tsk_infos()))
 85            .max(self.prepare_glwe_automorphism_key_tmp_bytes(&infos.atk_infos()))
 86    }
 87
 88    fn circuit_bootstrapping_key_prepare<DM, DR>(
 89        &self,
 90        res: &mut CircuitBootstrappingKeyPrepared<DM, BRA, BE>,
 91        other: &CircuitBootstrappingKey<DR, BRA>,
 92        scratch: &mut Scratch<BE>,
 93    ) where
 94        DM: DataMut,
 95        DR: DataRef,
 96    {
 97        res.brk.prepare(self, &other.brk, scratch);
 98        res.tsk.prepare(self, &other.tsk, scratch);
 99
100        for (k, a) in res.atk.iter_mut() {
101            a.prepare(self, other.atk.get(k).unwrap(), scratch);
102        }
103    }
104}
105
106pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
107    pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
108    pub(crate) tsk: GGLWEToGGSWKeyPrepared<D, B>,
109    pub(crate) atk: HashMap<i64, GLWEAutomorphismKeyPrepared<D, B>>,
110}
111
112impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> GLWEAutomorphismKeyHelper<GLWEAutomorphismKeyPrepared<D, BE>, BE>
113    for CircuitBootstrappingKeyPrepared<D, BRA, BE>
114{
115    fn get_automorphism_key(&self, k: i64) -> Option<&GLWEAutomorphismKeyPrepared<D, BE>> {
116        self.atk.get_automorphism_key(k)
117    }
118
119    fn automorphism_key_infos(&self) -> poulpy_core::layouts::GGLWELayout {
120        self.atk.automorphism_key_infos()
121    }
122}
123
124impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyPrepared<D, BRA, B> {
125    fn atk_infos(&self) -> GLWEAutomorphismKeyLayout {
126        let (_, atk) = self.atk.iter().next().expect("atk is empty");
127        GLWEAutomorphismKeyLayout {
128            n: atk.n(),
129            base2k: atk.base2k(),
130            k: atk.k(),
131            dnum: atk.dnum(),
132            dsize: atk.dsize(),
133            rank: atk.rank(),
134        }
135    }
136
137    fn brk_infos(&self) -> BlindRotationKeyLayout {
138        BlindRotationKeyLayout {
139            n_glwe: self.brk.n_glwe(),
140            n_lwe: self.brk.n_lwe(),
141            base2k: self.brk.base2k(),
142            k: self.brk.k(),
143            dnum: self.brk.dnum(),
144            rank: self.brk.rank(),
145        }
146    }
147
148    fn tsk_infos(&self) -> GGLWEToGGSWKeyLayout {
149        GGLWEToGGSWKeyLayout {
150            n: self.tsk.n(),
151            base2k: self.tsk.base2k(),
152            k: self.tsk.k(),
153            dnum: self.tsk.dnum(),
154            dsize: self.tsk.dsize(),
155            rank: self.tsk.rank(),
156        }
157    }
158}