1use std::marker::PhantomData;
  2
  3use crate::ffi::svp;
  4use crate::znx_base::ZnxInfos;
  5use crate::{alloc_aligned, Backend, DataView, DataViewMut, Module, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef, ZnxSliceSize, ZnxView, FFT64};
  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    pub fn as_vec_znx_dft(self) -> VecZnxDft<D, B>{
 97        VecZnxDft{
 98            data: self.data,
 99            n: self.n,
100            cols: self.cols,
101            size: 1,
102            _phantom: PhantomData,
103        }
104    }
105}
106
107pub type ScalarZnxDftOwned<B> = ScalarZnxDft<Vec<u8>, B>;
108
109pub trait ScalarZnxDftToRef<B: Backend> {
110    fn to_ref(&self) -> ScalarZnxDft<&[u8], B>;
111}
112
113pub trait ScalarZnxDftToMut<B: Backend> {
114    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B>;
115}
116
117impl<B: Backend> ScalarZnxDftToMut<B> for ScalarZnxDft<Vec<u8>, B> {
118    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
119        ScalarZnxDft {
120            data: self.data.as_mut_slice(),
121            n: self.n,
122            cols: self.cols,
123            _phantom: PhantomData,
124        }
125    }
126}
127
128impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<Vec<u8>, B> {
129    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
130        ScalarZnxDft {
131            data: self.data.as_slice(),
132            n: self.n,
133            cols: self.cols,
134            _phantom: PhantomData,
135        }
136    }
137}
138
139impl<B: Backend> ScalarZnxDftToMut<B> for ScalarZnxDft<&mut [u8], B> {
140    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
141        ScalarZnxDft {
142            data: self.data,
143            n: self.n,
144            cols: self.cols,
145            _phantom: PhantomData,
146        }
147    }
148}
149
150impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<&mut [u8], B> {
151    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
152        ScalarZnxDft {
153            data: self.data,
154            n: self.n,
155            cols: self.cols,
156            _phantom: PhantomData,
157        }
158    }
159}
160
161impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<&[u8], B> {
162    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
163        ScalarZnxDft {
164            data: self.data,
165            n: self.n,
166            cols: self.cols,
167            _phantom: PhantomData,
168        }
169    }
170}
171
172impl<B: Backend> VecZnxDftToMut<B> for ScalarZnxDft<Vec<u8>, B> {
173    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
174        VecZnxDft {
175            data: self.data.as_mut_slice(),
176            n: self.n,
177            cols: self.cols,
178            size: 1,
179            _phantom: PhantomData,
180        }
181    }
182}
183
184impl<B: Backend> VecZnxDftToRef<B> for ScalarZnxDft<Vec<u8>, B> {
185    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
186        VecZnxDft {
187            data: self.data.as_slice(),
188            n: self.n,
189            cols: self.cols,
190            size: 1,
191            _phantom: PhantomData,
192        }
193    }
194}
195
196impl<B: Backend> VecZnxDftToMut<B> for ScalarZnxDft<&mut [u8], B> {
197    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
198        VecZnxDft {
199            data: self.data,
200            n: self.n,
201            cols: self.cols,
202            size: 1,
203            _phantom: PhantomData,
204        }
205    }
206}
207
208impl<B: Backend> VecZnxDftToRef<B> for ScalarZnxDft<&mut [u8], B> {
209    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
210        VecZnxDft {
211            data: self.data,
212            n: self.n,
213            cols: self.cols,
214            size: 1,
215            _phantom: PhantomData,
216        }
217    }
218}
219
220impl<B: Backend> VecZnxDftToRef<B> for ScalarZnxDft<&[u8], B> {
221    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
222        VecZnxDft {
223            data: self.data,
224            n: self.n,
225            cols: self.cols,
226            size: 1,
227            _phantom: PhantomData,
228        }
229    }
230}