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(module, basek, k, rows, digits,1, rank));
 16        });
 17        Self { keys: keys }
 18    }
 19
 20    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, digits: 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, digits,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    pub fn digits(&self) -> usize {
 56        self.keys[0].digits()
 57    }
 58}
 59
 60impl TensorKey<Vec<u8>, FFT64> {
 61    pub fn generate_from_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
 62        GLWESecret::bytes_of(module, 1) + GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank)
 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: &GLWESecret<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.rank());
 79            assert_eq!(self.n(), module.n());
 80            assert_eq!(sk.n(), module.n());
 81        }
 82
 83        let rank: usize = self.rank();
 84
 85        let (mut sk_ij, scratch1) = scratch.tmp_sk(module, 1);
 86
 87        (0..rank).for_each(|i| {
 88            (i..rank).for_each(|j| {
 89                module.svp_apply(
 90                    &mut sk_ij.data_fourier,
 91                    0,
 92                    &sk.data_fourier,
 93                    i,
 94                    &sk.data_fourier,
 95                    j,
 96                );
 97                module.svp_idft(&mut sk_ij.data, 0, &sk_ij.data_fourier, 0, scratch1);
 98                self.at_mut(i, j)
 99                    .generate_from_sk(module, &sk_ij, sk, source_xa, source_xe, sigma, scratch1);
100            });
101        })
102    }
103
104    // Returns a mutable reference to GLWESwitchingKey_{s}(s[i] * s[j])
105    pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKey<DataSelf, FFT64> {
106        if i > j {
107            std::mem::swap(&mut i, &mut j);
108        };
109        let rank: usize = self.rank();
110        &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
111    }
112}
113
114impl<DataSelf: AsRef<[u8]>> TensorKey<DataSelf, FFT64> {
115    // Returns a reference to GLWESwitchingKey_{s}(s[i] * s[j])
116    pub fn at(&self, mut i: usize, mut j: usize) -> &GLWESwitchingKey<DataSelf, FFT64> {
117        if i > j {
118            std::mem::swap(&mut i, &mut j);
119        };
120        let rank: usize = self.rank();
121        &self.keys[i * rank + j - (i * (i + 1) / 2)]
122    }
123}