1use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
  2
  3use crate::layouts::{
  4    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GGLWEPrepared, GGLWEPreparedFactory, GGLWEPreparedToMut, GGLWEPreparedToRef,
  5    GGLWEToGGSWKey, GGLWEToGGSWKeyToRef, GLWEInfos, LWEInfos, Rank, TorusPrecision,
  6};
  7
  8pub struct GGLWEToGGSWKeyPrepared<D: Data, BE: Backend> {
  9    pub(crate) keys: Vec<GGLWEPrepared<D, BE>>,
 10}
 11
 12impl<D: Data, BE: Backend> LWEInfos for GGLWEToGGSWKeyPrepared<D, BE> {
 13    fn n(&self) -> Degree {
 14        self.keys[0].n()
 15    }
 16
 17    fn base2k(&self) -> Base2K {
 18        self.keys[0].base2k()
 19    }
 20
 21    fn k(&self) -> TorusPrecision {
 22        self.keys[0].k()
 23    }
 24
 25    fn size(&self) -> usize {
 26        self.keys[0].size()
 27    }
 28}
 29
 30impl<D: Data, BE: Backend> GLWEInfos for GGLWEToGGSWKeyPrepared<D, BE> {
 31    fn rank(&self) -> Rank {
 32        self.keys[0].rank_out()
 33    }
 34}
 35
 36impl<D: Data, BE: Backend> GGLWEInfos for GGLWEToGGSWKeyPrepared<D, BE> {
 37    fn rank_in(&self) -> Rank {
 38        self.rank_out()
 39    }
 40
 41    fn rank_out(&self) -> Rank {
 42        self.keys[0].rank_out()
 43    }
 44
 45    fn dsize(&self) -> Dsize {
 46        self.keys[0].dsize()
 47    }
 48
 49    fn dnum(&self) -> Dnum {
 50        self.keys[0].dnum()
 51    }
 52}
 53
 54pub trait GGLWEToGGSWKeyPreparedFactory<BE: Backend> {
 55    fn alloc_gglwe_to_ggsw_key_prepared_from_infos<A>(&self, infos: &A) -> GGLWEToGGSWKeyPrepared<Vec<u8>, BE>
 56    where
 57        A: GGLWEInfos;
 58
 59    fn alloc_gglwe_to_ggsw_key_prepared(
 60        &self,
 61        base2k: Base2K,
 62        k: TorusPrecision,
 63        rank: Rank,
 64        dnum: Dnum,
 65        dsize: Dsize,
 66    ) -> GGLWEToGGSWKeyPrepared<Vec<u8>, BE>;
 67
 68    fn bytes_of_gglwe_to_ggsw_from_infos<A>(&self, infos: &A) -> usize
 69    where
 70        A: GGLWEInfos;
 71
 72    fn bytes_of_gglwe_to_ggsw(&self, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize;
 73
 74    fn prepare_gglwe_to_ggsw_key_tmp_bytes<A>(&self, infos: &A) -> usize
 75    where
 76        A: GGLWEInfos;
 77
 78    fn prepare_gglwe_to_ggsw_key<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<BE>)
 79    where
 80        R: GGLWEToGGSWKeyPreparedToMut<BE>,
 81        O: GGLWEToGGSWKeyToRef;
 82}
 83
 84impl<BE: Backend> GGLWEToGGSWKeyPreparedFactory<BE> for Module<BE>
 85where
 86    Self: GGLWEPreparedFactory<BE>,
 87{
 88    fn alloc_gglwe_to_ggsw_key_prepared_from_infos<A>(&self, infos: &A) -> GGLWEToGGSWKeyPrepared<Vec<u8>, BE>
 89    where
 90        A: GGLWEInfos,
 91    {
 92        assert_eq!(
 93            infos.rank_in(),
 94            infos.rank_out(),
 95            "rank_in != rank_out is not supported for GGLWEToGGSWKeyPrepared"
 96        );
 97        self.alloc_gglwe_to_ggsw_key_prepared(infos.base2k(), infos.k(), infos.rank(), infos.dnum(), infos.dsize())
 98    }
 99
100    fn alloc_gglwe_to_ggsw_key_prepared(
101        &self,
102        base2k: Base2K,
103        k: TorusPrecision,
104        rank: Rank,
105        dnum: Dnum,
106        dsize: Dsize,
107    ) -> GGLWEToGGSWKeyPrepared<Vec<u8>, BE> {
108        GGLWEToGGSWKeyPrepared {
109            keys: (0..rank.as_usize())
110                .map(|_| self.alloc_gglwe_prepared(base2k, k, rank, rank, dnum, dsize))
111                .collect(),
112        }
113    }
114
115    fn bytes_of_gglwe_to_ggsw_from_infos<A>(&self, infos: &A) -> usize
116    where
117        A: GGLWEInfos,
118    {
119        assert_eq!(
120            infos.rank_in(),
121            infos.rank_out(),
122            "rank_in != rank_out is not supported for GGLWEToGGSWKeyPrepared"
123        );
124        self.bytes_of_gglwe_to_ggsw(infos.base2k(), infos.k(), infos.rank(), infos.dnum(), infos.dsize())
125    }
126
127    fn bytes_of_gglwe_to_ggsw(&self, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
128        rank.as_usize() * self.bytes_of_gglwe_prepared(base2k, k, rank, rank, dnum, dsize)
129    }
130
131    fn prepare_gglwe_to_ggsw_key_tmp_bytes<A>(&self, infos: &A) -> usize
132    where
133        A: GGLWEInfos,
134    {
135        self.prepare_gglwe_tmp_bytes(infos)
136    }
137
138    fn prepare_gglwe_to_ggsw_key<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<BE>)
139    where
140        R: GGLWEToGGSWKeyPreparedToMut<BE>,
141        O: GGLWEToGGSWKeyToRef,
142    {
143        let res: &mut GGLWEToGGSWKeyPrepared<&mut [u8], BE> = &mut res.to_mut();
144        let other: &GGLWEToGGSWKey<&[u8]> = &other.to_ref();
145
146        assert_eq!(res.keys.len(), other.keys.len());
147
148        for (a, b) in res.keys.iter_mut().zip(other.keys.iter()) {
149            self.prepare_gglwe(a, b, scratch);
150        }
151    }
152}
153
154impl<BE: Backend> GGLWEToGGSWKeyPrepared<Vec<u8>, BE> {
155    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
156    where
157        A: GGLWEInfos,
158        M: GGLWEToGGSWKeyPreparedFactory<BE>,
159    {
160        module.alloc_gglwe_to_ggsw_key_prepared_from_infos(infos)
161    }
162
163    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self
164    where
165        M: GGLWEToGGSWKeyPreparedFactory<BE>,
166    {
167        module.alloc_gglwe_to_ggsw_key_prepared(base2k, k, rank, dnum, dsize)
168    }
169
170    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
171    where
172        A: GGLWEInfos,
173        M: GGLWEToGGSWKeyPreparedFactory<BE>,
174    {
175        module.bytes_of_gglwe_to_ggsw_from_infos(infos)
176    }
177
178    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize
179    where
180        M: GGLWEToGGSWKeyPreparedFactory<BE>,
181    {
182        module.bytes_of_gglwe_to_ggsw(base2k, k, rank, dnum, dsize)
183    }
184}
185
186impl<D: DataMut, BE: Backend> GGLWEToGGSWKeyPrepared<D, BE> {
187    pub fn prepare<M, O>(&mut self, module: &M, other: &O, scratch: &mut Scratch<BE>)
188    where
189        M: GGLWEToGGSWKeyPreparedFactory<BE>,
190        O: GGLWEToGGSWKeyToRef,
191    {
192        module.prepare_gglwe_to_ggsw_key(self, other, scratch);
193    }
194}
195
196impl<D: DataMut, BE: Backend> GGLWEToGGSWKeyPrepared<D, BE> {
197    // Returns a mutable reference to GGLWEPrepared_{s}([s[i]*s[0], s[i]*s[1], ..., s[i]*s[rank]])
198    pub fn at_mut(&mut self, i: usize) -> &mut GGLWEPrepared<D, BE> {
199        assert!((i as u32) < self.rank());
200        &mut self.keys[i]
201    }
202}
203
204impl<D: DataRef, BE: Backend> GGLWEToGGSWKeyPrepared<D, BE> {
205    // Returns a reference to GGLWEPrepared_{s}([s[i]*s[0], s[i]*s[1], ..., s[i]*s[rank]])
206    pub fn at(&self, i: usize) -> &GGLWEPrepared<D, BE> {
207        assert!((i as u32) < self.rank());
208        &self.keys[i]
209    }
210}
211
212pub trait GGLWEToGGSWKeyPreparedToRef<BE: Backend> {
213    fn to_ref(&self) -> GGLWEToGGSWKeyPrepared<&[u8], BE>;
214}
215
216impl<D: DataRef, BE: Backend> GGLWEToGGSWKeyPreparedToRef<BE> for GGLWEToGGSWKeyPrepared<D, BE>
217where
218    GGLWEPrepared<D, BE>: GGLWEPreparedToRef<BE>,
219{
220    fn to_ref(&self) -> GGLWEToGGSWKeyPrepared<&[u8], BE> {
221        GGLWEToGGSWKeyPrepared {
222            keys: self.keys.iter().map(|c| c.to_ref()).collect(),
223        }
224    }
225}
226
227pub trait GGLWEToGGSWKeyPreparedToMut<BE: Backend> {
228    fn to_mut(&mut self) -> GGLWEToGGSWKeyPrepared<&mut [u8], BE>;
229}
230
231impl<D: DataMut, BE: Backend> GGLWEToGGSWKeyPreparedToMut<BE> for GGLWEToGGSWKeyPrepared<D, BE>
232where
233    GGLWEPrepared<D, BE>: GGLWEPreparedToMut<BE>,
234{
235    fn to_mut(&mut self) -> GGLWEToGGSWKeyPrepared<&mut [u8], BE> {
236        GGLWEToGGSWKeyPrepared {
237            keys: self.keys.iter_mut().map(|c| c.to_mut()).collect(),
238        }
239    }
240}