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 [LWE] to [GLWE].
 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> {
118    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
119    where
120        A: GGLWEInfos,
121        M: LWEToGLWESwitchingKeyPreparedAlloc<B>,
122    {
123        module.alloc_lwe_to_glwe_switching_key_prepared_from_infos(infos)
124    }
125
126    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> Self
127    where
128        M: LWEToGLWESwitchingKeyPreparedAlloc<B>,
129    {
130        module.alloc_lwe_to_glwe_switching_key_prepared(base2k, k, rank_out, dnum)
131    }
132
133    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
134    where
135        A: GGLWEInfos,
136        M: LWEToGLWESwitchingKeyPreparedAlloc<B>,
137    {
138        module.bytes_of_lwe_to_glwe_switching_key_prepared_from_infos(infos)
139    }
140
141    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> usize
142    where
143        M: LWEToGLWESwitchingKeyPreparedAlloc<B>,
144    {
145        module.bytes_of_lwe_to_glwe_switching_key_prepared(base2k, k, rank_out, dnum)
146    }
147}
148
149pub trait LWEToGLWESwitchingKeyPrepare<B: Backend>
150where
151    Self: GLWESwitchingKeyPrepare<B>,
152{
153    fn prepare_lwe_to_glwe_switching_key_tmp_bytes<A>(&self, infos: &A)
154    where
155        A: GGLWEInfos,
156    {
157        self.prepare_glwe_switching_key_tmp_bytes(infos);
158    }
159
160    fn prepare_lwe_to_glwe_switching_key<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
161    where
162        R: LWEToGLWESwitchingKeyPreparedToMut<B>,
163        O: LWEToGLWESwitchingKeyToRef,
164    {
165        self.prepare_glwe_switching(&mut res.to_mut().0, &other.to_ref().0, scratch);
166    }
167}
168
169impl<B: Backend> LWEToGLWESwitchingKeyPrepare<B> for Module<B> where Self: GLWESwitchingKeyPrepare<B> {}
170
171impl<B: Backend> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
172    pub fn prepare_tmp_bytes<A, M>(&self, module: &M, infos: &A)
173    where
174        A: GGLWEInfos,
175        M: LWEToGLWESwitchingKeyPrepare<B>,
176    {
177        module.prepare_lwe_to_glwe_switching_key_tmp_bytes(infos);
178    }
179}
180
181impl<D: DataMut, B: Backend> LWEToGLWESwitchingKeyPrepared<D, B> {
182    fn prepare<O, M>(&mut self, module: &M, other: &O, scratch: &mut Scratch<B>)
183    where
184        O: LWEToGLWESwitchingKeyToRef,
185        M: LWEToGLWESwitchingKeyPrepare<B>,
186    {
187        module.prepare_lwe_to_glwe_switching_key(self, other, scratch);
188    }
189}
190
191pub trait LWEToGLWESwitchingKeyPreparedToRef<B: Backend> {
192    fn to_ref(&self) -> LWEToGLWESwitchingKeyPrepared<&[u8], B>;
193}
194
195impl<D: DataRef, B: Backend> LWEToGLWESwitchingKeyPreparedToRef<B> for LWEToGLWESwitchingKeyPrepared<D, B>
196where
197    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToRef<B>,
198{
199    fn to_ref(&self) -> LWEToGLWESwitchingKeyPrepared<&[u8], B> {
200        LWEToGLWESwitchingKeyPrepared(self.0.to_ref())
201    }
202}
203
204pub trait LWEToGLWESwitchingKeyPreparedToMut<B: Backend> {
205    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyPrepared<&mut [u8], B>;
206}
207
208impl<D: DataMut, B: Backend> LWEToGLWESwitchingKeyPreparedToMut<B> for LWEToGLWESwitchingKeyPrepared<D, B>
209where
210    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToMut<B>,
211{
212    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyPrepared<&mut [u8], B> {
213        LWEToGLWESwitchingKeyPrepared(self.0.to_mut())
214    }
215}