1use backend::{
  2    Backend, FFT64, MatZnxDft, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft, ScalarZnxDftAlloc,
  3    ScalarZnxDftOps, ScalarZnxDftToRef, Scratch, VecZnxDftOps, VecZnxDftToRef,
  4};
  5use sampling::source::Source;
  6
  7use crate::{
  8    elem::Infos,
  9    keys::{SecretKey, SecretKeyFourier},
 10    keyswitch_key::GLWESwitchingKey,
 11};
 12
 13pub struct TensorKey<C, B: Backend> {
 14    pub(crate) keys: Vec<GLWESwitchingKey<C, B>>,
 15}
 16
 17impl TensorKey<Vec<u8>, FFT64> {
 18    pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
 19        let mut keys: Vec<GLWESwitchingKey<Vec<u8>, FFT64>> = Vec::new();
 20        let pairs: usize = ((rank + 1) * rank) >> 1;
 21        (0..pairs).for_each(|_| {
 22            keys.push(GLWESwitchingKey::alloc(module, basek, k, rows, 1, rank));
 23        });
 24        Self { keys: keys }
 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
 58impl TensorKey<Vec<u8>, FFT64> {
 59    pub fn generate_from_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
 60        module.bytes_of_scalar_znx_dft(1) + GLWESwitchingKey::encrypt_sk_scratch_space(module, rank, size)
 61    }
 62}
 63
 64impl<DataSelf> TensorKey<DataSelf, FFT64>
 65where
 66    MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64>,
 67{
 68    pub fn generate_from_sk<DataSk>(
 69        &mut self,
 70        module: &Module<FFT64>,
 71        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
 72        source_xa: &mut Source,
 73        source_xe: &mut Source,
 74        sigma: f64,
 75        scratch: &mut Scratch,
 76    ) where
 77        ScalarZnxDft<DataSk, FFT64>: VecZnxDftToRef<FFT64> + ScalarZnxDftToRef<FFT64>,
 78    {
 79        #[cfg(debug_assertions)]
 80        {
 81            assert_eq!(self.rank(), sk_dft.rank());
 82            assert_eq!(self.n(), module.n());
 83            assert_eq!(sk_dft.n(), module.n());
 84        }
 85
 86        let rank: usize = self.rank();
 87
 88        (0..rank).for_each(|i| {
 89            (i..rank).for_each(|j| {
 90                let (mut sk_ij_dft, scratch1) = scratch.tmp_scalar_znx_dft(module, 1);
 91                module.svp_apply(&mut sk_ij_dft, 0, &sk_dft.data, i, &sk_dft.data, j);
 92                let sk_ij: ScalarZnx<&mut [u8]> = module
 93                    .vec_znx_idft_consume(sk_ij_dft.as_vec_znx_dft())
 94                    .to_vec_znx_small()
 95                    .to_scalar_znx();
 96                let sk_ij: SecretKey<&mut [u8]> = SecretKey {
 97                    data: sk_ij,
 98                    dist: sk_dft.dist,
 99                };
100
101                self.at_mut(i, j).encrypt_sk(
102                    module, &sk_ij, sk_dft, source_xa, source_xe, sigma, scratch1,
103                );
104            });
105        })
106    }
107
108    // Returns a mutable reference to GLWESwitchingKey_{s}(s[i] * s[j])
109    pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKey<DataSelf, FFT64> {
110        if i > j {
111            std::mem::swap(&mut i, &mut j);
112        };
113        let rank: usize = self.rank();
114        &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
115    }
116}
117
118impl<DataSelf> TensorKey<DataSelf, FFT64>
119where
120    MatZnxDft<DataSelf, FFT64>: MatZnxDftToRef<FFT64>,
121{
122    // Returns a reference to GLWESwitchingKey_{s}(s[i] * s[j])
123    pub fn at(&self, mut i: usize, mut j: usize) -> &GLWESwitchingKey<DataSelf, FFT64> {
124        if i > j {
125            std::mem::swap(&mut i, &mut j);
126        };
127        let rank: usize = self.rank();
128        &self.keys[i * rank + j - (i * (i + 1) / 2)]
129    }
130}