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        let rank_out: usize = self.rank_out().as_usize();
 52        let pairs: usize = (((rank_out + 1) * rank_out) >> 1).max(1);
 53        pairs.into()
 54    }
 55
 56    fn rank_out(&self) -> Rank {
 57        self.0.rank_out()
 58    }
 59
 60    fn dsize(&self) -> Dsize {
 61        self.0.dsize()
 62    }
 63
 64    fn dnum(&self) -> Dnum {
 65        self.0.dnum()
 66    }
 67}
 68
 69impl LWEInfos for GLWETensorKeyLayout {
 70    fn n(&self) -> Degree {
 71        self.n
 72    }
 73
 74    fn base2k(&self) -> Base2K {
 75        self.base2k
 76    }
 77
 78    fn k(&self) -> TorusPrecision {
 79        self.k
 80    }
 81}
 82
 83impl GLWEInfos for GLWETensorKeyLayout {
 84    fn rank(&self) -> Rank {
 85        self.rank_out()
 86    }
 87}
 88
 89impl GGLWEInfos for GLWETensorKeyLayout {
 90    fn rank_in(&self) -> Rank {
 91        let rank_out: usize = self.rank_out().as_usize();
 92        let pairs: usize = (((rank_out + 1) * rank_out) >> 1).max(1);
 93        pairs.into()
 94    }
 95
 96    fn dsize(&self) -> Dsize {
 97        self.dsize
 98    }
 99
100    fn rank_out(&self) -> Rank {
101        self.rank
102    }
103
104    fn dnum(&self) -> Dnum {
105        self.dnum
106    }
107}
108
109impl<D: DataRef> fmt::Debug for GLWETensorKey<D> {
110    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
111        write!(f, "{self}")
112    }
113}
114
115impl<D: DataMut> FillUniform for GLWETensorKey<D> {
116    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
117        self.0.fill_uniform(log_bound, source)
118    }
119}
120
121impl<D: DataRef> fmt::Display for GLWETensorKey<D> {
122    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123        writeln!(f, "(GLWETensorKey)",)?;
124        write!(f, "{}", self.0)?;
125        Ok(())
126    }
127}
128
129impl GLWETensorKey<Vec<u8>> {
130    pub fn alloc_from_infos<A>(infos: &A) -> Self
131    where
132        A: GGLWEInfos,
133    {
134        Self::alloc(
135            infos.n(),
136            infos.base2k(),
137            infos.k(),
138            infos.rank(),
139            infos.dnum(),
140            infos.dsize(),
141        )
142    }
143
144    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
145        let pairs: u32 = (((rank.0 + 1) * rank.0) >> 1).max(1);
146        GLWETensorKey(GGLWE::alloc(n, base2k, k, Rank(pairs), rank, dnum, dsize))
147    }
148
149    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
150    where
151        A: GGLWEInfos,
152    {
153        Self::bytes_of(
154            infos.n(),
155            infos.base2k(),
156            infos.k(),
157            infos.rank(),
158            infos.dnum(),
159            infos.dsize(),
160        )
161    }
162
163    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
164        let pairs: u32 = (((rank.0 + 1) * rank.0) >> 1).max(1);
165        GGLWE::bytes_of(n, base2k, k, Rank(pairs), rank, dnum, dsize)
166    }
167}
168
169impl<D: DataMut> ReaderFrom for GLWETensorKey<D> {
170    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
171        self.0.read_from(reader)?;
172        Ok(())
173    }
174}
175
176impl<D: DataRef> WriterTo for GLWETensorKey<D> {
177    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
178        self.0.write_to(writer)?;
179        Ok(())
180    }
181}
182
183impl<D: DataRef> GGLWEToRef for GLWETensorKey<D>
184where
185    GGLWE<D>: GGLWEToRef,
186{
187    fn to_ref(&self) -> GGLWE<&[u8]> {
188        self.0.to_ref()
189    }
190}
191
192impl<D: DataMut> GGLWEToMut for GLWETensorKey<D>
193where
194    GGLWE<D>: GGLWEToMut,
195{
196    fn to_mut(&mut self) -> GGLWE<&mut [u8]> {
197        self.0.to_mut()
198    }
199}