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