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