1use std::{fmt, marker::PhantomData};
  2
  3use rand_distr::num_traits::Zero;
  4
  5use crate::{
  6    alloc_aligned,
  7    api::{DataView, DataViewMut, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero},
  8    layouts::{Backend, Data, DataMut, DataRef, VecZnxBig},
  9    oep::VecZnxBigAllocBytesImpl,
 10};
 11#[derive(PartialEq, Eq)]
 12pub struct VecZnxDft<D: Data, B: Backend> {
 13    pub data: D,
 14    pub n: usize,
 15    pub cols: usize,
 16    pub size: usize,
 17    pub max_size: usize,
 18    pub _phantom: PhantomData<B>,
 19}
 20
 21impl<D: Data, B: Backend> ZnxSliceSize for VecZnxDft<D, B> {
 22    fn sl(&self) -> usize {
 23        B::layout_prep_word_count() * self.n() * self.cols()
 24    }
 25}
 26
 27impl<D: DataRef, B: Backend> ZnxView for VecZnxDft<D, B> {
 28    type Scalar = B::ScalarPrep;
 29}
 30
 31impl<D: Data, B: Backend> VecZnxDft<D, B> {
 32    pub fn into_big(self) -> VecZnxBig<D, B> {
 33        VecZnxBig::<D, B>::from_data(self.data, self.n, self.cols, self.size)
 34    }
 35}
 36
 37impl<D: Data, B: Backend> ZnxInfos for VecZnxDft<D, B> {
 38    fn cols(&self) -> usize {
 39        self.cols
 40    }
 41
 42    fn rows(&self) -> usize {
 43        1
 44    }
 45
 46    fn n(&self) -> usize {
 47        self.n
 48    }
 49
 50    fn size(&self) -> usize {
 51        self.size
 52    }
 53}
 54
 55impl<D: Data, B: Backend> DataView for VecZnxDft<D, B> {
 56    type D = D;
 57    fn data(&self) -> &Self::D {
 58        &self.data
 59    }
 60}
 61
 62impl<D: Data, B: Backend> DataViewMut for VecZnxDft<D, B> {
 63    fn data_mut(&mut self) -> &mut Self::D {
 64        &mut self.data
 65    }
 66}
 67
 68impl<D: DataRef, B: Backend> VecZnxDft<D, B> {
 69    pub fn max_size(&self) -> usize {
 70        self.max_size
 71    }
 72}
 73
 74impl<D: DataMut, B: Backend> VecZnxDft<D, B> {
 75    pub fn set_size(&mut self, size: usize) {
 76        assert!(size <= self.max_size);
 77        self.size = size
 78    }
 79}
 80
 81impl<D: DataMut, B: Backend> ZnxZero for VecZnxDft<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> VecZnxDft<D, B>
 95where
 96    B: VecZnxBigAllocBytesImpl<B>,
 97{
 98    pub fn alloc(n: usize, cols: usize, size: usize) -> Self {
 99        let data: Vec<u8> = 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
124pub type VecZnxDftOwned<B> = VecZnxDft<Vec<u8>, B>;
125
126impl<D: Data, B: Backend> VecZnxDft<D, B> {
127    pub fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
128        Self {
129            data,
130            n,
131            cols,
132            size,
133            max_size: size,
134            _phantom: PhantomData,
135        }
136    }
137}
138
139pub trait VecZnxDftToRef<B: Backend> {
140    fn to_ref(&self) -> VecZnxDft<&[u8], B>;
141}
142
143impl<D: DataRef, B: Backend> VecZnxDftToRef<B> for VecZnxDft<D, B> {
144    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
145        VecZnxDft {
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 VecZnxDftToMut<B: Backend> {
157    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B>;
158}
159
160impl<D: DataMut, B: Backend> VecZnxDftToMut<B> for VecZnxDft<D, B> {
161    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
162        VecZnxDft {
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 VecZnxDft<D, B> {
174    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
175        writeln!(
176            f,
177            "VecZnxDft(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}