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