1use backend::{Backend, FFT64, Module, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, ZnxInfos};
 2
 3use crate::{GLWESecret, dist::Distribution};
 4
 5pub struct FourierGLWESecret<T, B: Backend> {
 6    pub(crate) data: ScalarZnxDft<T, B>,
 7    pub(crate) dist: Distribution,
 8}
 9
10impl<B: Backend> FourierGLWESecret<Vec<u8>, B> {
11    pub fn alloc(module: &Module<B>, rank: usize) -> Self {
12        Self {
13            data: module.new_scalar_znx_dft(rank),
14            dist: Distribution::NONE,
15        }
16    }
17
18    pub fn bytes_of(module: &Module<B>, rank: usize) -> usize {
19        module.bytes_of_scalar_znx_dft(rank)
20    }
21}
22
23impl FourierGLWESecret<Vec<u8>, FFT64> {
24    pub fn from<D>(module: &Module<FFT64>, sk: &GLWESecret<D>) -> Self
25    where
26        D: AsRef<[u8]>,
27    {
28        let mut sk_dft: FourierGLWESecret<Vec<u8>, FFT64> = Self::alloc(module, sk.rank());
29        sk_dft.set(module, sk);
30        sk_dft
31    }
32}
33
34impl<DataSelf, B: Backend> FourierGLWESecret<DataSelf, B> {
35    pub fn n(&self) -> usize {
36        self.data.n()
37    }
38
39    pub fn log_n(&self) -> usize {
40        self.data.log_n()
41    }
42
43    pub fn rank(&self) -> usize {
44        self.data.cols()
45    }
46}
47
48impl<S: AsMut<[u8]> + AsRef<[u8]>> FourierGLWESecret<S, FFT64> {
49    pub(crate) fn set<D>(&mut self, module: &Module<FFT64>, sk: &GLWESecret<D>)
50    where
51        D: AsRef<[u8]>,
52    {
53        (0..self.rank()).for_each(|i| {
54            module.svp_prepare(&mut self.data, i, &sk.data, i);
55        });
56        self.dist = sk.dist
57    }
58}