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