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