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