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