1use std::marker::PhantomData;
  2
  3use rand_distr::num_traits::Zero;
  4use std::fmt;
  5
  6use crate::{
  7    alloc_aligned,
  8    api::{DataView, DataViewMut, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero},
  9    layouts::{Backend, Data, DataMut, DataRef},
 10    oep::VecZnxBigAllocBytesImpl,
 11};
 12
 13#[derive(PartialEq, Eq)]
 14pub struct VecZnxBig<D: Data, B: Backend> {
 15    pub data: D,
 16    pub n: usize,
 17    pub cols: usize,
 18    pub size: usize,
 19    pub max_size: usize,
 20    pub _phantom: PhantomData<B>,
 21}
 22
 23impl<D: Data, B: Backend> ZnxSliceSize for VecZnxBig<D, B> {
 24    fn sl(&self) -> usize {
 25        B::layout_big_word_count() * self.n() * self.cols()
 26    }
 27}
 28
 29impl<D: DataRef, B: Backend> ZnxView for VecZnxBig<D, B> {
 30    type Scalar = B::ScalarBig;
 31}
 32
 33impl<D: Data, B: Backend> ZnxInfos for VecZnxBig<D, B> {
 34    fn cols(&self) -> usize {
 35        self.cols
 36    }
 37
 38    fn rows(&self) -> usize {
 39        1
 40    }
 41
 42    fn n(&self) -> usize {
 43        self.n
 44    }
 45
 46    fn size(&self) -> usize {
 47        self.size
 48    }
 49}
 50
 51impl<D: Data, B: Backend> DataView for VecZnxBig<D, B> {
 52    type D = D;
 53    fn data(&self) -> &Self::D {
 54        &self.data
 55    }
 56}
 57
 58impl<D: Data, B: Backend> DataViewMut for VecZnxBig<D, B> {
 59    fn data_mut(&mut self) -> &mut Self::D {
 60        &mut self.data
 61    }
 62}
 63
 64impl<D: DataMut, B: Backend> ZnxZero for VecZnxBig<D, B>
 65where
 66    Self: ZnxViewMut,
 67    <Self as ZnxView>::Scalar: Zero + Copy,
 68{
 69    fn zero(&mut self) {
 70        self.raw_mut().fill(<Self as ZnxView>::Scalar::zero())
 71    }
 72    fn zero_at(&mut self, i: usize, j: usize) {
 73        self.at_mut(i, j).fill(<Self as ZnxView>::Scalar::zero());
 74    }
 75}
 76
 77impl<D: DataRef + From<Vec<u8>>, B: Backend> VecZnxBig<D, B>
 78where
 79    B: VecZnxBigAllocBytesImpl<B>,
 80{
 81    pub fn alloc(n: usize, cols: usize, size: usize) -> Self {
 82        let data = alloc_aligned::<u8>(B::vec_znx_big_alloc_bytes_impl(n, cols, size));
 83        Self {
 84            data: data.into(),
 85            n,
 86            cols,
 87            size,
 88            max_size: size,
 89            _phantom: PhantomData,
 90        }
 91    }
 92
 93    pub fn from_bytes(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
 94        let data: Vec<u8> = bytes.into();
 95        assert!(data.len() == B::vec_znx_big_alloc_bytes_impl(n, cols, size));
 96        Self {
 97            data: data.into(),
 98            n,
 99            cols,
100            size,
101            max_size: size,
102            _phantom: PhantomData,
103        }
104    }
105}
106
107impl<D: Data, B: Backend> VecZnxBig<D, B> {
108    pub fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
109        Self {
110            data,
111            n,
112            cols,
113            size,
114            max_size: size,
115            _phantom: PhantomData,
116        }
117    }
118}
119
120pub type VecZnxBigOwned<B> = VecZnxBig<Vec<u8>, B>;
121
122pub trait VecZnxBigToRef<B: Backend> {
123    fn to_ref(&self) -> VecZnxBig<&[u8], B>;
124}
125
126impl<D: DataRef, B: Backend> VecZnxBigToRef<B> for VecZnxBig<D, B> {
127    fn to_ref(&self) -> VecZnxBig<&[u8], B> {
128        VecZnxBig {
129            data: self.data.as_ref(),
130            n: self.n,
131            cols: self.cols,
132            size: self.size,
133            max_size: self.max_size,
134            _phantom: std::marker::PhantomData,
135        }
136    }
137}
138
139pub trait VecZnxBigToMut<B: Backend> {
140    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B>;
141}
142
143impl<D: DataMut, B: Backend> VecZnxBigToMut<B> for VecZnxBig<D, B> {
144    fn to_mut(&mut self) -> VecZnxBig<&mut [u8], B> {
145        VecZnxBig {
146            data: self.data.as_mut(),
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
156impl<D: DataRef, B: Backend> fmt::Display for VecZnxBig<D, B> {
157    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158        writeln!(
159            f,
160            "VecZnxBig(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}