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