1use backend::hal::{
 2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
 3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat},
 4};
 5
 6use crate::layouts::{
 7    GLWEToLWESwitchingKey, Infos,
 8    prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc},
 9};
10
11#[derive(PartialEq, Eq)]
12pub struct GLWEToLWESwitchingKeyPrepared<D: Data, B: Backend>(pub(crate) GGLWESwitchingKeyPrepared<D, B>);
13
14impl<D: Data, B: Backend> Infos for GLWEToLWESwitchingKeyPrepared<D, B> {
15    type Inner = VmpPMat<D, B>;
16
17    fn inner(&self) -> &Self::Inner {
18        &self.0.inner()
19    }
20
21    fn basek(&self) -> usize {
22        self.0.basek()
23    }
24
25    fn k(&self) -> usize {
26        self.0.k()
27    }
28}
29
30impl<D: Data, B: Backend> GLWEToLWESwitchingKeyPrepared<D, B> {
31    pub fn digits(&self) -> usize {
32        self.0.digits()
33    }
34
35    pub fn rank(&self) -> usize {
36        self.0.rank()
37    }
38
39    pub fn rank_in(&self) -> usize {
40        self.0.rank_in()
41    }
42
43    pub fn rank_out(&self) -> usize {
44        self.0.rank_out()
45    }
46}
47
48impl<B: Backend> GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> {
49    pub fn alloc(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, rank_in: usize) -> Self
50    where
51        Module<B>: VmpPMatAlloc<B>,
52    {
53        Self(GGLWESwitchingKeyPrepared::alloc(
54            module, n, basek, k, rows, 1, rank_in, 1,
55        ))
56    }
57
58    pub fn bytes_of(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize) -> usize
59    where
60        Module<B>: VmpPMatAllocBytes,
61    {
62        GGLWESwitchingKeyPrepared::<Vec<u8>, B>::bytes_of(module, n, basek, k, rows, digits, rank_in, 1)
63    }
64}
65
66impl<D: DataRef, B: Backend> PrepareAlloc<B, GLWEToLWESwitchingKeyPrepared<Vec<u8>, B>> for GLWEToLWESwitchingKey<D>
67where
68    Module<B>: VmpPrepare<B> + VmpPMatAlloc<B>,
69{
70    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> {
71        let mut ksk_prepared: GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> = GLWEToLWESwitchingKeyPrepared::alloc(
72            module,
73            self.0.n(),
74            self.0.basek(),
75            self.0.k(),
76            self.0.rows(),
77            self.0.rank_in(),
78        );
79        ksk_prepared.prepare(module, self, scratch);
80        ksk_prepared
81    }
82}
83
84impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, GLWEToLWESwitchingKey<DR>> for GLWEToLWESwitchingKeyPrepared<DM, B>
85where
86    Module<B>: VmpPrepare<B>,
87{
88    fn prepare(&mut self, module: &Module<B>, other: &GLWEToLWESwitchingKey<DR>, scratch: &mut Scratch<B>) {
89        self.0.prepare(module, &other.0, scratch);
90    }
91}