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