1use backend::{Backend, FFT64, MatZnxDft, Module};
 2
 3use crate::{GLWESwitchingKey, Infos};
 4
 5pub struct GLWETensorKey<C, B: Backend> {
 6    pub(crate) keys: Vec<GLWESwitchingKey<C, B>>,
 7}
 8
 9impl GLWETensorKey<Vec<u8>, FFT64> {
10    pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self {
11        let mut keys: Vec<GLWESwitchingKey<Vec<u8>, FFT64>> = Vec::new();
12        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
13        (0..pairs).for_each(|_| {
14            keys.push(GLWESwitchingKey::alloc(
15                module, basek, k, rows, digits, 1, rank,
16            ));
17        });
18        Self { keys: keys }
19    }
20
21    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize {
22        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
23        pairs * GLWESwitchingKey::<Vec<u8>, FFT64>::bytes_of(module, basek, k, rows, digits, 1, rank)
24    }
25}
26
27impl<T, B: Backend> Infos for GLWETensorKey<T, B> {
28    type Inner = MatZnxDft<T, B>;
29
30    fn inner(&self) -> &Self::Inner {
31        &self.keys[0].inner()
32    }
33
34    fn basek(&self) -> usize {
35        self.keys[0].basek()
36    }
37
38    fn k(&self) -> usize {
39        self.keys[0].k()
40    }
41}
42
43impl<T, B: Backend> GLWETensorKey<T, B> {
44    pub fn rank(&self) -> usize {
45        self.keys[0].rank()
46    }
47
48    pub fn rank_in(&self) -> usize {
49        self.keys[0].rank_in()
50    }
51
52    pub fn rank_out(&self) -> usize {
53        self.keys[0].rank_out()
54    }
55
56    pub fn digits(&self) -> usize {
57        self.keys[0].digits()
58    }
59}
60
61impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWETensorKey<DataSelf, FFT64> {
62    // Returns a mutable reference to GLWESwitchingKey_{s}(s[i] * s[j])
63    pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKey<DataSelf, FFT64> {
64        if i > j {
65            std::mem::swap(&mut i, &mut j);
66        };
67        let rank: usize = self.rank();
68        &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
69    }
70}
71
72impl<DataSelf: AsRef<[u8]>> GLWETensorKey<DataSelf, FFT64> {
73    // Returns a reference to GLWESwitchingKey_{s}(s[i] * s[j])
74    pub fn at(&self, mut i: usize, mut j: usize) -> &GLWESwitchingKey<DataSelf, FFT64> {
75        if i > j {
76            std::mem::swap(&mut i, &mut j);
77        };
78        let rank: usize = self.rank();
79        &self.keys[i * rank + j - (i * (i + 1) / 2)]
80    }
81}