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!(
191            f,
192            "VecZnx(n={}, cols={}, size={})",
193            self.n, self.cols, self.size
194        )?;
195
196        for col in 0..self.cols {
197            writeln!(f, "Column {col}:")?;
198            for size in 0..self.size {
199                let coeffs = self.at(col, size);
200                write!(f, "  Size {size}: [")?;
201
202                let max_show = 100;
203                let show_count = coeffs.len().min(max_show);
204
205                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
206                    if i > 0 {
207                        write!(f, ", ")?;
208                    }
209                    write!(f, "{coeff}")?;
210                }
211
212                if coeffs.len() > max_show {
213                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
214                }
215
216                writeln!(f, "]")?;
217            }
218        }
219        Ok(())
220    }
221}
222
223impl<D: DataMut> FillUniform for VecZnx<D> {
224    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
225        match log_bound {
226            64 => source.fill_bytes(self.data.as_mut()),
227            0 => panic!("invalid log_bound, cannot be zero"),
228            _ => {
229                let mask: u64 = (1u64 << log_bound) - 1;
230                for x in self.raw_mut().iter_mut() {
231                    let r = source.next_u64() & mask;
232                    *x = ((r << (64 - log_bound)) as i64) >> (64 - log_bound);
233                }
234            }
235        }
236    }
237}
238
239pub type VecZnxOwned = VecZnx<Vec<u8>>;
240pub type VecZnxMut<'a> = VecZnx<&'a mut [u8]>;
241pub type VecZnxRef<'a> = VecZnx<&'a [u8]>;
242
243pub trait VecZnxToRef {
244    fn to_ref(&self) -> VecZnx<&[u8]>;
245}
246
247impl<D: DataRef> VecZnxToRef for VecZnx<D> {
248    fn to_ref(&self) -> VecZnx<&[u8]> {
249        VecZnx {
250            data: self.data.as_ref(),
251            n: self.n,
252            cols: self.cols,
253            size: self.size,
254            max_size: self.max_size,
255        }
256    }
257}
258
259pub trait VecZnxToMut {
260    fn to_mut(&mut self) -> VecZnx<&mut [u8]>;
261}
262
263impl<D: DataMut> VecZnxToMut for VecZnx<D> {
264    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
265        VecZnx {
266            data: self.data.as_mut(),
267            n: self.n,
268            cols: self.cols,
269            size: self.size,
270            max_size: self.max_size,
271        }
272    }
273}
274
275impl<D: DataMut> ReaderFrom for VecZnx<D> {
276    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
277        self.n = reader.read_u64::<LittleEndian>()? as usize;
278        self.cols = reader.read_u64::<LittleEndian>()? as usize;
279        self.size = reader.read_u64::<LittleEndian>()? as usize;
280        self.max_size = reader.read_u64::<LittleEndian>()? as usize;
281        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
282        let buf: &mut [u8] = self.data.as_mut();
283        if buf.len() != len {
284            return Err(std::io::Error::new(
285                std::io::ErrorKind::UnexpectedEof,
286                format!("self.data.len()={} != read len={}", buf.len(), len),
287            ));
288        }
289        reader.read_exact(&mut buf[..len])?;
290        Ok(())
291    }
292}
293
294impl<D: DataRef> WriterTo for VecZnx<D> {
295    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
296        writer.write_u64::<LittleEndian>(self.n as u64)?;
297        writer.write_u64::<LittleEndian>(self.cols as u64)?;
298        writer.write_u64::<LittleEndian>(self.size as u64)?;
299        writer.write_u64::<LittleEndian>(self.max_size as u64)?;
300        let buf: &[u8] = self.data.as_ref();
301        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
302        writer.write_all(buf)?;
303        Ok(())
304    }
305}