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, ToOwnedDeep, WriterTo, ZnxInfos,
 10        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
207pub type ZnOwned = Zn<Vec<u8>>;
208pub type ZnMut<'a> = Zn<&'a mut [u8]>;
209pub type ZnRef<'a> = Zn<&'a [u8]>;
210
211pub trait ZnToRef {
212    fn to_ref(&self) -> Zn<&[u8]>;
213}
214
215impl<D: DataRef> ZnToRef for Zn<D> {
216    fn to_ref(&self) -> Zn<&[u8]> {
217        Zn {
218            data: self.data.as_ref(),
219            n: self.n,
220            cols: self.cols,
221            size: self.size,
222            max_size: self.max_size,
223        }
224    }
225}
226
227pub trait ZnToMut {
228    fn to_mut(&mut self) -> Zn<&mut [u8]>;
229}
230
231impl<D: DataMut> ZnToMut for Zn<D> {
232    fn to_mut(&mut self) -> Zn<&mut [u8]> {
233        Zn {
234            data: self.data.as_mut(),
235            n: self.n,
236            cols: self.cols,
237            size: self.size,
238            max_size: self.max_size,
239        }
240    }
241}
242
243impl<D: DataMut> ReaderFrom for Zn<D> {
244    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
245        self.n = reader.read_u64::<LittleEndian>()? as usize;
246        self.cols = reader.read_u64::<LittleEndian>()? as usize;
247        self.size = reader.read_u64::<LittleEndian>()? as usize;
248        self.max_size = reader.read_u64::<LittleEndian>()? as usize;
249        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
250        let buf: &mut [u8] = self.data.as_mut();
251        if buf.len() != len {
252            return Err(std::io::Error::new(
253                std::io::ErrorKind::UnexpectedEof,
254                format!("self.data.len()={} != read len={}", buf.len(), len),
255            ));
256        }
257        reader.read_exact(&mut buf[..len])?;
258        Ok(())
259    }
260}
261
262impl<D: DataRef> WriterTo for Zn<D> {
263    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
264        writer.write_u64::<LittleEndian>(self.n as u64)?;
265        writer.write_u64::<LittleEndian>(self.cols as u64)?;
266        writer.write_u64::<LittleEndian>(self.size as u64)?;
267        writer.write_u64::<LittleEndian>(self.max_size as u64)?;
268        let buf: &[u8] = self.data.as_ref();
269        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
270        writer.write_all(buf)?;
271        Ok(())
272    }
273}