1use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
  2
  3use crate::layouts::{
  4    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, LWEToGLWESwitchingKeyToRef, Rank, TorusPrecision,
  5    prepared::{
  6        GLWESwitchingKeyPrepare, GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedAlloc, GLWESwitchingKeyPreparedToMut,
  7        GLWESwitchingKeyPreparedToRef,
  8    },
  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) GLWESwitchingKeyPrepared<D, B>);
 14
 15impl<D: Data, B: Backend> LWEInfos for LWEToGLWESwitchingKeyPrepared<D, B> {
 16    fn base2k(&self) -> Base2K {
 17        self.0.base2k()
 18    }
 19
 20    fn k(&self) -> TorusPrecision {
 21        self.0.k()
 22    }
 23
 24    fn n(&self) -> Degree {
 25        self.0.n()
 26    }
 27
 28    fn size(&self) -> usize {
 29        self.0.size()
 30    }
 31}
 32
 33impl<D: Data, B: Backend> GLWEInfos for LWEToGLWESwitchingKeyPrepared<D, B> {
 34    fn rank(&self) -> Rank {
 35        self.rank_out()
 36    }
 37}
 38
 39impl<D: Data, B: Backend> GGLWEInfos for LWEToGLWESwitchingKeyPrepared<D, B> {
 40    fn dsize(&self) -> Dsize {
 41        self.0.dsize()
 42    }
 43
 44    fn rank_in(&self) -> Rank {
 45        self.0.rank_in()
 46    }
 47
 48    fn rank_out(&self) -> Rank {
 49        self.0.rank_out()
 50    }
 51
 52    fn dnum(&self) -> Dnum {
 53        self.0.dnum()
 54    }
 55}
 56
 57pub trait LWEToGLWESwitchingKeyPreparedAlloc<B: Backend>
 58where
 59    Self: GLWESwitchingKeyPreparedAlloc<B>,
 60{
 61    fn alloc_lwe_to_glwe_switching_key_prepared(
 62        &self,
 63        base2k: Base2K,
 64        k: TorusPrecision,
 65        rank_out: Rank,
 66        dnum: Dnum,
 67    ) -> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
 68        LWEToGLWESwitchingKeyPrepared(self.alloc_glwe_switching_key_prepared(base2k, k, Rank(1), rank_out, dnum, Dsize(1)))
 69    }
 70    fn alloc_lwe_to_glwe_switching_key_prepared_from_infos<A>(&self, infos: &A) -> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B>
 71    where
 72        A: GGLWEInfos,
 73    {
 74        debug_assert_eq!(
 75            infos.rank_in().0,
 76            1,
 77            "rank_in > 1 is not supported for LWEToGLWESwitchingKey"
 78        );
 79        debug_assert_eq!(
 80            infos.dsize().0,
 81            1,
 82            "dsize > 1 is not supported for LWEToGLWESwitchingKey"
 83        );
 84        self.alloc_lwe_to_glwe_switching_key_prepared(infos.base2k(), infos.k(), infos.rank_out(), infos.dnum())
 85    }
 86
 87    fn bytes_of_lwe_to_glwe_switching_key_prepared(
 88        &self,
 89        base2k: Base2K,
 90        k: TorusPrecision,
 91        rank_out: Rank,
 92        dnum: Dnum,
 93    ) -> usize {
 94        self.bytes_of_glwe_switching_key_prepared(base2k, k, Rank(1), rank_out, dnum, Dsize(1))
 95    }
 96
 97    fn bytes_of_lwe_to_glwe_switching_key_prepared_from_infos<A>(&self, infos: &A) -> usize
 98    where
 99        A: GGLWEInfos,
