1use backend::{Backend, FFT64, MatZnxDft, Module, ScalarZnxDftOps, Scratch};
  2use sampling::source::Source;
  3
  4use crate::{GLWESecret, GLWESwitchingKey, Infos, ScratchCore};
  5
  6pub struct TensorKey<C, B: Backend> {
  7    pub(crate) keys: Vec<GLWESwitchingKey<C, B>>,
  8}
  9
 10impl TensorKey<Vec<u8>, FFT64> {
 11    pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
 12        let mut keys: Vec<GLWESwitchingKey<Vec<u8>, FFT64>> = Vec::new();
 13        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
 14        (0..pairs).for_each(|_| {
 15            keys.push(GLWESwitchingKey::alloc(module, basek, k, rows, 1, rank));
 16        });
 17        Self { keys: keys }
 18    }
 19
 20    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> usize {
 21        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
 22        pairs * GLWESwitchingKey::<Vec<u8>, FFT64>::bytes_of(module, basek, k, rows, 1, rank)
 23    }
 24}
 25
 26impl<T, B: Backend> Infos for TensorKey<T, B> {
 27    type Inner = MatZnxDft<T, B>;
 28
 29    fn inner(&self) -> &Self::Inner {
 30        &self.keys[0].inner()
 31    }
 32
 33    fn basek(&self) -> usize {
 34        self.keys[0].basek()
 35    }
 36
 37    fn k(&self) -> usize {
 38        self.keys[0].k()
 39    }
 40}
 41
 42impl<T, B: Backend> TensorKey<T, B> {
 43    pub fn rank(&self) -> usize {
 44        self.keys[0].rank()
 45    }
 46
 47    pub fn rank_in(&self) -> usize {
 48        self.keys[0].rank_in()
 49    }
 50
 51    pub fn rank_out(&self) -> usize {
 52        self.keys[0].rank_out()
 53    }
 54}
 55
 56impl TensorKey<Vec<u8>, FFT64> {
 57    pub fn generate_from_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
 58        GLWESecret::bytes_of(module, 1) + GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank)
 59    }
 60}
 61
 62impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> TensorKey<DataSelf, FFT64> {
 63    pub fn generate_from_sk<DataSk: AsRef<[u8]>>(
 64        &mut self,
 65        module: &Module<FFT64>,
 66        sk: &GLWESecret<DataSk, FFT64>,
 67        source_xa: &mut Source,
 68        source_xe: &mut Source,
 69        sigma: f64,
 70        scratch: &mut Scratch,
 71    ) {
 72        #[cfg(debug_assertions)]
 73        {
 74            assert_eq!(self.rank(), sk.rank());
 75            assert_eq!(self.n(), module.n());
 76            assert_eq!(sk.n(), module.n());
 77        }
 78
 79        let rank: usize = self.rank();
 80
 81        let (mut sk_ij, scratch1) = scratch.tmp_sk(module, 1);
 82
 83        (0..rank).for_each(|i| {
 84            (i..rank).for_each(|j| {
 85                module.svp_apply(
 86                    &mut sk_ij.data_fourier,
 87                    0,
 88                    &sk.data_fourier,
 89                    i,
 90                    &sk.data_fourier,
 91                    j,
 92                );
 93                module.svp_idft(&mut sk_ij.data, 0, &sk_ij.data_fourier, 0, scratch1);
 94                self.at_mut(i, j)
 95                    .generate_from_sk(module, &sk_ij, sk, source_xa, source_xe, sigma, scratch1);
 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}