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