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