1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::layouts::{
  7    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, LWESwitchingKeyToRef, Rank, TorusPrecision,
  8    prepared::{GLWESwitchingKeyPrepare, GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToMut, GLWESwitchingKeyPreparedToRef},
  9};
 10
 11#[derive(PartialEq, Eq)]
 12pub struct LWESwitchingKeyPrepared<D: Data, B: Backend>(pub(crate) GLWESwitchingKeyPrepared<D, B>);
 13
 14impl<D: Data, B: Backend> LWEInfos for LWESwitchingKeyPrepared<D, B> {
 15    fn base2k(&self) -> Base2K {
 16        self.0.base2k()
 17    }
 18
 19    fn k(&self) -> TorusPrecision {
 20        self.0.k()
 21    }
 22
 23    fn n(&self) -> Degree {
 24        self.0.n()
 25    }
 26
 27    fn size(&self) -> usize {
 28        self.0.size()
 29    }
 30}
 31impl<D: Data, B: Backend> GLWEInfos for LWESwitchingKeyPrepared<D, B> {
 32    fn rank(&self) -> Rank {
 33        self.rank_out()
 34    }
 35}
 36
 37impl<D: Data, B: Backend> GGLWEInfos for LWESwitchingKeyPrepared<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
 55impl<B: Backend> LWESwitchingKeyPrepared<Vec<u8>, B> {
 56    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
 57    where
 58        A: GGLWEInfos,
 59        Module<B>: VmpPMatAlloc<B>,
 60    {
 61        debug_assert_eq!(
 62            infos.dsize().0,
 63            1,
 64            "dsize > 1 is not supported for LWESwitchingKey"
 65        );
 66        debug_assert_eq!(
 67            infos.rank_in().0,
 68            1,
 69            "rank_in > 1 is not supported for LWESwitchingKey"
 70        );
 71        debug_assert_eq!(
 72            infos.rank_out().0,
 73            1,
 74            "rank_out > 1 is not supported for LWESwitchingKey"
 75        );
 76        Self(GLWESwitchingKeyPrepared::alloc_from_infos(module, infos))
 77    }
 78
 79    pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum) -> Self
 80    where
 81        Module<B>: VmpPMatAlloc<B>,
 82    {
 83        Self(GLWESwitchingKeyPrepared::alloc(
 84            module,
 85            base2k,
 86            k,
 87            Rank(1),
 88            Rank(1),
 89            dnum,
 90            Dsize(1),
 91        ))
 92    }
 93
 94    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
 95    where
 96        A: GGLWEInfos,
 97        Module<B>: VmpPMatAllocBytes,
 98    {
 99        debug_assert_eq!(
100            infos.dsize().0,
101            1,
102            "dsize > 1 is not supported for LWESwitchingKey"
103        );
104        debug_assert_eq!(
105            infos.rank_in().0,
106            1,
107            "rank_in > 1 is not supported for LWESwitchingKey"
108        );
109        debug_assert_eq!(
110            infos.rank_out().0,
111            1,
112            "rank_out > 1 is not supported for LWESwitchingKey"
113        );
114        GLWESwitchingKeyPrepared::alloc_bytes_from_infos(module, infos)
115    }
116
117    pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum) -> usize
118    where
119        Module<B>: VmpPMatAllocBytes,
120    {
121        GLWESwitchingKeyPrepared::alloc_bytes(module, base2k, k, Rank(1), Rank(1), dnum, Dsize(1))
122    }
123}
124
125pub trait LWESwitchingKeyPrepareTmpBytes {
126    fn lwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
127    where
128        A: GGLWEInfos;
129}
130
131impl<B: Backend> LWESwitchingKeyPrepareTmpBytes for Module<B>
132where
133    Module<B>: LWESwitchingKeyPrepareTmpBytes,
134{
135    fn lwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
136    where
137        A: GGLWEInfos,
138    {
139        self.lwe_switching_key_prepare_tmp_bytes(infos);
140    }
141}
142
143impl<B: Backend> LWESwitchingKeyPrepared<Vec<u8>, B> {
144    pub fn prepare_tmp_bytes<A>(&self, module: &Module<B>, infos: &A)
145    where
146        A: GLWEInfos,
147        Module<B>: LWESwitchingKeyPrepareTmpBytes,
148    {
149        module.glwe_secret_prepare_tmp_bytes(infos);
150    }
151}
152
153pub trait LWESwitchingKeyPrepare<B: Backend> {
154    fn lwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
155    where
156        R: LWESwitchingKeyPreparedToMut<B>,
157        O: LWESwitchingKeyToRef;
158}
159
160impl<B: Backend> LWESwitchingKeyPrepare<B> for Module<B>
161where
162    Module<B>: GLWESwitchingKeyPrepare<B>,
163{
164    fn lwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
165    where
166        R: LWESwitchingKeyPreparedToMut<B>,
167        O: LWESwitchingKeyToRef,
168    {
169        self.glwe_switching_prepare(&mut res.to_mut().0, other, scratch);
170    }
171}
172
173impl<D: DataMut, B: Backend> LWESwitchingKeyPrepared<D, B> {
174    fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &Scratch<B>)
175    where
176        O: LWESwitchingKeyToRef,
177        Module<B>: LWESwitchingKeyPrepare<B>,
178    {
179        module.lwe_switching_key_prepare(self, other, scratch);
180    }
181}
182
183pub trait LWESwitchingKeyPrepareAlloc<B: Backend> {
184    fn lwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
185    where
186        O: LWESwitchingKeyToRef;
187}
188
189impl<B: Backend> LWESwitchingKeyPrepareAlloc<B> for Module<B>
190where
191    Module<B>: LWESwitchingKeyPrepare<B>,
192{
193    fn lwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
194    where
195        O: LWESwitchingKeyToRef,
196    {
197        let mut ct_prep: LWESwitchingKeyPrepared<Vec<u8>, B> = LWESwitchingKeyPrepared::alloc(self, &other.to_ref());
198        self.lwe_switching_key_prepare(&mut ct_prep, other, scratch);
199        ct_prep
200    }
201}
202
203pub trait LWESwitchingKeyPreparedToRef<B: Backend> {
204    fn to_ref(&self) -> LWESwitchingKeyPrepared<&[u8], B>;
205}
206
207impl<D: DataRef, B: Backend> LWESwitchingKeyPreparedToRef<B> for LWESwitchingKeyPrepared<D, B>
208where
209    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToRef<B>,
210{
211    fn to_ref(&self) -> LWESwitchingKeyPrepared<&[u8], B> {
212        LWESwitchingKeyPrepared(self.0.to_ref())
213    }
214}
215
216pub trait LWESwitchingKeyPreparedToMut<B: Backend> {
217    fn to_mut(&mut self) -> LWESwitchingKeyPrepared<&mut [u8], B>;
218}
219
220impl<D: DataMut, B: Backend> LWESwitchingKeyPreparedToMut<B> for LWESwitchingKeyPrepared<D, B>
221where
222    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToMut<B>,
223{
224    fn to_mut(&mut self) -> LWESwitchingKeyPrepared<&mut [u8], B> {
225        LWESwitchingKeyPrepared(self.0.to_mut())
226    }
227}