1use std::collections::HashMap;
  2
  3use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
  4
  5use crate::layouts::{
  6    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GGLWELayout, GGLWEPrepared, GGLWEPreparedFactory, GGLWEPreparedToMut,
  7    GGLWEPreparedToRef, GGLWEToRef, GLWEAutomorphismKeyHelper, GLWEInfos, GetGaloisElement, LWEInfos, Rank, SetGaloisElement,
  8    TorusPrecision,
  9};
 10
 11impl<K, BE: Backend> GLWEAutomorphismKeyHelper<K, BE> for HashMap<i64, K>
 12where
 13    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 14{
 15    fn get_automorphism_key(&self, k: i64) -> Option<&K> {
 16        self.get(&k)
 17    }
 18
 19    fn automorphism_key_infos(&self) -> GGLWELayout {
 20        self.get(self.keys().next().unwrap())
 21            .unwrap()
 22            .gglwe_layout()
 23    }
 24}
 25
 26#[derive(PartialEq, Eq)]
 27pub struct GLWEAutomorphismKeyPrepared<D: Data, B: Backend> {
 28    pub(crate) key: GGLWEPrepared<D, B>,
 29    pub(crate) p: i64,
 30}
 31
 32impl<D: Data, B: Backend> LWEInfos for GLWEAutomorphismKeyPrepared<D, B> {
 33    fn n(&self) -> Degree {
 34        self.key.n()
 35    }
 36
 37    fn base2k(&self) -> Base2K {
 38        self.key.base2k()
 39    }
 40
 41    fn k(&self) -> TorusPrecision {
 42        self.key.k()
 43    }
 44
 45    fn size(&self) -> usize {
 46        self.key.size()
 47    }
 48}
 49
 50impl<D: Data, B: Backend> GetGaloisElement for GLWEAutomorphismKeyPrepared<D, B> {
 51    fn p(&self) -> i64 {
 52        self.p
 53    }
 54}
 55
 56impl<D: Data, B: Backend> SetGaloisElement for GLWEAutomorphismKeyPrepared<D, B> {
 57    fn set_p(&mut self, p: i64) {
 58        self.p = p
 59    }
 60}
 61
 62impl<D: Data, B: Backend> GLWEInfos for GLWEAutomorphismKeyPrepared<D, B> {
 63    fn rank(&self) -> Rank {
 64        self.rank_out()
 65    }
 66}
 67
 68impl<D: Data, B: Backend> GGLWEInfos for GLWEAutomorphismKeyPrepared<D, B> {
 69    fn rank_in(&self) -> Rank {
 70        self.key.rank_in()
 71    }
 72
 73    fn rank_out(&self) -> Rank {
 74        self.key.rank_out()
 75    }
 76
 77    fn dsize(&self) -> Dsize {
 78        self.key.dsize()
 79    }
 80
 81    fn dnum(&self) -> Dnum {
 82        self.key.dnum()
 83    }
 84}
 85
 86pub trait GLWEAutomorphismKeyPreparedFactory<B: Backend>
 87where
 88    Self: GGLWEPreparedFactory<B>,
 89{
 90    fn alloc_glwe_automorphism_key_prepared(
 91        &self,
 92        base2k: Base2K,
 93        k: TorusPrecision,
 94        rank: Rank,
 95        dnum: Dnum,
 96        dsize: Dsize,
 97    ) -> GLWEAutomorphismKeyPrepared<Vec<u8>, B> {
 98        GLWEAutomorphismKeyPrepared::<Vec<u8>, B> {
 99            key: self.alloc_gglwe_prepared(base2k, k, rank, rank, dnum, dsize),
100            p: 0,
101        }
102    }
103
104    fn alloc_glwe_automorphism_key_prepared_from_infos<A>(&self, infos: &A) -> GLWEAutomorphismKeyPrepared<Vec<u8>, B>
105    where
106        A: GGLWEInfos,
107    {
108        assert_eq!(
109            infos.rank_in(),
110            infos.rank_out(),
111            "rank_in != rank_out is not supported for AutomorphismKeyPrepared"
112        );
113        self.alloc_glwe_automorphism_key_prepared(
114            infos.base2k(),
115            infos.k(),
116            infos.rank(),
117            infos.dnum(),
118            infos.dsize(),
119        )
120    }
121
122    fn bytes_of_glwe_automorphism_key_prepared(
123        &self,
124        base2k: Base2K,
125        k: TorusPrecision,
126        rank: Rank,
127        dnum: Dnum,
128        dsize: Dsize,
129    ) -> usize {
130        self.bytes_of_gglwe_prepared(base2k, k, rank, rank, dnum, dsize)
131    }
132
133    fn bytes_of_glwe_automorphism_key_prepared_from_infos<A>(&self, infos: &A) -> usize
134    where
135        A: GGLWEInfos,
136    {
137        assert_eq!(
138            infos.rank_in(),
139            infos.rank_out(),
140            "rank_in != rank_out is not supported for AutomorphismKeyPrepared"
141        );
142        self.bytes_of_glwe_automorphism_key_prepared(
143            infos.base2k(),
144            infos.k(),
145            infos.rank(),
146            infos.dnum(),
147            infos.dsize(),
148        )
149    }
150
151    fn prepare_glwe_automorphism_key_tmp_bytes<A>(&self, infos: &A) -> usize
152    where
153        A: GGLWEInfos,
154    {
155        self.prepare_gglwe_tmp_bytes(infos)
156    }
157
158    fn prepare_glwe_automorphism_key<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
159    where
160        R: GGLWEPreparedToMut<B> + SetGaloisElement,
161        O: GGLWEToRef + GetGaloisElement,
162    {
163        self.prepare_gglwe(res, other, scratch);
164        res.set_p(other.p());
165    }
166}
167
168impl<B: Backend> GLWEAutomorphismKeyPreparedFactory<B> for Module<B> where Module<B>: GGLWEPreparedFactory<B> {}
169
170impl<B: Backend> GLWEAutomorphismKeyPrepared<Vec<u8>, B> {
171    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
172    where
173        A: GGLWEInfos,
174        M: GLWEAutomorphismKeyPreparedFactory<B>,
175    {
176        module.alloc_glwe_automorphism_key_prepared_from_infos(infos)
177    }
178
179    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self
180    where
181        M: GLWEAutomorphismKeyPreparedFactory<B>,
182    {
183        module.alloc_glwe_automorphism_key_prepared(base2k, k, rank, dnum, dsize)
184    }
185
186    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
187    where
188        A: GGLWEInfos,
189        M: GLWEAutomorphismKeyPreparedFactory<B>,
190    {
191        module.bytes_of_glwe_automorphism_key_prepared_from_infos(infos)
192    }
193
194    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize
195    where
196        M: GLWEAutomorphismKeyPreparedFactory<B>,
197    {
198        module.bytes_of_glwe_automorphism_key_prepared(base2k, k, rank, dnum, dsize)
199    }
200}
201
202impl<B: Backend> GLWEAutomorphismKeyPrepared<Vec<u8>, B> {
203    pub fn prepare_tmp_bytes<M>(&self, module: &M) -> usize
204    where
205        M: GLWEAutomorphismKeyPreparedFactory<B>,
206    {
207        module.prepare_glwe_automorphism_key_tmp_bytes(self)
208    }
209}
210
211impl<D: DataMut, B: Backend> GLWEAutomorphismKeyPrepared<D, B> {
212    pub fn prepare<O, M>(&mut self, module: &M, other: &O, scratch: &mut Scratch<B>)
213    where
214        O: GGLWEToRef + GetGaloisElement,
215        M: GLWEAutomorphismKeyPreparedFactory<B>,
216    {
217        module.prepare_glwe_automorphism_key(self, other, scratch);
218    }
219}
220
221impl<D: DataMut, B: Backend> GGLWEPreparedToMut<B> for GLWEAutomorphismKeyPrepared<D, B> {
222    fn to_mut(&mut self) -> GGLWEPrepared<&mut [u8], B> {
223        self.key.to_mut()
224    }
225}
226
227impl<D: DataRef, BE: Backend> GGLWEPreparedToRef<BE> for GLWEAutomorphismKeyPrepared<D, BE> {
228    fn to_ref(&self) -> GGLWEPrepared<&[u8], BE> {
229        self.key.to_ref()
230    }
231}