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