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