1use backend::{Backend, FFT64, MatZnxDft, Module, ScalarZnx, ScalarZnxDftAlloc, ScalarZnxDftOps, Scratch, VecZnxDftOps};
  2use sampling::source::Source;
  3
  4use crate::{
  5    elem::Infos,
  6    keys::{SecretKey, SecretKeyFourier},
  7    keyswitch_key::GLWESwitchingKey,
  8};
  9
 10pub struct TensorKey<C, B: Backend> {
 11    pub(crate) keys: Vec<GLWESwitchingKey<C, B>>,
 12}
 13
 14impl TensorKey<Vec<u8>, FFT64> {
 15    pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
 16        let mut keys: Vec<GLWESwitchingKey<Vec<u8>, FFT64>> = Vec::new();
 17        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
 18        (0..pairs).for_each(|_| {
 19            keys.push(GLWESwitchingKey::alloc(module, basek, k, rows, 1, rank));
 20        });
 21        Self { keys: keys }
 22    }
 23
 24    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> usize {
 25        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
 26        pairs * GLWESwitchingKey::<Vec<u8>, FFT64>::bytes_of(module, basek, k, rows, 1, rank)
 27    }
 28}
 29
 30impl<T, B: Backend> Infos for TensorKey<T, B> {
 31    type Inner = MatZnxDft<T, B>;
 32
 33    fn inner(&self) -> &Self::Inner {
 34        &self.keys[0].inner()
 35    }
 36
 37    fn basek(&self) -> usize {
 38        self.keys[0].basek()
 39    }
 40
 41    fn k(&self) -> usize {
 42        self.keys[0].k()
 43    }
 44}
 45
 46impl<T, B: Backend> TensorKey<T, B> {
 47    pub fn rank(&self) -> usize {
 48        self.keys[0].rank()
 49    }
 50
 51    pub fn rank_in(&self) -> usize {
 52        self.keys[0].rank_in()
 53    }
 54
 55    pub fn rank_out(&self) -> usize {
 56        self.keys[0].rank_out()
 57    }
 58}
 59
 60impl TensorKey<Vec<u8>, FFT64> {
 61    pub fn generate_from_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
 62        module.bytes_of_scalar_znx_dft(1) + GLWESwitchingKey::encrypt_sk_scratch_space(module, rank, size)
 63    }
 64}
 65
 66impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> TensorKey<DataSelf, FFT64> {
 67    pub fn generate_from_sk<DataSk: AsRef<[u8]>>(
 68        &mut self,
 69        module: &Module<FFT64>,
 70        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
 71        source_xa: &mut Source,
 72        source_xe: &mut Source,
 73        sigma: f64,
 74        scratch: &mut Scratch,
 75    ) {
 76        #[cfg(debug_assertions)]
 77        {
 78            assert_eq!(self.rank(), sk_dft.rank());
 79            assert_eq!(self.n(), module.n());
 80            assert_eq!(sk_dft.n(), module.n());
 81        }
 82
 83        let rank: usize = self.rank();
 84
 85        (0..rank).for_each(|i| {
 86            (i..rank).for_each(|j| {
 87                let (mut sk_ij_dft, scratch1) = scratch.tmp_scalar_znx_dft(module, 1);
 88                module.svp_apply(&mut sk_ij_dft, 0, &sk_dft.data, i, &sk_dft.data, j);
 89                let sk_ij: ScalarZnx<&mut [u8]> = module
 90                    .vec_znx_idft_consume(sk_ij_dft.as_vec_znx_dft())
 91                    .to_vec_znx_small()
 92                    .to_scalar_znx();
 93                let sk_ij: SecretKey<&mut [u8]> = SecretKey {
 94                    data: sk_ij,
 95                    dist: sk_dft.dist,
 96                };
 97
 98                self.at_mut(i, j).generate_from_sk(
 99                    module, &sk_ij, sk_dft, source_xa, source_xe, sigma, scratch1,
100                );
101            });
102        })
103    }
104
105    // Returns a mutable reference to GLWESwitchingKey_{s}(s[i] * s[j])
106    pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKey<DataSelf, FFT64> {
107        if i > j {
108            std::mem::swap(&mut i, &mut j);
109        };
110        let rank: usize = self.rank();
111        &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
112    }
113}
114
115impl<DataSelf: AsRef<[u8]>> TensorKey<DataSelf, FFT64> {
116    // Returns a reference to GLWESwitchingKey_{s}(s[i] * s[j])
117    pub fn at(&self, mut i: usize, mut j: usize) -> &GLWESwitchingKey<DataSelf, FFT64> {
118        if i > j {
119            std::mem::swap(&mut i, &mut j);
120        };
121        let rank: usize = self.rank();
122        &self.keys[i * rank + j - (i * (i + 1) / 2)]
123    }
124}