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    data: D,
 10    n: usize,
 11    cols: usize,
 12    size: usize,
 13    _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
 63pub(crate) fn bytes_of_vec_znx_dft<B: Backend>(module: &Module<B>, cols: usize, size: usize) -> usize {
 64    unsafe { vec_znx_dft::bytes_of_vec_znx_dft(module.ptr, size as u64) as usize * cols }
 65}
 66
 67impl<D: From<Vec<u8>>, B: Backend> VecZnxDft<D, B> {
 68    pub(crate) fn new(module: &Module<B>, cols: usize, size: usize) -> Self {
 69        let data = alloc_aligned::<u8>(bytes_of_vec_znx_dft(module, cols, size));
 70        Self {
 71            data: data.into(),
 72            n: module.n(),
 73            cols,
 74            size,
 75            _phantom: PhantomData,
 76        }
 77    }
 78
 79    pub(crate) fn new_from_bytes(module: &Module<B>, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
 80        let data: Vec<u8> = bytes.into();
 81        assert!(data.len() == bytes_of_vec_znx_dft(module, cols, size));
 82        Self {
 83            data: data.into(),
 84            n: module.n(),
 85            cols,
 86            size,
 87            _phantom: PhantomData,
 88        }
 89    }
 90}
 91
 92pub type VecZnxDftOwned<B> = VecZnxDft<Vec<u8>, B>;
 93
 94impl<D, B: Backend> VecZnxDft<D, B> {
 95    pub(crate) fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
 96        Self {
 97            data,
 98            n,
 99            cols,
100            size,
101            _phantom: PhantomData,
102        }
103    }
104}
105
106pub trait VecZnxDftToRef<B: Backend> {
107    fn to_ref(&self) -> VecZnxDft<&[u8], B>;
108}
109
110pub trait VecZnxDftToMut<B: Backend> {
111    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B>;
112}
113
114impl<B: Backend> VecZnxDftToMut<B> for VecZnxDft<Vec<u8>, B> {
115    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
116        VecZnxDft {
117            data: self.data.as_mut_slice(),
118            n: self.n,
119            cols: self.cols,
120            size: self.size,
121            _phantom: PhantomData,
122        }
123    }
124}
125
126impl<B: Backend> VecZnxDftToRef<B> for VecZnxDft<Vec<u8>, B> {
127    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
128        VecZnxDft {
129            data: self.data.as_slice(),
130            n: self.n,
131            cols: self.cols,
132            size: self.size,
133            _phantom: PhantomData,
134        }
135    }
136}
137
138impl<B: Backend> VecZnxDftToMut<B> for VecZnxDft<&mut [u8], B> {
139    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
140        VecZnxDft {
141            data: self.data,
142            n: self.n,
143            cols: self.cols,
144            size: self.size,
145            _phantom: PhantomData,
146        }
147    }
148}
149
150impl<B: Backend> VecZnxDftToRef<B> for VecZnxDft<&mut [u8], B> {
151    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
152        VecZnxDft {
153            data: self.data,
154            n: self.n,
155            cols: self.cols,
156            size: self.size,
157            _phantom: PhantomData,
158        }
159    }
160}
161
162impl<B: Backend> VecZnxDftToRef<B> for VecZnxDft<&[u8], B> {
163    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
164        VecZnxDft {
165            data: self.data,
166            n: self.n,
167            cols: self.cols,
168            size: self.size,
169            _phantom: PhantomData,
170        }
171    }
172}
173
174impl<D: AsRef<[u8]>> fmt::Display for VecZnxDft<D, FFT64> {
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}