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(
 99            n,
100            base2k,
101            k,
102            Rank(pairs),
103            rank,
104            dnum,
105            dsize,
106        ))
107    }
108
109    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
110    where
111        A: GGLWEInfos,
112    {
113        Self::bytes_of(
114            infos.n(),
115            infos.base2k(),
116            infos.k(),
117            infos.rank(),
118            infos.dnum(),
119            infos.dsize(),
120        )
121    }
122
123    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
124        let pairs: u32 = (((rank.as_u32() + 1) * rank.as_u32()) >> 1).max(1);
125        GGLWECompressed::bytes_of(n, base2k, k, Rank(pairs), dnum, dsize)
126    }
127}
128
129impl<D: DataMut> ReaderFrom for GLWETensorKeyCompressed<D> {
130    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
131        self.0.read_from(reader)?;
132        Ok(())
133    }
134}
135
136impl<D: DataRef> WriterTo for GLWETensorKeyCompressed<D> {
137    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
138        self.0.write_to(writer)?;
139        Ok(())
140    }
141}
142
143pub trait GLWETensorKeyDecompress
144where
145    Self: GGLWEDecompress,
146{
147    fn decompress_tensor_key<R, O>(&self, res: &mut R, other: &O)
148    where
149        R: GGLWEToMut,
150        O: GGLWECompressedToRef,
151    {
152        self.decompress_gglwe(res, other);
153    }
154}
155
156impl<B: Backend> GLWETensorKeyDecompress for Module<B> where Self: GGLWEDecompress {}
157
158impl<D: DataMut> GLWETensorKey<D> {
159    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
160    where
161        O: GGLWECompressedToRef,
162        M: GLWETensorKeyDecompress,
163    {
164        module.decompress_tensor_key(self, other);
165    }
166}
167
168impl<D: DataMut> GGLWECompressedToMut for GLWETensorKeyCompressed<D>
169where
170    GGLWECompressed<D>: GGLWECompressedToMut,
171{
172    fn to_mut(&mut self) -> GGLWECompressed<&mut [u8]> {
173        self.0.to_mut()
174    }
175}
176
177impl<D: DataRef> GGLWECompressedToRef for GLWETensorKeyCompressed<D>
178where
179    GGLWECompressed<D>: GGLWECompressedToRef,
180{
181    fn to_ref(&self) -> GGLWECompressed<&[u8]> {
182        self.0.to_ref()
183    }
184}