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