1use backend::hal::{
 2    api::{SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, ZnxInfos},
 3    layouts::{Backend, Data, DataMut, DataRef, Module, SvpPPol},
 4};
 5
 6use crate::{dist::Distribution, layouts::GLWESecret, trait_families::GLWESecretExecModuleFamily};
 7
 8pub struct GLWESecretExec<D: Data, B: Backend> {
 9    pub(crate) data: SvpPPol<D, B>,
10    pub(crate) dist: Distribution,
11}
12
13impl<B: Backend> GLWESecretExec<Vec<u8>, B> {
14    pub fn alloc(module: &Module<B>, n: usize, rank: usize) -> Self
15    where
16        Module<B>: GLWESecretExecModuleFamily<B>,
17    {
18        Self {
19            data: module.svp_ppol_alloc(n, rank),
20            dist: Distribution::NONE,
21        }
22    }
23
24    pub fn bytes_of(module: &Module<B>, n: usize, rank: usize) -> usize
25    where
26        Module<B>: GLWESecretExecModuleFamily<B>,
27    {
28        module.svp_ppol_alloc_bytes(n, rank)
29    }
30}
31
32impl<B: Backend> GLWESecretExec<Vec<u8>, B> {
33    pub fn from<D>(module: &Module<B>, sk: &GLWESecret<D>) -> Self
34    where
35        D: DataRef,
36        Module<B>: GLWESecretExecModuleFamily<B>,
37    {
38        let mut sk_dft: GLWESecretExec<Vec<u8>, B> = Self::alloc(module, sk.n(), sk.rank());
39        sk_dft.prepare(module, sk);
40        sk_dft
41    }
42}
43
44impl<D: Data, B: Backend> GLWESecretExec<D, B> {
45    pub fn n(&self) -> usize {
46        self.data.n()
47    }
48
49    pub fn log_n(&self) -> usize {
50        self.data.log_n()
51    }
52
53    pub fn rank(&self) -> usize {
54        self.data.cols()
55    }
56}
57
58impl<D: DataMut, B: Backend> GLWESecretExec<D, B> {
59    pub(crate) fn prepare<O>(&mut self, module: &Module<B>, sk: &GLWESecret<O>)
60    where
61        O: DataRef,
62        Module<B>: GLWESecretExecModuleFamily<B>,
63    {
64        (0..self.rank()).for_each(|i| {
65            module.svp_prepare(&mut self.data, i, &sk.data, i);
66        });
67        self.dist = sk.dist
68    }
69}