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::{GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedFactory},
  7};
  8
  9/// A special [GLWESwitchingKey] required to for the conversion from [LWE] to [GLWE].
 10#[derive(PartialEq, Eq)]
 11pub struct LWEToGLWEKeyPrepared<D: Data, B: Backend>(pub(crate) GLWESwitchingKeyPrepared<D, B>);
 12
 13impl<D: Data, B: Backend> LWEInfos for LWEToGLWEKeyPrepared<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 LWEToGLWEKeyPrepared<D, B> {
 32    fn rank(&self) -> Rank {
 33        self.rank_out()
 34    }
 35}
 36
 37impl<D: Data, B: Backend> GGLWEInfos for LWEToGLWEKeyPrepared<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 LWEToGLWEKeyPreparedFactory<B: Backend>
 56where
 57    Self: GLWESwitchingKeyPreparedFactory<B>,
 58{
 59    fn alloc_lwe_to_glwe_key_prepared(
 60        &self,
 61        base2k: Base2K,
 62        k: TorusPrecision,
 63        rank_out: Rank,
 64        dnum: Dnum,
 65    ) -> LWEToGLWEKeyPrepared<Vec<u8>, B> {
 66        LWEToGLWEKeyPrepared(self.alloc_glwe_switching_key_prepared(base2k, k, Rank(1), rank_out, dnum, Dsize(1)))
 67    }
 68    fn alloc_lwe_to_glwe_key_prepared_from_infos<A>(&self, infos: &A) -> LWEToGLWEKeyPrepared<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 LWEToGLWEKey"
 76        );
 77        debug_assert_eq!(
 78            infos.dsize().0,
 79            1,
 80            "dsize > 1 is not supported for LWEToGLWEKey"
 81        );
 82        self.alloc_lwe_to_glwe_key_prepared(infos.base2k(), infos.k(), infos.rank_out(), infos.dnum())
 83    }
 84
 85    fn bytes_of_lwe_to_glwe_key_prepared(&self, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> usize {
 86        self.bytes_of_glwe_key_prepared(base2k, k, Rank(1), rank_out, dnum, Dsize(1))
 87    }
 88
 89    fn bytes_of_lwe_to_glwe_key_prepared_from_infos<A>(&self, infos: &A) -> usize
 90    where
 91        A: GGLWEInfos,
 92    {
 93        debug_assert_eq!(
 94            infos.rank_in().0,
 95            1,
 96            "rank_in > 1 is not supported for LWEToGLWEKey"
 97        );
 98        debug_assert_eq!(
 99            infos.dsize().0,
100            1,
101            "dsize > 1 is not supported for LWEToGLWEKey"
102        );
103        self.bytes_of_lwe_to_glwe_key_prepared(infos.base2k(), infos.k(), infos.rank_out(), infos.dnum())
104    }
105
106    fn prepare_lwe_to_glwe_key_tmp_bytes<A>(&self, infos: &A)
107    where
108        A: GGLWEInfos,
109    {
110        self.prepare_glwe_switching_key_tmp_bytes(infos);
111    }
112
113    fn prepare_lwe_to_glwe_key<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
114    where
115        R: GGLWEPreparedToMut<B> + GLWESwitchingKeyDegreesMut,
116        O: GGLWEToRef + GLWESwitchingKeyDegrees,
117    {
118        self.prepare_glwe_switching(res, other, scratch);
119    }
120}
121
122impl<B: Backend> LWEToGLWEKeyPreparedFactory<B> for Module<B> where Self: GLWESwitchingKeyPreparedFactory<B> {}
123
124impl<B: Backend> LWEToGLWEKeyPrepared<Vec<u8>, B> {
125    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
126    where
127        A: GGLWEInfos,
128        M: LWEToGLWEKeyPreparedFactory<B>,
129    {
130        module.alloc_lwe_to_glwe_key_prepared_from_infos(infos)
131    }
132
133    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> Self
134    where
135        M: LWEToGLWEKeyPreparedFactory<B>,
136    {
137        module.alloc_lwe_to_glwe_key_prepared(base2k, k, rank_out, dnum)
138    }
139
140    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
141    where
142        A: GGLWEInfos,
143        M: LWEToGLWEKeyPreparedFactory<B>,
144    {
145        module.bytes_of_lwe_to_glwe_key_prepared_from_infos(infos)
146    }
147
148    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> usize
149    where
150        M: LWEToGLWEKeyPreparedFactory<B>,
151    {
152        module.bytes_of_lwe_to_glwe_key_prepared(base2k, k, rank_out, dnum)
153    }
154}
155
156impl<B: Backend> LWEToGLWEKeyPrepared<Vec<u8>, B> {
157    pub fn prepare_tmp_bytes<A, M>(&self, module: &M, infos: &A)
158    where
159        A: GGLWEInfos,
160        M: LWEToGLWEKeyPreparedFactory<B>,
161    {
162        module.prepare_lwe_to_glwe_key_tmp_bytes(infos);
163    }
164}
165
166impl<D: DataMut, B: Backend> LWEToGLWEKeyPrepared<D, B> {
167    pub fn prepare<O, M>(&mut self, module: &M, other: &O, scratch: &mut Scratch<B>)
168    where
169        O: GGLWEToRef + GLWESwitchingKeyDegrees,
170        M: LWEToGLWEKeyPreparedFactory<B>,
171    {
172        module.prepare_lwe_to_glwe_key(self, other, scratch);
173    }
174}
175
176impl<D: DataRef, B: Backend> GGLWEPreparedToRef<B> for LWEToGLWEKeyPrepared<D, B>
177where
178    GLWESwitchingKeyPrepared<D, B>: GGLWEPreparedToRef<B>,
179{
180    fn to_ref(&self) -> GGLWEPrepared<&[u8], B> {
181        self.0.to_ref()
182    }
183}
184
185impl<D: DataMut, B: Backend> GGLWEPreparedToMut<B> for LWEToGLWEKeyPrepared<D, B>
186where
187    GLWESwitchingKeyPrepared<D, B>: GGLWEPreparedToMut<B>,
188{
189    fn to_mut(&mut self) -> GGLWEPrepared<&mut [u8], B> {
190        self.0.to_mut()
191    }
192}
193
194impl<D: DataMut, B: Backend> GLWESwitchingKeyDegreesMut for LWEToGLWEKeyPrepared<D, B> {
195    fn input_degree(&mut self) -> &mut Degree {
196        &mut self.0.input_degree
197    }
198
199    fn output_degree(&mut self) -> &mut Degree {
200        &mut self.0.output_degree
201    }
202}