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