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(
 98            infos.base2k(),
 99            infos.k(),
100            infos.rank(),
101            infos.dnum(),
102            infos.dsize(),
103        )
104    }
105
106    fn alloc_gglwe_to_ggsw_key_prepared(
107        &self,
108        base2k: Base2K,
109        k: TorusPrecision,
110        rank: Rank,
111        dnum: Dnum,
112        dsize: Dsize,
113    ) -> GGLWEToGGSWKeyPrepared<Vec<u8>, BE> {
114        GGLWEToGGSWKeyPrepared {
115            keys: (0..rank.as_usize())
116                .map(|_| self.alloc_gglwe_prepared(base2k, k, rank, rank, dnum, dsize))
117                .collect(),
118        }
119    }
120
121    fn bytes_of_gglwe_to_ggsw_from_infos<A>(&self, infos: &A) -> usize
122    where
123        A: GGLWEInfos,
124    {
125        assert_eq!(
126            infos.rank_in(),
127            infos.rank_out(),
128            "rank_in != rank_out is not supported for GGLWEToGGSWKeyPrepared"
129        );
130        self.bytes_of_gglwe_to_ggsw(
131            infos.base2k(),
132            infos.k(),
133            infos.rank(),
134            infos.dnum(),
135            infos.dsize(),
136        )
137    }
138
139    fn bytes_of_gglwe_to_ggsw(&self, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
140        rank.as_usize() * self.bytes_of_gglwe_prepared(base2k, k, rank, rank, dnum, dsize)
141    }
142
143    fn prepare_gglwe_to_ggsw_key_tmp_bytes<A>(&self, infos: &A) -> usize
144    where
145        A: GGLWEInfos,
146    {
147        self.prepare_gglwe_tmp_bytes(infos)
148    }
149
150    fn prepare_gglwe_to_ggsw_key<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<BE>)
151    where
152        R: GGLWEToGGSWKeyPreparedToMut<BE>,
153        O: GGLWEToGGSWKeyToRef,
154    {
155        let res: &mut GGLWEToGGSWKeyPrepared<&mut [u8], BE> = &mut res.to_mut();
156        let other: &GGLWEToGGSWKey<&[u8]> = &other.to_ref();
157
158        assert_eq!(res.keys.len(), other.keys.len());
159
160        for (a, b) in res.keys.iter_mut().zip(other.keys.iter()) {
161            self.prepare_gglwe(a, b, scratch);
162        }
163    }
164}
165
166impl<BE: Backend> GGLWEToGGSWKeyPrepared<Vec<u8>, BE> {
167    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
168    where
169        A: GGLWEInfos,
170        M: GGLWEToGGSWKeyPreparedFactory<BE>,
171    {
172        module.alloc_gglwe_to_ggsw_key_prepared_from_infos(infos)
173    }
174
175    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self
176    where
177        M: GGLWEToGGSWKeyPreparedFactory<BE>,
178    {
179        module.alloc_gglwe_to_ggsw_key_prepared(base2k, k, rank, dnum, dsize)
180    }
181
182    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
183    where
184        A: GGLWEInfos,
185        M: GGLWEToGGSWKeyPreparedFactory<BE>,
186    {
187        module.bytes_of_gglwe_to_ggsw_from_infos(infos)
188    }
189
190    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize
191    where
192        M: GGLWEToGGSWKeyPreparedFactory<BE>,
193    {
194        module.bytes_of_gglwe_to_ggsw(base2k, k, rank, dnum, dsize)
195    }
196}
197
198impl<D: DataMut, BE: Backend> GGLWEToGGSWKeyPrepared<D, BE> {
199    pub fn prepare<M, O>(&mut self, module: &M, other: &O, scratch: &mut Scratch<BE>)
200    where
201        M: GGLWEToGGSWKeyPreparedFactory<BE>,
202        O: GGLWEToGGSWKeyToRef,
203    {
204        module.prepare_gglwe_to_ggsw_key(self, other, scratch);
205    }
206}
207
208impl<D: DataMut, BE: Backend> GGLWEToGGSWKeyPrepared<D, BE> {
209    // Returns a mutable reference to GGLWEPrepared_{s}([s[i]*s[0], s[i]*s[1], ..., s[i]*s[rank]])
210    pub fn at_mut(&mut self, i: usize) -> &mut GGLWEPrepared<D, BE> {
211        assert!((i as u32) < self.rank());
212        &mut self.keys[i]
213    }
214}
215
216impl<D: DataRef, BE: Backend> GGLWEToGGSWKeyPrepared<D, BE> {
217    // Returns a reference to GGLWEPrepared_{s}([s[i]*s[0], s[i]*s[1], ..., s[i]*s[rank]])
218    pub fn at(&self, i: usize) -> &GGLWEPrepared<D, BE> {
219        assert!((i as u32) < self.rank());
220        &self.keys[i]
221    }
222}
223
224pub trait GGLWEToGGSWKeyPreparedToRef<BE: Backend> {
225    fn to_ref(&self) -> GGLWEToGGSWKeyPrepared<&[u8], BE>;
226}
227
228impl<D: DataRef, BE: Backend> GGLWEToGGSWKeyPreparedToRef<BE> for GGLWEToGGSWKeyPrepared<D, BE>
229where
230    GGLWEPrepared<D, BE>: GGLWEPreparedToRef<BE>,
231{
232    fn to_ref(&self) -> GGLWEToGGSWKeyPrepared<&[u8], BE> {
233        GGLWEToGGSWKeyPrepared {
234            keys: self.keys.iter().map(|c| c.to_ref()).collect(),
235        }
236    }
237}
238
239pub trait GGLWEToGGSWKeyPreparedToMut<BE: Backend> {
240    fn to_mut(&mut self) -> GGLWEToGGSWKeyPrepared<&mut [u8], BE>;
241}
242
243impl<D: DataMut, BE: Backend> GGLWEToGGSWKeyPreparedToMut<BE> for GGLWEToGGSWKeyPrepared<D, BE>
244where
245    GGLWEPrepared<D, BE>: GGLWEPreparedToMut<BE>,
246{
247    fn to_mut(&mut self) -> GGLWEToGGSWKeyPrepared<&mut [u8], BE> {
248        GGLWEToGGSWKeyPrepared {
249            keys: self.keys.iter_mut().map(|c| c.to_mut()).collect(),
250        }
251    }
252}