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