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