1use std::fmt;
  2
  3use crate::{
  4    alloc_aligned,
  5    hal::{
  6        api::{DataView, DataViewMut, FillUniform, Reset, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero},
  7        layouts::{Data, DataMut, DataRef, ReaderFrom, ToOwnedDeep, WriterTo},
  8    },
  9};
 10
 11#[derive(PartialEq, Eq, Clone, Copy)]
 12pub struct VecZnx<D: Data> {
 13    pub(crate) data: D,
 14    pub(crate) n: usize,
 15    pub(crate) cols: usize,
 16    pub(crate) size: usize,
 17    pub(crate) max_size: usize,
 18}
 19
 20impl<D: DataRef> ToOwnedDeep for VecZnx<D> {
 21    type Owned = VecZnx<Vec<u8>>;
 22    fn to_owned_deep(&self) -> Self::Owned {
 23        VecZnx {
 24            data: self.data.as_ref().to_vec(),
 25            n: self.n,
 26            cols: self.cols,
 27            size: self.size,
 28            max_size: self.max_size,
 29        }
 30    }
 31}
 32
 33impl<D: DataRef> fmt::Debug for VecZnx<D> {
 34    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 35        write!(f, "{}", self)
 36    }
 37}
 38
 39impl<D: Data> ZnxInfos for VecZnx<D> {
 40    fn cols(&self) -> usize {
 41        self.cols
 42    }
 43
 44    fn rows(&self) -> usize {
 45        1
 46    }
 47
 48    fn n(&self) -> usize {
 49        self.n
 50    }
 51
 52    fn size(&self) -> usize {
 53        self.size
 54    }
 55}
 56
 57impl<D: Data> ZnxSliceSize for VecZnx<D> {
 58    fn sl(&self) -> usize {
 59        self.n() * self.cols()
 60    }
 61}
 62
 63impl<D: Data> DataView for VecZnx<D> {
 64    type D = D;
 65    fn data(&self) -> &Self::D {
 66        &self.data
 67    }
 68}
 69
 70impl<D: Data> DataViewMut for VecZnx<D> {
 71    fn data_mut(&mut self) -> &mut Self::D {
 72        &mut self.data
 73    }
 74}
 75
 76impl<D: DataRef> ZnxView for VecZnx<D> {
 77    type Scalar = i64;
 78}
 79
 80impl VecZnx<Vec<u8>> {
 81    pub fn rsh_scratch_space(n: usize) -> usize {
 82        n * std::mem::size_of::<i64>()
 83    }
 84}
 85
 86impl<D: DataMut> ZnxZero for VecZnx<D> {
 87    fn zero(&mut self) {
 88        self.raw_mut().fill(0)
 89    }
 90    fn zero_at(&mut self, i: usize, j: usize) {
 91        self.at_mut(i, j).fill(0);
 92    }
 93}
 94
 95impl VecZnx<Vec<u8>> {
 96    pub fn alloc_bytes(n: usize, cols: usize, size: usize) -> usize {
 97        n * cols * size * size_of::<i64>()
 98    }
 99
100    pub fn alloc(n: usize, cols: usize, size: usize) -> Self {
101        let data: Vec<u8> = alloc_aligned::<u8>(Self::alloc_bytes(n, cols, size));
102        Self {
103            data,
104            n,
105            cols,
106            size,
107            max_size: size,
108        }
109    }
110
111    pub fn from_bytes<Scalar: Sized>(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
112        let data: Vec<u8> = bytes.into();
113        assert!(data.len() == Self::alloc_bytes(n, cols, size));
114        Self {
115            data,
116            n,
117            cols,
118            size,
119            max_size: size,
120        }
121    }
122}
123
124impl<D: Data> VecZnx<D> {
125    pub fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
126        Self {
127            data,
128            n,
129            cols,
130            size,
131            max_size: size,
132        }
133    }
134}
135
136impl<D: DataRef> fmt::Display for VecZnx<D> {
137    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
138        writeln!(
139            f,
140            "VecZnx(n={}, cols={}, size={})",
141            self.n, self.cols, self.size
142        )?;
143
144        for col in 0..self.cols {
145            writeln!(f, "Column {}:", col)?;
146            for size in 0..self.size {
147                let coeffs = self.at(col, size);
148                write!(f, "  Size {}: [", size)?;
149
150                let max_show = 100;
151                let show_count = coeffs.len().min(max_show);
152
153                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
154                    if i > 0 {
155                        write!(f, ", ")?;
156                    }
157                    write!(f, "{}", coeff)?;
158                }
159
160                if coeffs.len() > max_show {
161                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
162                }
163
164                writeln!(f, "]")?;
165            }
166        }
167        Ok(())
168    }
169}
170
171impl<D: DataMut> FillUniform for VecZnx<D> {
172    fn fill_uniform(&mut self, source: &mut Source) {
173        source.fill_bytes(self.data.as_mut());
174    }
175}
176
177impl<D: DataMut> Reset for VecZnx<D> {
178    fn reset(&mut self) {
179        self.zero();
180        self.n = 0;
181        self.cols = 0;
182        self.size = 0;
183        self.max_size = 0;
184    }
185}
186
187pub type VecZnxOwned = VecZnx<Vec<u8>>;
188pub type VecZnxMut<'a> = VecZnx<&'a mut [u8]>;
189pub type VecZnxRef<'a> = VecZnx<&'a [u8]>;
190
191pub trait VecZnxToRef {
192    fn to_ref(&self) -> VecZnx<&[u8]>;
193}
194
195impl<D: DataRef> VecZnxToRef for VecZnx<D> {
196    fn to_ref(&self) -> VecZnx<&[u8]> {
197        VecZnx {
198            data: self.data.as_ref(),
199            n: self.n,
200            cols: self.cols,
201            size: self.size,
202            max_size: self.max_size,
203        }
204    }
205}
206
207pub trait VecZnxToMut {
208    fn to_mut(&mut self) -> VecZnx<&mut [u8]>;
209}
210
211impl<D: DataMut> VecZnxToMut for VecZnx<D> {
212    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
213        VecZnx {
214            data: self.data.as_mut(),
215            n: self.n,
216            cols: self.cols,
217            size: self.size,
218            max_size: self.max_size,
219        }
220    }
221}
222
223use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
224use rand::RngCore;
225use sampling::source::Source;
226
227impl<D: DataMut> ReaderFrom for VecZnx<D> {
228    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
229        self.n = reader.read_u64::<LittleEndian>()? as usize;
230        self.cols = reader.read_u64::<LittleEndian>()? as usize;
231        self.size = reader.read_u64::<LittleEndian>()? as usize;
232        self.max_size = reader.read_u64::<LittleEndian>()? as usize;
233        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
234        let buf: &mut [u8] = self.data.as_mut();
235        if buf.len() != len {
236            return Err(std::io::Error::new(
237                std::io::ErrorKind::UnexpectedEof,
238                format!("self.data.len()={} != read len={}", buf.len(), len),
239            ));
240        }
241        reader.read_exact(&mut buf[..len])?;
242        Ok(())
243    }
244}
245
246impl<D: DataRef> WriterTo for VecZnx<D> {
247    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
248        writer.write_u64::<LittleEndian>(self.n as u64)?;
249        writer.write_u64::<LittleEndian>(self.cols as u64)?;
250        writer.write_u64::<LittleEndian>(self.size as u64)?;
251        writer.write_u64::<LittleEndian>(self.max_size as u64)?;
252        let buf: &[u8] = self.data.as_ref();
253        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
254        writer.write_all(buf)?;
255        Ok(())
256    }
257}