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
  7// const 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
100impl<B> VecZnxDft<Vec<u8>, B> {
101    pub fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
102        VecZnxDft {
103            data: self.data.as_mut_slice(),
104            n: self.n,
105            cols: self.cols,
106            size: self.size,
107            _phantom: PhantomData,
108        }
109    }
110
111    pub fn to_ref(&self) -> VecZnxDft<&[u8], B> {
112        VecZnxDft {
113            data: self.data.as_slice(),
114            n: self.n,
115            cols: self.cols,
116            size: self.size,
117            _phantom: PhantomData,
118        }
119    }
120}
121
122impl<B> VecZnxDft<&mut [u8], B> {
123    pub fn to_ref(&self) -> VecZnxDft<&[u8], B> {
124        VecZnxDft {
125            data: &self.data,
126            n: self.n,
127            cols: self.cols,
128            size: self.size,
129            _phantom: PhantomData,
130        }
131    }
132}