1use poulpy_hal::{
  2    layouts::{Data, DataMut, DataRef, FillUniform, ReaderFrom, WriterTo},
  3    source::Source,
  4};
  5
  6use crate::layouts::{
  7    Base2K, Degree, Dnum, Dsize, GGLWE, GGLWEInfos, GGLWEToMut, GGLWEToRef, GLWEInfos, LWEInfos, Rank, TorusPrecision,
  8};
  9
 10use std::fmt;
 11
 12#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 13pub struct GLWETensorKeyLayout {
 14    pub n: Degree,
 15    pub base2k: Base2K,
 16    pub k: TorusPrecision,
 17    pub rank: Rank,
 18    pub dnum: Dnum,
 19    pub dsize: Dsize,
 20}
 21
 22#[derive(PartialEq, Eq, Clone)]
 23pub struct GLWETensorKey<D: Data>(pub(crate) GGLWE<D>);
 24
 25impl<D: Data> LWEInfos for GLWETensorKey<D> {
 26    fn n(&self) -> Degree {
 27        self.0.n()
 28    }
 29
 30    fn base2k(&self) -> Base2K {
 31        self.0.base2k()
 32    }
 33
 34    fn k(&self) -> TorusPrecision {
 35        self.0.k()
 36    }
 37
 38    fn size(&self) -> usize {
 39        self.0.size()
 40    }
 41}
 42
 43impl<D: Data> GLWEInfos for GLWETensorKey<D> {
 44    fn rank(&self) -> Rank {
 45        self.0.rank_out()
 46    }
 47}
 48
 49impl<D: Data> GGLWEInfos for GLWETensorKey<D> {
 50    fn rank_in(&self) -> Rank {
 51        self.rank_out()
 52    }
 53
 54    fn rank_out(&self) -> Rank {
 55        self.0.rank_out()
 56    }
 57
 58    fn dsize(&self) -> Dsize {
 59        self.0.dsize()
 60    }
 61
 62    fn dnum(&self) -> Dnum {
 63        self.0.dnum()
 64    }
 65}
 66
 67impl LWEInfos for GLWETensorKeyLayout {
 68    fn n(&self) -> Degree {
 69        self.n
 70    }
 71
 72    fn base2k(&self) -> Base2K {
 73        self.base2k
 74    }
 75
 76    fn k(&self) -> TorusPrecision {
 77        self.k
 78    }
 79}
 80
 81impl GLWEInfos for GLWETensorKeyLayout {
 82    fn rank(&self) -> Rank {
 83        self.rank_out()
 84    }
 85}
 86
 87impl GGLWEInfos for GLWETensorKeyLayout {
 88    fn rank_in(&self) -> Rank {
 89        self.rank
 90    }
 91
 92    fn dsize(&self) -> Dsize {
 93        self.dsize
 94    }
 95
 96    fn rank_out(&self) -> Rank {
 97        self.rank
 98    }
 99
100    fn dnum(&self) -> Dnum {
101        self.dnum
102    }
103}
104
105impl<D: DataRef> fmt::Debug for GLWETensorKey<D> {
106    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
107        write!(f, "{self}")
108    }
109}
110
111impl<D: DataMut> FillUniform for GLWETensorKey<D> {
112    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
113        self.0.fill_uniform(log_bound, source)
114    }
115}
116
117impl<D: DataRef> fmt::Display for GLWETensorKey<D> {
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        writeln!(f, "(GLWETensorKey)",)?;
120        write!(f, "{}", self.0)?;
121        Ok(())
122    }
123}
124
125impl GLWETensorKey<Vec<u8>> {
126    pub fn alloc_from_infos<A>(infos: &A) -> Self
127    where
128        A: GGLWEInfos,
129    {
130        assert_eq!(
131            infos.rank_in(),
132            infos.rank_out(),
133            "rank_in != rank_out is not supported for GGLWETensorKey"
134        );
135        Self::alloc(
136            infos.n(),
137            infos.base2k(),
138            infos.k(),
139            infos.rank(),
140            infos.dnum(),
141            infos.dsize(),
142        )
143    }
144
145    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
146        let pairs: u32 = (((rank.0 + 1) * rank.0) >> 1).max(1);
147        GLWETensorKey(GGLWE::alloc(n, base2k, k, Rank(pairs), rank, dnum, dsize))
148    }
149
150    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
151    where
152        A: GGLWEInfos,
153    {
154        assert_eq!(
155            infos.rank_in(),
156            infos.rank_out(),
157            "rank_in != rank_out is not supported for GGLWETensorKey"
158        );
159        Self::bytes_of(
160            infos.n(),
161            infos.base2k(),
162            infos.k(),
163            infos.rank(),
164            infos.dnum(),
165            infos.dsize(),
166        )
167    }
168
169    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
170        let pairs: u32 = (((rank.0 + 1) * rank.0) >> 1).max(1);
171        GGLWE::bytes_of(n, base2k, k, Rank(pairs), rank, dnum, dsize)
172    }
173}
174
175impl<D: DataMut> ReaderFrom for GLWETensorKey<D> {
176    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
177        self.0.read_from(reader)?;
178        Ok(())
179    }
180}
181
182impl<D: DataRef> WriterTo for GLWETensorKey<D> {
183    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
184        self.0.write_to(writer)?;
185        Ok(())
186    }
187}
188
189impl<D: DataRef> GGLWEToRef for GLWETensorKey<D>
190where
191    GGLWE<D>: GGLWEToRef,
192{
193    fn to_ref(&self) -> GGLWE<&[u8]> {
194        self.0.to_ref()
195    }
196}
197
198impl<D: DataMut> GGLWEToMut for GLWETensorKey<D>
199where
200    GGLWE<D>: GGLWEToMut,
201{
202    fn to_mut(&mut self) -> GGLWE<&mut [u8]> {
203        self.0.to_mut()
204    }
205}