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