1use poulpy_hal::{
  2    layouts::{Backend, Data, DataMut, DataRef, FillUniform, Module, ReaderFrom, WriterTo},
  3    source::Source,
  4};
  5
  6use crate::layouts::{
  7    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, GLWESwitchingKey, GLWESwitchingKeyAlloc, GLWESwitchingKeyToMut,
  8    GLWESwitchingKeyToRef, LWEInfos, Rank, TorusPrecision,
  9};
 10use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 11
 12use std::fmt;
 13
 14#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 15pub struct TensorKeyLayout {
 16    pub n: Degree,
 17    pub base2k: Base2K,
 18    pub k: TorusPrecision,
 19    pub rank: Rank,
 20    pub dnum: Dnum,
 21    pub dsize: Dsize,
 22}
 23
 24#[derive(PartialEq, Eq, Clone)]
 25pub struct TensorKey<D: Data> {
 26    pub(crate) keys: Vec<GLWESwitchingKey<D>>,
 27}
 28
 29impl<D: Data> LWEInfos for TensorKey<D> {
 30    fn n(&self) -> Degree {
 31        self.keys[0].n()
 32    }
 33
 34    fn base2k(&self) -> Base2K {
 35        self.keys[0].base2k()
 36    }
 37
 38    fn k(&self) -> TorusPrecision {
 39        self.keys[0].k()
 40    }
 41
 42    fn size(&self) -> usize {
 43        self.keys[0].size()
 44    }
 45}
 46
 47impl<D: Data> GLWEInfos for TensorKey<D> {
 48    fn rank(&self) -> Rank {
 49        self.keys[0].rank_out()
 50    }
 51}
 52
 53impl<D: Data> GGLWEInfos for TensorKey<D> {
 54    fn rank_in(&self) -> Rank {
 55        self.rank_out()
 56    }
 57
 58    fn rank_out(&self) -> Rank {
 59        self.keys[0].rank_out()
 60    }
 61
 62    fn dsize(&self) -> Dsize {
 63        self.keys[0].dsize()
 64    }
 65
 66    fn dnum(&self) -> Dnum {
 67        self.keys[0].dnum()
 68    }
 69}
 70
 71impl LWEInfos for TensorKeyLayout {
 72    fn n(&self) -> Degree {
 73        self.n
 74    }
 75
 76    fn base2k(&self) -> Base2K {
 77        self.base2k
 78    }
 79
 80    fn k(&self) -> TorusPrecision {
 81        self.k
 82    }
 83}
 84
 85impl GLWEInfos for TensorKeyLayout {
 86    fn rank(&self) -> Rank {
 87        self.rank_out()
 88    }
 89}
 90
 91impl GGLWEInfos for TensorKeyLayout {
 92    fn rank_in(&self) -> Rank {
 93        self.rank
 94    }
 95
 96    fn dsize(&self) -> Dsize {
 97        self.dsize
 98    }
 99
100    fn rank_out(&self) -> Rank {
101        self.rank
102    }
103
104    fn dnum(&self) -> Dnum {
105        self.dnum
106    }
107}
108
109impl<D: DataRef> fmt::Debug for TensorKey<D> {
110    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
111        write!(f, "{self}")
112    }
113}
114
115impl<D: DataMut> FillUniform for TensorKey<D> {
116    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
117        self.keys
118            .iter_mut()
119            .for_each(|key: &mut GLWESwitchingKey<D>| key.fill_uniform(log_bound, source))
120    }
121}
122
123impl<D: DataRef> fmt::Display for TensorKey<D> {
124    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
125        writeln!(f, "(GLWETensorKey)",)?;
126        for (i, key) in self.keys.iter().enumerate() {
127            write!(f, "{i}: {key}")?;
128        }
129        Ok(())
130    }
131}
132
133pub trait TensorKeyAlloc
134where
135    Self: GLWESwitchingKeyAlloc,
136{
137    fn alloc_tensor_key(&self, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> TensorKey<Vec<u8>> {
138        let pairs: u32 = (((rank.0 + 1) * rank.0) >> 1).max(1);
139        TensorKey {
140            keys: (0..pairs)
141                .map(|_| self.alloc_glwe_switching_key(base2k, k, Rank(1), rank, dnum, dsize))
142                .collect(),
143        }
144    }
145
146    fn alloc_tensor_key_from_infos<A>(&self, infos: &A) -> TensorKey<Vec<u8>>
147    where
148        A: GGLWEInfos,
149    {
150        assert_eq!(
151            infos.rank_in(),
152            infos.rank_out(),
153            "rank_in != rank_out is not supported for GGLWETensorKey"
154        );
155        self.alloc_tensor_key(
156            infos.base2k(),
157            infos.k(),
158            infos.rank(),
159            infos.dnum(),
160            infos.dsize(),
161        )
162    }
163
164    fn bytes_of_tensor_key(&self, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
165        let pairs: usize = (((rank.0 + 1) * rank.0) >> 1).max(1) as usize;
166        pairs * self.bytes_of_glwe_switching_key(base2k, k, Rank(1), rank, dnum, dsize)
167    }
168
169    fn bytes_of_tensor_key_from_infos<A>(&self, infos: &A) -> usize
170    where
171        A: GGLWEInfos,
172    {
173        assert_eq!(
174            infos.rank_in(),
175            infos.rank_out(),
176            "rank_in != rank_out is not supported for GGLWETensorKey"
177        );
178        self.bytes_of_tensor_key(
179            infos.base2k(),
180            infos.k(),
181            infos.rank(),
182            infos.dnum(),
183            infos.dsize(),
184        )
185    }
186}
187
188impl<B: Backend> TensorKeyAlloc for Module<B> where Self: GLWESwitchingKeyAlloc {}
189
190impl TensorKey<Vec<u8>> {
191    pub fn alloc_from_infos<A, B: Backend>(module: Module<B>, infos: &A) -> Self
192    where
193        A: GGLWEInfos,
194        Module<B>: TensorKeyAlloc,
195    {
196        module.alloc_tensor_key_from_infos(infos)
197    }
198
199    pub fn alloc<B: Backend>(module: Module<B>, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self
200    where
201        Module<B>: TensorKeyAlloc,
202    {
203        module.alloc_tensor_key(base2k, k, rank, dnum, dsize)
204    }
205
206    pub fn bytes_of_from_infos<A, B: Backend>(module: Module<B>, infos: &A) -> usize
207    where
208        A: GGLWEInfos,
209        Module<B>: TensorKeyAlloc,
210    {
211        module.bytes_of_tensor_key_from_infos(infos)
212    }
213
214    pub fn bytes_of<B: Backend>(
215        module: Module<B>,
216        base2k: Base2K,
217        k: TorusPrecision,
218        rank: Rank,
219        dnum: Dnum,
220        dsize: Dsize,
221    ) -> usize
222    where
223        Module<B>: TensorKeyAlloc,
224    {
225        module.bytes_of_tensor_key(base2k, k, rank, dnum, dsize)
226    }
227}
228
229impl<D: DataMut> TensorKey<D> {
230    // Returns a mutable reference to GLWESwitchingKey_{s}(s[i] * s[j])
231    pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKey<D> {
232        if i > j {
233            std::mem::swap(&mut i, &mut j);
234        };
235        let rank: usize = self.rank_out().into();
236        &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
237    }
238}
239
240impl<D: DataRef> TensorKey<D> {
241    // Returns a reference to GLWESwitchingKey_{s}(s[i] * s[j])
242    pub fn at(&self, mut i: usize, mut j: usize) -> &GLWESwitchingKey<D> {
243        if i > j {
244            std::mem::swap(&mut i, &mut j);
245        };
246        let rank: usize = self.rank_out().into();
247        &self.keys[i * rank + j - (i * (i + 1) / 2)]
248    }
249}
250
251impl<D: DataMut> ReaderFrom for TensorKey<D> {
252    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
253        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
254        if self.keys.len() != len {
255            return Err(std::io::Error::new(
256                std::io::ErrorKind::InvalidData,
257                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
258            ));
259        }
260        for key in &mut self.keys {
261            key.read_from(reader)?;
262        }
263        Ok(())
264    }
265}
266
267impl<D: DataRef> WriterTo for TensorKey<D> {
268    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
269        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
270        for key in &self.keys {
271            key.write_to(writer)?;
272        }
273        Ok(())
274    }
275}
276
277pub trait TensorKeyToRef {
278    fn to_ref(&self) -> TensorKey<&[u8]>;
279}
280
281impl<D: DataRef> TensorKeyToRef for TensorKey<D>
282where
283    GLWESwitchingKey<D>: GLWESwitchingKeyToRef,
284{
285    fn to_ref(&self) -> TensorKey<&[u8]> {
286        TensorKey {
287            keys: self.keys.iter().map(|c| c.to_ref()).collect(),
288        }
289    }
290}
291
292pub trait TensorKeyToMut {
293    fn to_mut(&mut self) -> TensorKey<&mut [u8]>;
294}
295
296impl<D: DataMut> TensorKeyToMut for TensorKey<D>
297where
298    GLWESwitchingKey<D>: GLWESwitchingKeyToMut,
299{
300    fn to_mut(&mut self) -> TensorKey<&mut [u8]> {
301        TensorKey {
302            keys: self.keys.iter_mut().map(|c| c.to_mut()).collect(),
303        }
304    }
305}