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, GGLWECompressed, GGLWECompressedToMut, GGLWECompressedToRef, GGLWEDecompress, GGLWEInfos,
  8    GLWEInfos, GLWETensorKey, GLWETensorKeyToMut, LWEInfos, Rank, TorusPrecision,
  9};
 10use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 11use std::fmt;
 12
 13#[derive(PartialEq, Eq, Clone)]
 14pub struct GLWETensorKeyCompressed<D: Data> {
 15    pub(crate) keys: Vec<GGLWECompressed<D>>,
 16}
 17
 18impl<D: Data> LWEInfos for GLWETensorKeyCompressed<D> {
 19    fn n(&self) -> Degree {
 20        self.keys[0].n()
 21    }
 22
 23    fn base2k(&self) -> Base2K {
 24        self.keys[0].base2k()
 25    }
 26
 27    fn k(&self) -> TorusPrecision {
 28        self.keys[0].k()
 29    }
 30    fn size(&self) -> usize {
 31        self.keys[0].size()
 32    }
 33}
 34impl<D: Data> GLWEInfos for GLWETensorKeyCompressed<D> {
 35    fn rank(&self) -> Rank {
 36        self.rank_out()
 37    }
 38}
 39
 40impl<D: Data> GGLWEInfos for GLWETensorKeyCompressed<D> {
 41    fn rank_in(&self) -> Rank {
 42        self.rank_out()
 43    }
 44
 45    fn rank_out(&self) -> Rank {
 46        self.keys[0].rank_out()
 47    }
 48
 49    fn dsize(&self) -> Dsize {
 50        self.keys[0].dsize()
 51    }
 52
 53    fn dnum(&self) -> Dnum {
 54        self.keys[0].dnum()
 55    }
 56}
 57
 58impl<D: DataRef> fmt::Debug for GLWETensorKeyCompressed<D> {
 59    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 60        write!(f, "{self}")
 61    }
 62}
 63
 64impl<D: DataMut> FillUniform for GLWETensorKeyCompressed<D> {
 65    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 66        self.keys
 67            .iter_mut()
 68            .for_each(|key: &mut GGLWECompressed<D>| key.fill_uniform(log_bound, source))
 69    }
 70}
 71
 72impl<D: DataRef> fmt::Display for GLWETensorKeyCompressed<D> {
 73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 74        writeln!(f, "(GLWETensorKeyCompressed)",)?;
 75        for (i, key) in self.keys.iter().enumerate() {
 76            write!(f, "{i}: {key}")?;
 77        }
 78        Ok(())
 79    }
 80}
 81
 82impl GLWETensorKeyCompressed<Vec<u8>> {
 83    pub fn alloc_from_infos<A>(infos: &A) -> Self
 84    where
 85        A: GGLWEInfos,
 86    {
 87        Self::alloc(
 88            infos.n(),
 89            infos.base2k(),
 90            infos.k(),
 91            infos.rank(),
 92            infos.dnum(),
 93            infos.dsize(),
 94        )
 95    }
 96
 97    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
 98        let pairs: u32 = (((rank.as_u32() + 1) * rank.as_u32()) >> 1).max(1);
 99        GLWETensorKeyCompressed {
100            keys: (0..pairs)
101                .map(|_| GGLWECompressed::alloc(n, base2k, k, Rank(1), rank, dnum, dsize))
102                .collect(),
103        }
104    }
105
106    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
107    where
108        A: GGLWEInfos,
109    {
110        Self::bytes_of(
111            infos.n(),
112            infos.base2k(),
113            infos.k(),
114            infos.rank(),
115            infos.dnum(),
116            infos.dsize(),
117        )
118    }
119
120    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
121        let pairs: usize = (((rank.0 + 1) * rank.0) >> 1).max(1) as usize;
122        pairs * GGLWECompressed::bytes_of(n, base2k, k, Rank(1), dnum, dsize)
123    }
124}
125
126impl<D: DataMut> ReaderFrom for GLWETensorKeyCompressed<D> {
127    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
128        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
129        if self.keys.len() != len {
130            return Err(std::io::Error::new(
131                std::io::ErrorKind::InvalidData,
132                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
133            ));
134        }
135        for key in &mut self.keys {
136            key.read_from(reader)?;
137        }
138        Ok(())
139    }
140}
141
142impl<D: DataRef> WriterTo for GLWETensorKeyCompressed<D> {
143    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
144        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
145        for key in &self.keys {
146            key.write_to(writer)?;
147        }
148        Ok(())
149    }
150}
151
152pub trait GLWETensorKeyCompressedAtRef<D: DataRef> {
153    fn at(&self, i: usize, j: usize) -> &GGLWECompressed<D>;
154}
155
156impl<D: DataRef> GLWETensorKeyCompressedAtRef<D> for GLWETensorKeyCompressed<D> {
157    fn at(&self, mut i: usize, mut j: usize) -> &GGLWECompressed<D> {
158        if i > j {
159            std::mem::swap(&mut i, &mut j);
160        };
161        let rank: usize = self.rank_out().into();
162        &self.keys[i * rank + j - (i * (i + 1) / 2)]
163    }
164}
165
166pub trait GLWETensorKeyCompressedAtMut<D: DataMut> {
167    fn at_mut(&mut self, i: usize, j: usize) -> &mut GGLWECompressed<D>;
168}
169
170impl<D: DataMut> GLWETensorKeyCompressedAtMut<D> for GLWETensorKeyCompressed<D> {
171    fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GGLWECompressed<D> {
172        if i > j {
173            std::mem::swap(&mut i, &mut j);
174        };
175        let rank: usize = self.rank_out().into();
176        &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
177    }
178}
179
180pub trait GLWETensorKeyDecompress
181where
182    Self: GGLWEDecompress,
183{
184    fn decompress_tensor_key<R, O>(&self, res: &mut R, other: &O)
185    where
186        R: GLWETensorKeyToMut,
187        O: GLWETensorKeyCompressedToRef,
188    {
189        let res: &mut GLWETensorKey<&mut [u8]> = &mut res.to_mut();
190        let other: &GLWETensorKeyCompressed<&[u8]> = &other.to_ref();
191
192        assert_eq!(
193            res.keys.len(),
194            other.keys.len(),
195            "invalid receiver: res.keys.len()={} != other.keys.len()={}",
196            res.keys.len(),
197            other.keys.len()
198        );
199
200        for (a, b) in res.keys.iter_mut().zip(other.keys.iter()) {
201            self.decompress_gglwe(a, b);
202        }
203    }
204}
205
206impl<B: Backend> GLWETensorKeyDecompress for Module<B> where Self: GGLWEDecompress {}
207
208impl<D: DataMut> GLWETensorKey<D> {
209    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
210    where
211        O: GLWETensorKeyCompressedToRef,
212        M: GLWETensorKeyDecompress,
213    {
214        module.decompress_tensor_key(self, other);
215    }
216}
217
218pub trait GLWETensorKeyCompressedToMut {
219    fn to_mut(&mut self) -> GLWETensorKeyCompressed<&mut [u8]>;
220}
221
222impl<D: DataMut> GLWETensorKeyCompressedToMut for GLWETensorKeyCompressed<D>
223where
224    GGLWECompressed<D>: GGLWECompressedToMut,
225{
226    fn to_mut(&mut self) -> GLWETensorKeyCompressed<&mut [u8]> {
227        GLWETensorKeyCompressed {
228            keys: self.keys.iter_mut().map(|c| c.to_mut()).collect(),
229        }
230    }
231}
232
233pub trait GLWETensorKeyCompressedToRef {
234    fn to_ref(&self) -> GLWETensorKeyCompressed<&[u8]>;
235}
236
237impl<D: DataRef> GLWETensorKeyCompressedToRef for GLWETensorKeyCompressed<D>
238where
239    GGLWECompressed<D>: GGLWECompressedToRef,
240{
241    fn to_ref(&self) -> GLWETensorKeyCompressed<&[u8]> {
242        GLWETensorKeyCompressed {
243            keys: self.keys.iter().map(|c| c.to_ref()).collect(),
244        }
245    }
246}