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