1use backend::hal::{
 2    api::{VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftFromVecZnx},
 3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VecZnxDft},
 4};
 5
 6use crate::{
 7    dist::Distribution,
 8    layouts::{GLWEPublicKey, Infos},
 9};
10
11#[derive(PartialEq, Eq)]
12pub struct GLWEPublicKeyExec<D: Data, B: Backend> {
13    pub(crate) data: VecZnxDft<D, B>,
14    pub(crate) basek: usize,
15    pub(crate) k: usize,
16    pub(crate) dist: Distribution,
17}
18
19impl<D: Data, B: Backend> Infos for GLWEPublicKeyExec<D, B> {
20    type Inner = VecZnxDft<D, B>;
21
22    fn inner(&self) -> &Self::Inner {
23        &self.data
24    }
25
26    fn basek(&self) -> usize {
27        self.basek
28    }
29
30    fn k(&self) -> usize {
31        self.k
32    }
33}
34
35impl<D: Data, B: Backend> GLWEPublicKeyExec<D, B> {
36    pub fn rank(&self) -> usize {
37        self.cols() - 1
38    }
39}
40
41impl<B: Backend> GLWEPublicKeyExec<Vec<u8>, B> {
42    pub fn alloc(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> Self
43    where
44        Module<B>: VecZnxDftAlloc<B>,
45    {
46        Self {
47            data: module.vec_znx_dft_alloc(n, rank + 1, k.div_ceil(basek)),
48            basek: basek,
49            k: k,
50            dist: Distribution::NONE,
51        }
52    }
53
54    pub fn bytes_of(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
55    where
56        Module<B>: VecZnxDftAllocBytes,
57    {
58        module.vec_znx_dft_alloc_bytes(n, rank + 1, k.div_ceil(basek))
59    }
60
61    pub fn from<DataOther>(module: &Module<B>, other: &GLWEPublicKey<DataOther>, scratch: &mut Scratch<B>) -> Self
62    where
63        DataOther: DataRef,
64        Module<B>: VecZnxDftAlloc<B> + VecZnxDftFromVecZnx<B>,
65    {
66        let mut pk_exec: GLWEPublicKeyExec<Vec<u8>, B> =
67            GLWEPublicKeyExec::alloc(module, other.n(), other.basek(), other.k(), other.rank());
68        pk_exec.prepare(module, other, scratch);
69        pk_exec
70    }
71}
72
73impl<D: DataMut, B: Backend> GLWEPublicKeyExec<D, B> {
74    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GLWEPublicKey<DataOther>, _scratch: &mut Scratch<B>)
75    where
76        DataOther: DataRef,
77        Module<B>: VecZnxDftFromVecZnx<B>,
78    {
79        #[cfg(debug_assertions)]
80        {
81            assert_eq!(self.n(), other.n());
82            assert_eq!(self.size(), other.size());
83        }
84
85        (0..self.cols()).for_each(|i| {
86            module.vec_znx_dft_from_vec_znx(1, 0, &mut self.data, i, &other.data, i);
87        });
88        self.k = other.k;
89        self.basek = other.basek;
90        self.dist = other.dist;
91    }
92}