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