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
 55pub(crate) fn bytes_of_scalar_znx_dft<B: Backend>(module: &Module<B>, cols: usize) -> usize {
 56    ScalarZnxDftOwned::bytes_of(module, cols)
 57}
 58
 59impl<D: From<Vec<u8>>, B: Backend> ScalarZnxDft<D, B> {
 60    pub(crate) fn bytes_of(module: &Module<B>, cols: usize) -> usize {
 61        unsafe { svp::bytes_of_svp_ppol(module.ptr) as usize * cols }
 62    }
 63
 64    pub(crate) fn new(module: &Module<B>, cols: usize) -> Self {
 65        let data = alloc_aligned::<u8>(Self::bytes_of(module, cols));
 66        Self {
 67            data: data.into(),
 68            n: module.n(),
 69            cols,
 70            _phantom: PhantomData,
 71        }
 72    }
 73
 74    pub(crate) fn new_from_bytes(module: &Module<B>, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
 75        let data: Vec<u8> = bytes.into();
 76        assert!(data.len() == Self::bytes_of(module, cols));
 77        Self {
 78            data: data.into(),
 79            n: module.n(),
 80            cols,
 81            _phantom: PhantomData,
 82        }
 83    }
 84}
 85
 86impl<D, B: Backend> ScalarZnxDft<D, B> {
 87    pub(crate) fn from_data(data: D, n: usize, cols: usize) -> Self {
 88        Self {
 89            data,
 90            n,
 91            cols,
 92            _phantom: PhantomData,
 93        }
 94    }
 95}
 96
 97pub type ScalarZnxDftOwned<B> = ScalarZnxDft<Vec<u8>, B>;
 98
 99pub trait ScalarZnxDftToRef<B: Backend> {
100    fn to_ref(&self) -> ScalarZnxDft<&[u8], B>;
101}
102
103pub trait ScalarZnxDftToMut<B: Backend> {
104    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B>;
105}
106
107impl<B: Backend> ScalarZnxDftToMut<B> for ScalarZnxDft<Vec<u8>, B> {
108    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
109        ScalarZnxDft {
110            data: self.data.as_mut_slice(),
111            n: self.n,
112            cols: self.cols,
113            _phantom: PhantomData,
114        }
115    }
116}
117
118impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<Vec<u8>, B> {
119    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
120        ScalarZnxDft {
121            data: self.data.as_slice(),
122            n: self.n,
123            cols: self.cols,
124            _phantom: PhantomData,
125        }
126    }
127}
128
129impl<B: Backend> ScalarZnxDftToMut<B> for ScalarZnxDft<&mut [u8], B> {
130    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
131        ScalarZnxDft {
132            data: self.data,
133            n: self.n,
134            cols: self.cols,
135            _phantom: PhantomData,
136        }
137    }
138}
139
140impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<&mut [u8], B> {
141    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
142        ScalarZnxDft {
143            data: self.data,
144            n: self.n,
145            cols: self.cols,
146            _phantom: PhantomData,
147        }
148    }
149}
150
151impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<&[u8], B> {
152    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
153        ScalarZnxDft {
154            data: self.data,
155            n: self.n,
156            cols: self.cols,
157            _phantom: PhantomData,
158        }
159    }
160}