1use std::marker::PhantomData;
  2
  3use crate::ffi::svp;
  4use crate::znx_base::ZnxInfos;
  5use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxSliceSize, ZnxView, alloc_aligned};
  6
  7pub struct ScalarZnxDft<D, B: Backend> {
  8    data: D,
  9    n: usize,
 10    cols: usize,
 11    _phantom: PhantomData<B>,
 12}
 13
 14impl<D, B: Backend> ZnxInfos for ScalarZnxDft<D, B> {
 15    fn cols(&self) -> usize {
 16        self.cols
 17    }
 18
 19    fn rows(&self) -> usize {
 20        1
 21    }
 22
 23    fn n(&self) -> usize {
 24        self.n
 25    }
 26
 27    fn size(&self) -> usize {
 28        1
 29    }
 30}
 31
 32impl<D> ZnxSliceSize for ScalarZnxDft<D, FFT64> {
 33    fn sl(&self) -> usize {
 34        self.n()
 35    }
 36}
 37
 38impl<D, B: Backend> DataView for ScalarZnxDft<D, B> {
 39    type D = D;
 40    fn data(&self) -> &Self::D {
 41        &self.data
 42    }
 43}
 44
 45impl<D, B: Backend> DataViewMut for ScalarZnxDft<D, B> {
 46    fn data_mut(&mut self) -> &mut Self::D {
 47        &mut self.data
 48    }
 49}
 50
 51impl<D: AsRef<[u8]>> ZnxView for ScalarZnxDft<D, FFT64> {
 52    type Scalar = f64;
 53}
 54
 55impl<D: From<Vec<u8>>, B: Backend> ScalarZnxDft<D, B> {
 56    pub(crate) fn bytes_of(module: &Module<B>, cols: usize) -> usize {
 57        unsafe { svp::bytes_of_svp_ppol(module.ptr) as usize * cols }
 58    }
 59
 60    pub(crate) fn new(module: &Module<B>, cols: usize) -> Self {
 61        let data = alloc_aligned::<u8>(Self::bytes_of(module, cols));
 62        Self {
 63            data: data.into(),
 64            n: module.n(),
 65            cols,
 66            _phantom: PhantomData,
 67        }
 68    }
 69
 70    pub(crate) fn new_from_bytes(module: &Module<B>, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
 71        let data: Vec<u8> = bytes.into();
 72        assert!(data.len() == Self::bytes_of(module, cols));
 73        Self {
 74            data: data.into(),
 75            n: module.n(),
 76            cols,
 77            _phantom: PhantomData,
 78        }
 79    }
 80}
 81
 82pub type ScalarZnxDftOwned<B> = ScalarZnxDft<Vec<u8>, B>;
 83
 84pub trait ScalarZnxDftToRef<B: Backend> {
 85    fn to_ref(&self) -> ScalarZnxDft<&[u8], B>;
 86}
 87
 88pub trait ScalarZnxDftToMut<B: Backend> {
 89    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B>;
 90}
 91
 92impl<B: Backend> ScalarZnxDftToMut<B> for ScalarZnxDft<Vec<u8>, B> {
 93    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
 94        ScalarZnxDft {
 95            data: self.data.as_mut_slice(),
 96            n: self.n,
 97            cols: self.cols,
 98            _phantom: PhantomData,
 99        }
100    }
101}
102
103impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<Vec<u8>, B> {
104    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
105        ScalarZnxDft {
106            data: self.data.as_slice(),
107            n: self.n,
108            cols: self.cols,
109            _phantom: PhantomData,
110        }
111    }
112}
113
114impl<B: Backend> ScalarZnxDftToMut<B> for ScalarZnxDft<&mut [u8], B> {
115    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
116        ScalarZnxDft {
117            data: self.data,
118            n: self.n,
119            cols: self.cols,
120            _phantom: PhantomData,
121        }
122    }
123}
124
125impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<&mut [u8], B> {
126    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
127        ScalarZnxDft {
128            data: self.data,
129            n: self.n,
130            cols: self.cols,
131            _phantom: PhantomData,
132        }
133    }
134}
135
136impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<&[u8], B> {
137    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
138        ScalarZnxDft {
139            data: self.data,
140            n: self.n,
141            cols: self.cols,
142            _phantom: PhantomData,
143        }
144    }
145}