1use std::{
  2    hash::{DefaultHasher, Hasher},
  3    marker::PhantomData,
  4};
  5
  6use rand_distr::num_traits::Zero;
  7use std::fmt;
  8
  9use crate::{
 10    alloc_aligned,
 11    layouts::{
 12        Backend, Data, DataMut, DataRef, DataView, DataViewMut, DigestU64, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero,
 13    },
 14    oep::VecZnxBigAllocBytesImpl,
 15};
 16
 17#[repr(C)]
 18#[derive(PartialEq, Eq, Hash)]
 19pub struct VecZnxBig<D: Data, B: Backend> {
 20    pub data: D,
 21    pub n: usize,
 22    pub cols: usize,
 23    pub size: usize,
 24    pub max_size: usize,
 25    pub _phantom: PhantomData<B>,
 26}
 27
 28impl<D: DataRef, B: Backend> DigestU64 for VecZnxBig<D, B> {
 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: Data, B: Backend> ZnxSliceSize for VecZnxBig<D, B> {
 41    fn sl(&self) -> usize {
 42        B::layout_big_word_count() * self.n() * self.cols()
 43    }
 44}
 45
 46impl<D: DataRef, B: Backend> ZnxView for VecZnxBig<D, B> {
 47    type Scalar = B::ScalarBig;
 48}
 49
 50impl<D: Data, B: Backend> ZnxInfos for VecZnxBig<D, B> {
 51    fn cols(&self) -> usize {
 52        self.cols
 53    }
 54
 55    fn rows(&self) -> usize {
 56        1
 57    }
 58
 59    fn n(&self) -> usize {
 60        self.n
 61    }
 62
 63    fn size(&self) -> usize {
 64        self.size
 65    }
 66}
 67
 68impl<D: Data, B: Backend> DataView for VecZnxBig<D, B> {
 69    type D = D;
 70    fn data(&self) -> &Self::D {
 71        &self.data
 72    }
 73}
 74
 75impl<D: Data, B: Backend> DataViewMut for VecZnxBig<D, B> {
 76    fn data_mut(&mut self) -> &mut Self::D {
 77        &mut self.data
 78    }
 79}
 80
 81impl<D: DataMut, B: Backend> ZnxZero for VecZnxBig<D, B>
 82where
 83    Self: ZnxViewMut,
 84    <Self as ZnxView>::Scalar: Zero + Copy,
 85{
 86    fn zero(&mut self) {
 87        self.raw_mut().fill(<Self as ZnxView>::Scalar::zero())
 88    }
 89    fn zero_at(&mut self, i: usize, j: usize) {
 90        self.at_mut(i, j).fill(<Self as ZnxView>::Scalar::zero());
 91    }
 92}
 93
 94impl<D: DataRef + From<Vec<u8>>, B: Backend> VecZnxBig<D, B>
 95where
 96    B: VecZnxBigAllocBytesImpl<B>,
 97{
 98    pub fn alloc(n: usize, cols: usize, size: usize) -> Self {
 99        let data = alloc_aligned::<u8>(B::vec_znx_big_alloc_bytes_impl(n, cols, size));
100        Self {
101            data: data.into(),
102            n,
103            cols,
104            size,
105            max_size: size,
106            _phantom: PhantomData,
107        }
108    }
109
110    pub fn from_bytes(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
111        let data: Vec<u8> = bytes.into();
112        assert!(data.len() == B::vec_znx_big_alloc_bytes_impl(n, cols, size));
113        Self {
114            data: data.into(),
115            n,
116            cols,
117            size,
118            max_size: size,
119            _phantom: PhantomData,
120        }
121    }
122}
123
124impl<D: Data, B: Backend> VecZnxBig<D, B> {
125    pub fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
126        Self {
127            data,
128            n,
129            cols,
130            size,
131            max_size: size,
132            _phantom: PhantomData,
133        }
134    }
135}
136
137pub type VecZnxBigOwned<B> = VecZnxBig<Vec<u8>, B>;
138
139pub trait VecZnxBigToRef<B: Backend> {
140    fn to_ref(&self) -> VecZnxBig<&[u8], B>;
141}
142
143impl<D: DataRef, B: Backend> VecZnxBigToRef<B> for VecZnxBig<D, B> {
144    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
145        VecZnxBig {
146            data: self.data.as_ref(),
147            n: self.n,
148            cols: self.cols,
149            size: self.size,
150            max_size: self.max_size,
151            _phantom: std::marker::PhantomData,
152        }
153    }
154}
155
156pub trait VecZnxBigToMut<B: Backend> {
157    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B>;
158}
159
160impl<D: DataMut, B: Backend> VecZnxBigToMut<B> for VecZnxBig<D, B> {
161    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
162        VecZnxBig {
163            data: self.data.as_mut(),
164            n: self.n,
165            cols: self.cols,
166            size: self.size,
167            max_size: self.max_size,
168            _phantom: std::marker::PhantomData,
169        }
170    }
171}
172
173impl<D: DataRef, B: Backend> fmt::Display for VecZnxBig<D, B> {
174    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
175        writeln!(
176            f,
177            "VecZnxBig(n={}, cols={}, size={})",
178            self.n, self.cols, self.size
179        )?;
180
181        for col in 0..self.cols {
182            writeln!(f, "Column {col}:")?;
183            for size in 0..self.size {
184                let coeffs = self.at(col, size);
185                write!(f, "  Size {size}: [")?;
186
187                let max_show = 100;
188                let show_count = coeffs.len().min(max_show);
189
190                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
191                    if i > 0 {
192                        write!(f, ", ")?;
193                    }
194                    write!(f, "{coeff}")?;
195                }
196
197                if coeffs.len() > max_show {
198                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
199                }
200
201                writeln!(f, "]")?;
202            }
203        }
204        Ok(())
205    }
206}