1use std::{
  2    fmt,
  3    hash::{DefaultHasher, Hasher},
  4    marker::PhantomData,
  5};
  6
  7use crate::{
  8    alloc_aligned,
  9    layouts::{
 10        Backend, Data, DataMut, DataRef, DataView, DataViewMut, DigestU64, ReaderFrom, WriterTo, ZnxInfos, ZnxSliceSize, ZnxView,
 11    },
 12    oep::SvpPPolAllocBytesImpl,
 13};
 14
 15#[repr(C)]
 16#[derive(PartialEq, Eq, Hash)]
 17pub struct SvpPPol<D: Data, B: Backend> {
 18    pub data: D,
 19    pub n: usize,
 20    pub cols: usize,
 21    pub _phantom: PhantomData<B>,
 22}
 23
 24impl<D: DataRef, B: Backend> DigestU64 for SvpPPol<D, B> {
 25    fn digest_u64(&self) -> u64 {
 26        let mut h: DefaultHasher = DefaultHasher::new();
 27        h.write(self.data.as_ref());
 28        h.write_usize(self.n);
 29        h.write_usize(self.cols);
 30        h.finish()
 31    }
 32}
 33
 34impl<D: Data, B: Backend> ZnxSliceSize for SvpPPol<D, B> {
 35    fn sl(&self) -> usize {
 36        B::layout_prep_word_count() * self.n()
 37    }
 38}
 39
 40impl<D: DataRef, B: Backend> ZnxView for SvpPPol<D, B> {
 41    type Scalar = B::ScalarPrep;
 42}
 43
 44impl<D: Data, B: Backend> ZnxInfos for SvpPPol<D, B> {
 45    fn cols(&self) -> usize {
 46        self.cols
 47    }
 48
 49    fn rows(&self) -> usize {
 50        1
 51    }
 52
 53    fn n(&self) -> usize {
 54        self.n
 55    }
 56
 57    fn size(&self) -> usize {
 58        1
 59    }
 60}
 61
 62impl<D: Data, B: Backend> DataView for SvpPPol<D, B> {
 63    type D = D;
 64    fn data(&self) -> &Self::D {
 65        &self.data
 66    }
 67}
 68
 69impl<D: Data, B: Backend> DataViewMut for SvpPPol<D, B> {
 70    fn data_mut(&mut self) -> &mut Self::D {
 71        &mut self.data
 72    }
 73}
 74
 75impl<D: Data + From<Vec<u8>>, B: Backend> SvpPPol<D, B>
 76where
 77    B: SvpPPolAllocBytesImpl<B>,
 78{
 79    pub fn alloc(n: usize, cols: usize) -> Self {
 80        let data: Vec<u8> = alloc_aligned::<u8>(B::svp_ppol_alloc_bytes_impl(n, cols));
 81        Self {
 82            data: data.into(),
 83            n,
 84            cols,
 85            _phantom: PhantomData,
 86        }
 87    }
 88
 89    pub fn from_bytes(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
 90        let data: Vec<u8> = bytes.into();
 91        assert!(data.len() == B::svp_ppol_alloc_bytes_impl(n, cols));
 92        Self {
 93            data: data.into(),
 94            n,
 95            cols,
 96            _phantom: PhantomData,
 97        }
 98    }
 99}
100
101pub type SvpPPolOwned<B> = SvpPPol<Vec<u8>, B>;
102
103pub trait SvpPPolToRef<B: Backend> {
104    fn to_ref(&self) -> SvpPPol<&[u8], B>;
105}
106
107impl<D: DataRef, B: Backend> SvpPPolToRef<B> for SvpPPol<D, B> {
108    fn to_ref(&self) -> SvpPPol<&[u8], B> {
109        SvpPPol {
110            data: self.data.as_ref(),
111            n: self.n,
112            cols: self.cols,
113            _phantom: PhantomData,
114        }
115    }
116}
117
118pub trait SvpPPolToMut<B: Backend> {
119    fn to_mut(&mut self) -> SvpPPol<&mut [u8], B>;
120}
121
122impl<D: DataMut, B: Backend> SvpPPolToMut<B> for SvpPPol<D, B> {
123    fn to_mut(&mut self) -> SvpPPol<&mut [u8], B> {
124        SvpPPol {
125            data: self.data.as_mut(),
126            n: self.n,
127            cols: self.cols,
128            _phantom: PhantomData,
129        }
130    }
131}
132
133impl<D: Data, B: Backend> SvpPPol<D, B> {
134    pub fn from_data(data: D, n: usize, cols: usize) -> Self {
135        Self {
136            data,
137            n,
138            cols,
139            _phantom: PhantomData,
140        }
141    }
142}
143
144use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
145
146impl<D: DataMut, B: Backend> ReaderFrom for SvpPPol<D, B> {
147    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
148        self.n = reader.read_u64::<LittleEndian>()? as usize;
149        self.cols = reader.read_u64::<LittleEndian>()? as usize;
150        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
151        let buf: &mut [u8] = self.data.as_mut();
152        if buf.len() != len {
153            return Err(std::io::Error::new(
154                std::io::ErrorKind::UnexpectedEof,
155                format!("self.data.len()={} != read len={}", buf.len(), len),
156            ));
157        }
158        reader.read_exact(&mut buf[..len])?;
159        Ok(())
160    }
161}
162
163impl<D: DataRef, B: Backend> WriterTo for SvpPPol<D, B> {
164    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
165        writer.write_u64::<LittleEndian>(self.n as u64)?;
166        writer.write_u64::<LittleEndian>(self.cols as u64)?;
167        let buf: &[u8] = self.data.as_ref();
168        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
169        writer.write_all(buf)?;
170        Ok(())
171    }
172}
173
174impl<D: DataRef, B: Backend> fmt::Display for SvpPPol<D, B> {
175    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
176        writeln!(f, "SvpPPol(n={}, cols={})", self.n, self.cols)?;
177
178        for col in 0..self.cols {
179            writeln!(f, "Column {}:", col)?;
180            let coeffs = self.at(col, 0);
181            write!(f, "[")?;
182
183            let max_show = 100;
184            let show_count = coeffs.len().min(max_show);
185
186            for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
187                if i > 0 {
188                    write!(f, ", ")?;
189                }
190                write!(f, "{}", coeff)?;
191            }
192
193            if coeffs.len() > max_show {
194                write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
195            }
196
197            writeln!(f, "]")?;
198        }
199        Ok(())
200    }
201}