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