1use std::marker::PhantomData;
  2
  3use crate::ffi::vec_znx_dft;
  4use crate::znx_base::ZnxInfos;
  5use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxView, alloc_aligned};
  6
  7const VEC_ZNX_DFT_ROWS: usize = 1;
  8
  9// VecZnxDft is `Backend` dependent denoted with generic `B`
 10pub struct VecZnxDft<D, B> {
 11    data: D,
 12    n: usize,
 13    cols: usize,
 14    size: usize,
 15    _phantom: PhantomData<B>,
 16}
 17
 18impl<D, B> ZnxInfos for VecZnxDft<D, B> {
 19    fn cols(&self) -> usize {
 20        self.cols
 21    }
 22
 23    fn rows(&self) -> usize {
 24        1
 25    }
 26
 27    fn n(&self) -> usize {
 28        self.n
 29    }
 30
 31    fn size(&self) -> usize {
 32        self.size
 33    }
 34
 35    fn sl(&self) -> usize {
 36        self.cols() * self.n()
 37    }
 38}
 39
 40impl<D, B> DataView for VecZnxDft<D, B> {
 41    type D = D;
 42    fn data(&self) -> &Self::D {
 43        &self.data
 44    }
 45}
 46
 47impl<D, B> 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> 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
100// impl<B: Backend> ZnxAlloc<B> for VecZnxDft<B> {
101//     type Scalar = u8;
102
103//     fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> Self {
104//         debug_assert_eq!(bytes.len(), Self::bytes_of(module, _rows, cols, size));
105//         Self {
106//             inner: ZnxBase::from_bytes_borrow(module.n(), VEC_ZNX_DFT_ROWS, cols, size, bytes),
107//             _marker: PhantomData,
108//         }
109//     }
110
111//     fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, size: usize) -> usize {
112//         debug_assert_eq!(
113//             _rows, VEC_ZNX_DFT_ROWS,
114//             "rows != {} not supported for VecZnxDft",
115//             VEC_ZNX_DFT_ROWS
116//         );
117//         unsafe { vec_znx_dft::bytes_of_vec_znx_dft(module.ptr, size as u64) as usize * cols }
118//     }
119// }
120
121// impl VecZnxDft<FFT64> {
122//     pub fn print(&self, n: usize, col: usize) {
123//         (0..self.size()).for_each(|i| println!("{}: {:?}", i, &self.at(col, i)[..n]));
124//     }
125// }
126
127// impl<B: Backend> VecZnxDft<B> {
128//     /// Cast a [VecZnxDft] into a [VecZnxBig].
129//     /// The returned [VecZnxBig] shares the backing array
130//     /// with the original [VecZnxDft].
131//     pub fn alias_as_vec_znx_big(&mut self) -> VecZnxBig<B> {
132//         assert!(
133//             self.data().len() == 0,
134//             "cannot alias VecZnxDft into VecZnxBig if it owns the data"
135//         );
136//         VecZnxBig::<B> {
137//             inner: ZnxBase {
138//                 data: Vec::new(),
139//                 ptr: self.ptr(),
140//                 n: self.n(),
141//                 rows: self.rows(),
142//                 cols: self.cols(),
143//                 size: self.size(),
144//             },
145//             _marker: PhantomData,
146//         }
147//     }
148// }