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