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, LWEToGLWESwitchingKeyToRef, Rank, TorusPrecision,
  8    prepared::{GLWESwitchingKeyPrepare, GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToMut, GLWESwitchingKeyPreparedToRef},
  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
 57impl<B: Backend> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
 58    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
 59    where
 60        A: GGLWEInfos,
 61        Module<B>: VmpPMatAlloc<B>,
 62    {
 63        debug_assert_eq!(
 64            infos.rank_in().0,
 65            1,
 66            "rank_in > 1 is not supported for LWEToGLWESwitchingKey"
 67        );
 68        debug_assert_eq!(
 69            infos.dsize().0,
 70            1,
 71            "dsize > 1 is not supported for LWEToGLWESwitchingKey"
 72        );
 73        Self(GLWESwitchingKeyPrepared::alloc_from_infos(module, infos))
 74    }
 75
 76    pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> Self
 77    where
 78        Module<B>: VmpPMatAlloc<B>,
 79    {
 80        Self(GLWESwitchingKeyPrepared::alloc(
 81            module,
 82            base2k,
 83            k,
 84            Rank(1),
 85            rank_out,
 86            dnum,
 87            Dsize(1),
 88        ))
 89    }
 90
 91    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
 92    where
 93        A: GGLWEInfos,
 94        Module<B>: VmpPMatAllocBytes,
 95    {
 96        debug_assert_eq!(
 97            infos.rank_in().0,
 98            1,
 99            "rank_in > 1 is not supported for LWEToGLWESwitchingKey"
100        );
101        debug_assert_eq!(
102            infos.dsize().0,
103            1,
104            "dsize > 1 is not supported for LWEToGLWESwitchingKey"
105        );
106        GLWESwitchingKeyPrepared::alloc_bytes_from_infos(module, infos)
107    }
108
109    pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum, rank_out: Rank) -> usize
110    where
111        Module<B>: VmpPMatAllocBytes,
112    {
113        GLWESwitchingKeyPrepared::alloc_bytes(module, base2k, k, Rank(1), rank_out, dnum, Dsize(1))
114    }
115}
116
117pub trait LWEToGLWESwitchingKeyPrepareTmpBytes {
118    fn lwe_to_glwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
119    where
120        A: GGLWEInfos;
121}
122
123impl<B: Backend> LWEToGLWESwitchingKeyPrepareTmpBytes for Module<B>
124where
125    Module<B>: LWEToGLWESwitchingKeyPrepareTmpBytes,
126{
127    fn lwe_to_glwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
128    where
129        A: GGLWEInfos,
130    {
131        self.lwe_to_glwe_switching_key_prepare_tmp_bytes(infos);
132    }
133}
134
135impl<B: Backend> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
136    pub fn prepare_tmp_bytes<A>(&self, module: &Module<B>, infos: &A)
137    where
138        A: GLWEInfos,
139        Module<B>: LWEToGLWESwitchingKeyPrepareTmpBytes,
140    {
141        module.glwe_secret_prepare_tmp_bytes(infos);
142    }
143}
144
145pub trait LWEToGLWESwitchingKeyPrepare<B: Backend> {
146    fn lwe_to_glwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
147    where
148        R: LWEToGLWESwitchingKeyPreparedToMut<B>,
149        O: LWEToGLWESwitchingKeyToRef;
150}
151
152impl<B: Backend> LWEToGLWESwitchingKeyPrepare<B> for Module<B>
153where
154    Module<B>: GLWESwitchingKeyPrepare<B>,
155{
156    fn lwe_to_glwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
157    where
158        R: LWEToGLWESwitchingKeyPreparedToMut<B>,
159        O: LWEToGLWESwitchingKeyToRef,
160    {
161        self.glwe_switching_prepare(&mut res.to_mut().0, other, scratch);
162    }
163}
164
165impl<D: DataMut, B: Backend> LWEToGLWESwitchingKeyPrepared<D, B> {
166    fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &Scratch<B>)
167    where
168        O: LWEToGLWESwitchingKeyToRef,
169        Module<B>: LWEToGLWESwitchingKeyPrepare<B>,
170    {
171        module.lwe_to_glwe_switching_key_prepare(self, other, scratch);
172    }
173}
174
175pub trait LWEToGLWESwitchingKeyPrepareAlloc<B: Backend> {
176    fn lwe_to_glwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
177    where
178        O: LWEToGLWESwitchingKeyToRef;
179}
180
181impl<B: Backend> LWEToGLWESwitchingKeyPrepareAlloc<B> for Module<B>
182where
183    Module<B>: LWEToGLWESwitchingKeyPrepare<B>,
184{
185    fn lwe_to_glwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
186    where
187        O: LWEToGLWESwitchingKeyToRef,
188    {
189        let mut ct_prep: LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> = LWEToGLWESwitchingKeyPrepared::alloc(self, &other.to_ref());
190        self.lwe_to_glwe_switching_key_prepare(&mut ct_prep, other, scratch);
191        ct_prep
192    }
193}
194
195pub trait LWEToGLWESwitchingKeyPreparedToRef<B: Backend> {
196    fn to_ref(&self) -> LWEToGLWESwitchingKeyPrepared<&[u8], B>;
197}
198
199impl<D: DataRef, B: Backend> LWEToGLWESwitchingKeyPreparedToRef<B> for LWEToGLWESwitchingKeyPrepared<D, B>
200where
201    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToRef<B>,
202{
203    fn to_ref(&self) -> LWEToGLWESwitchingKeyPrepared<&[u8], B> {
204        LWEToGLWESwitchingKeyPrepared(self.0.to_ref())
205    }
206}
207
208pub trait LWEToGLWESwitchingKeyPreparedToMut<B: Backend> {
209    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyPrepared<&mut [u8], B>;
210}
211
212impl<D: DataMut, B: Backend> LWEToGLWESwitchingKeyPreparedToMut<B> for LWEToGLWESwitchingKeyPrepared<D, B>
213where
214    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToMut<B>,
215{
216    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyPrepared<&mut [u8], B> {
217        LWEToGLWESwitchingKeyPrepared(self.0.to_mut())
218    }
219}