100    {
101        debug_assert_eq!(
102            infos.rank_in().0,
103            1,
104            "rank_in > 1 is not supported for LWEToGLWESwitchingKey"
105        );
106        debug_assert_eq!(
107            infos.dsize().0,
108            1,
109            "dsize > 1 is not supported for LWEToGLWESwitchingKey"
110        );
111        self.bytes_of_lwe_to_glwe_switching_key_prepared(infos.base2k(), infos.k(), infos.rank_out(), infos.dnum())
112    }
113}
114
115impl<B: Backend> LWEToGLWESwitchingKeyPreparedAlloc<B> for Module<B> where Self: GLWESwitchingKeyPreparedAlloc<B> {}
116
117impl<B: Backend> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B>
118where
119    Module<B>: LWEToGLWESwitchingKeyPreparedAlloc<B>,
120{
121    pub fn alloc_from_infos<A>(module: &Module<B>, infos: &A) -> Self
122    where
123        A: GGLWEInfos,
124    {
125        module.alloc_lwe_to_glwe_switching_key_prepared_from_infos(infos)
126    }
127
128    pub fn alloc(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> Self {
129        module.alloc_lwe_to_glwe_switching_key_prepared(base2k, k, rank_out, dnum)
130    }
131
132    pub fn bytes_of_from_infos<A>(module: &Module<B>, infos: &A) -> usize
133    where
134        A: GGLWEInfos,
135    {
136        module.bytes_of_lwe_to_glwe_switching_key_prepared_from_infos(infos)
137    }
138
139    pub fn bytes_of(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> usize {
140        module.bytes_of_lwe_to_glwe_switching_key_prepared(base2k, k, rank_out, dnum)
141    }
142}
143
144pub trait LWEToGLWESwitchingKeyPrepare<B: Backend>
145where
146    Self: GLWESwitchingKeyPrepare<B>,
147{
148    fn prepare_lwe_to_glwe_switching_key_tmp_bytes<A>(&self, infos: &A)
149    where
150        A: GGLWEInfos,
151    {
152        self.prepare_glwe_switching_key_tmp_bytes(infos);
153    }
154
155    fn prepare_lwe_to_glwe_switching_key<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
156    where
157        R: LWEToGLWESwitchingKeyPreparedToMut<B>,
158        O: LWEToGLWESwitchingKeyToRef,
159    {
160        self.prepare_glwe_switching(&mut res.to_mut().0, &other.to_ref().0, scratch);
161    }
162}
163
164impl<B: Backend> LWEToGLWESwitchingKeyPrepare<B> for Module<B> where Self: GLWESwitchingKeyPrepare<B> {}
165
166impl<B: Backend> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
167    pub fn prepare_tmp_bytes<A>(&self, module: &Module<B>, infos: &A)
168    where
169        A: GGLWEInfos,
170        Module<B>: LWEToGLWESwitchingKeyPrepare<B>,
171    {
172        module.prepare_lwe_to_glwe_switching_key_tmp_bytes(infos);
173    }
174}
175
176impl<D: DataMut, B: Backend> LWEToGLWESwitchingKeyPrepared<D, B> {
177    fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &mut Scratch<B>)
178    where
179        O: LWEToGLWESwitchingKeyToRef,
180        Module<B>: LWEToGLWESwitchingKeyPrepare<B>,
181    {
182        module.prepare_lwe_to_glwe_switching_key(self, other, scratch);
183    }
184}
185
186pub trait LWEToGLWESwitchingKeyPreparedToRef<B: Backend> {
187    fn to_ref(&self) -> LWEToGLWESwitchingKeyPrepared<&[u8], B>;
188}
189
190impl<D: DataRef, B: Backend> LWEToGLWESwitchingKeyPreparedToRef<B> for LWEToGLWESwitchingKeyPrepared<D, B>
191where
192    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToRef<B>,
193{
194    fn to_ref(&self) -> LWEToGLWESwitchingKeyPrepared<&[u8], B> {
195        LWEToGLWESwitchingKeyPrepared(self.0.to_ref())
196    }
197}
198
199pub trait LWEToGLWESwitchingKeyPreparedToMut<B: Backend> {
200    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyPrepared<&mut [u8], B>;
201}
202
203impl<D: DataMut, B: Backend> LWEToGLWESwitchingKeyPreparedToMut<B> for LWEToGLWESwitchingKeyPrepared<D, B>
204where
205    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToMut<B>,
206{
207    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyPrepared<&mut [u8], B> {
208        LWEToGLWESwitchingKeyPrepared(self.0.to_mut())
209    }
210}