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