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, GGLWECompressedSeedMut, GGLWECompressedToMut, GGLWECompressedToRef,
  8    GGLWEDecompress, GGLWEInfos, GGLWEToMut, GLWEInfos, GLWETensorKey, LWEInfos, Rank, TorusPrecision,
  9};
 10use std::fmt;
 11
 12#[derive(PartialEq, Eq, Clone)]
 13pub struct GLWETensorKeyCompressed<D: Data>(pub(crate) GGLWECompressed<D>);
 14
 15impl<D: DataMut> GGLWECompressedSeedMut for GLWETensorKeyCompressed<D> {
 16    fn seed_mut(&mut self) -> &mut Vec<[u8; 32]> {
 17        &mut self.0.seed
 18    }
 19}
 20
 21impl<D: Data> LWEInfos for GLWETensorKeyCompressed<D> {
 22    fn n(&self) -> Degree {
 23        self.0.n()
 24    }
 25
 26    fn base2k(&self) -> Base2K {
 27        self.0.base2k()
 28    }
 29
 30    fn k(&self) -> TorusPrecision {
 31        self.0.k()
 32    }
 33    fn size(&self) -> usize {
 34        self.0.size()
 35    }
 36}
 37impl<D: Data> GLWEInfos for GLWETensorKeyCompressed<D> {
 38    fn rank(&self) -> Rank {
 39        self.rank_out()
 40    }
 41}
 42
 43impl<D: Data> GGLWEInfos for GLWETensorKeyCompressed<D> {
 44    fn rank_in(&self) -> Rank {
 45        self.rank_out()
 46    }
 47
 48    fn rank_out(&self) -> Rank {
 49        self.0.rank_out()
 50    }
 51
 52    fn dsize(&self) -> Dsize {
 53        self.0.dsize()
 54    }
 55
 56    fn dnum(&self) -> Dnum {
 57        self.0.dnum()
 58    }
 59}
 60
 61impl<D: DataRef> fmt::Debug for GLWETensorKeyCompressed<D> {
 62    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 63        write!(f, "{self}")
 64    }
 65}
 66
 67impl<D: DataMut> FillUniform for GLWETensorKeyCompressed<D> {
 68    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 69        self.0.fill_uniform(log_bound, source);
 70    }
 71}
 72
 73impl<D: DataRef> fmt::Display for GLWETensorKeyCompressed<D> {
 74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 75        writeln!(f, "(GLWETensorKeyCompressed)",)?;
 76        write!(f, "{}", self.0)?;
 77        Ok(())
 78    }
 79}
 80
 81impl GLWETensorKeyCompressed<Vec<u8>> {
 82    pub fn alloc_from_infos<A>(infos: &A) -> Self
 83    where
 84        A: GGLWEInfos,
 85    {
 86        Self::alloc(
 87            infos.n(),
 88            infos.base2k(),
 89            infos.k(),
 90            infos.rank(),
 91            infos.dnum(),
 92            infos.dsize(),
 93        )
 94    }
 95
 96    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
 97        let pairs: u32 = (((rank.as_u32() + 1) * rank.as_u32()) >> 1).max(1);
 98        GLWETensorKeyCompressed(GGLWECompressed::alloc(n, base2k, k, Rank(pairs), rank, dnum, dsize))
 99    }
100
101    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
102    where
103        A: GGLWEInfos,
104    {
105        Self::bytes_of(
106            infos.n(),
107            infos.base2k(),
108            infos.k(),
109            infos.rank(),
110            infos.dnum(),
111            infos.dsize(),
112        )
113    }
114
115    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
116        let pairs: u32 = (((rank.as_u32() + 1) * rank.as_u32()) >> 1).max(1);
117        GGLWECompressed::bytes_of(n, base2k, k, Rank(pairs), dnum, dsize)
118    }
119}
120
121impl<D: DataMut> ReaderFrom for GLWETensorKeyCompressed<D> {
122    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
123        self.0.read_from(reader)?;
124        Ok(())
125    }
126}
127
128impl<D: DataRef> WriterTo for GLWETensorKeyCompressed<D> {
129    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
130        self.0.write_to(writer)?;
131        Ok(())
132    }
133}
134
135pub trait GLWETensorKeyDecompress
136where
137    Self: GGLWEDecompress,
138{
139    fn decompress_tensor_key<R, O>(&self, res: &mut R, other: &O)
140    where
141        R: GGLWEToMut,
142        O: GGLWECompressedToRef,
143    {
144        self.decompress_gglwe(res, other);
145    }
146}
147
148impl<B: Backend> GLWETensorKeyDecompress for Module<B> where Self: GGLWEDecompress {}
149
150impl<D: DataMut> GLWETensorKey<D> {
151    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
152    where
153        O: GGLWECompressedToRef,
154        M: GLWETensorKeyDecompress,
155    {
156        module.decompress_tensor_key(self, other);
157    }
158}
159
160impl<D: DataMut> GGLWECompressedToMut for GLWETensorKeyCompressed<D>
161where
162    GGLWECompressed<D>: GGLWECompressedToMut,
163{
164    fn to_mut(&mut self) -> GGLWECompressed<&mut [u8]> {
165        self.0.to_mut()
166    }
167}
168
169impl<D: DataRef> GGLWECompressedToRef for GLWETensorKeyCompressed<D>
170where
171    GGLWECompressed<D>: GGLWECompressedToRef,
172{
173    fn to_ref(&self) -> GGLWECompressed<&[u8]> {
174        self.0.to_ref()
175    }
176}