1use std::marker::PhantomData;
  2
  3use crate::{
  4    alloc_aligned,
  5    api::{DataView, DataViewMut, ZnxInfos, ZnxSliceSize, ZnxView},
  6    layouts::{Backend, Data, DataMut, DataRef, ReaderFrom, WriterTo},
  7    oep::SvpPPolAllocBytesImpl,
  8};
  9
 10#[derive(PartialEq, Eq)]
 11pub struct SvpPPol<D: Data, B: Backend> {
 12    pub data: D,
 13    pub n: usize,
 14    pub cols: usize,
 15    pub _phantom: PhantomData<B>,
 16}
 17
 18impl<D: Data, B: Backend> ZnxSliceSize for SvpPPol<D, B> {
 19    fn sl(&self) -> usize {
 20        B::layout_prep_word_count() * self.n()
 21    }
 22}
 23
 24impl<D: DataRef, B: Backend> ZnxView for SvpPPol<D, B> {
 25    type Scalar = B::ScalarPrep;
 26}
 27
 28impl<D: Data, B: Backend> ZnxInfos for SvpPPol<D, B> {
 29    fn cols(&self) -> usize {
 30        self.cols
 31    }
 32
 33    fn rows(&self) -> usize {
 34        1
 35    }
 36
 37    fn n(&self) -> usize {
 38        self.n
 39    }
 40
 41    fn size(&self) -> usize {
 42        1
 43    }
 44}
 45
 46impl<D: Data, B: Backend> DataView for SvpPPol<D, B> {
 47    type D = D;
 48    fn data(&self) -> &Self::D {
 49        &self.data
 50    }
 51}
 52
 53impl<D: Data, B: Backend> DataViewMut for SvpPPol<D, B> {
 54    fn data_mut(&mut self) -> &mut Self::D {
 55        &mut self.data
 56    }
 57}
 58
 59impl<D: Data + From<Vec<u8>>, B: Backend> SvpPPol<D, B>
 60where
 61    B: SvpPPolAllocBytesImpl<B>,
 62{
 63    pub fn alloc(n: usize, cols: usize) -> Self {
 64        let data: Vec<u8> = alloc_aligned::<u8>(B::svp_ppol_alloc_bytes_impl(n, cols));
 65        Self {
 66            data: data.into(),
 67            n,
 68            cols,
 69            _phantom: PhantomData,
 70        }
 71    }
 72
 73    pub fn from_bytes(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
 74        let data: Vec<u8> = bytes.into();
 75        assert!(data.len() == B::svp_ppol_alloc_bytes_impl(n, cols));
 76        Self {
 77            data: data.into(),
 78            n,
 79            cols,
 80            _phantom: PhantomData,
 81        }
 82    }
 83}
 84
 85pub type SvpPPolOwned<B> = SvpPPol<Vec<u8>, B>;
 86
 87pub trait SvpPPolToRef<B: Backend> {
 88    fn to_ref(&self) -> SvpPPol<&[u8], B>;
 89}
 90
 91impl<D: DataRef, B: Backend> SvpPPolToRef<B> for SvpPPol<D, B> {
 92    fn to_ref(&self) -> SvpPPol<&[u8], B> {
 93        SvpPPol {
 94            data: self.data.as_ref(),
 95            n: self.n,
 96            cols: self.cols,
 97            _phantom: PhantomData,
 98        }
 99    }
100}
101
102pub trait SvpPPolToMut<B: Backend> {
103    fn to_mut(&mut self) -> SvpPPol<&mut [u8], B>;
104}
105
106impl<D: DataMut, B: Backend> SvpPPolToMut<B> for SvpPPol<D, B> {
107    fn to_mut(&mut self) -> SvpPPol<&mut [u8], B> {
108        SvpPPol {
109            data: self.data.as_mut(),
110            n: self.n,
111            cols: self.cols,
112            _phantom: PhantomData,
113        }
114    }
115}
116
117impl<D: Data, B: Backend> SvpPPol<D, B> {
118    pub fn from_data(data: D, n: usize, cols: usize) -> Self {
119        Self {
120            data,
121            n,
122            cols,
123            _phantom: PhantomData,
124        }
125    }
126}
127
128use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
129
130impl<D: DataMut, B: Backend> ReaderFrom for SvpPPol<D, B> {
131    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
132        self.n = reader.read_u64::<LittleEndian>()? as usize;
133        self.cols = reader.read_u64::<LittleEndian>()? as usize;
134        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
135        let buf: &mut [u8] = self.data.as_mut();
136        if buf.len() != len {
137            return Err(std::io::Error::new(
138                std::io::ErrorKind::UnexpectedEof,
139                format!("self.data.len()={} != read len={}", buf.len(), len),
140            ));
141        }
142        reader.read_exact(&mut buf[..len])?;
143        Ok(())
144    }
145}
146
147impl<D: DataRef, B: Backend> WriterTo for SvpPPol<D, B> {
148    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
149        writer.write_u64::<LittleEndian>(self.n as u64)?;
150        writer.write_u64::<LittleEndian>(self.cols as u64)?;
151        let buf: &[u8] = self.data.as_ref();
152        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
153        writer.write_all(buf)?;
154        Ok(())
155    }
156}