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
116pub trait ScalarZnxDftToMut<B: Backend> {
117    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B>;
118}
119
120impl<B: Backend> ScalarZnxDftToMut<B> for ScalarZnxDft<Vec<u8>, B> {
121    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
122        ScalarZnxDft {
123            data: self.data.as_mut_slice(),
124            n: self.n,
125            cols: self.cols,
126            _phantom: PhantomData,
127        }
128    }
129}
130
131impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<Vec<u8>, B> {
132    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
133        ScalarZnxDft {
134            data: self.data.as_slice(),
135            n: self.n,
136            cols: self.cols,
137            _phantom: PhantomData,
138        }
139    }
140}
141
142impl<B: Backend> ScalarZnxDftToMut<B> for ScalarZnxDft<&mut [u8], B> {
143    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
144        ScalarZnxDft {
145            data: self.data,
146            n: self.n,
147            cols: self.cols,
148            _phantom: PhantomData,
149        }
150    }
151}
152
153impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<&mut [u8], B> {
154    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
155        ScalarZnxDft {
156            data: self.data,
157            n: self.n,
158            cols: self.cols,
159            _phantom: PhantomData,
160        }
161    }
162}
163
164impl<B: Backend> ScalarZnxDftToRef<B> for ScalarZnxDft<&[u8], B> {
165    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
166        ScalarZnxDft {
167            data: self.data,
168            n: self.n,
169            cols: self.cols,
170            _phantom: PhantomData,
171        }
172    }
173}
174
175impl<B: Backend> VecZnxDftToMut<B> for ScalarZnxDft<Vec<u8>, B> {
176    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
177        VecZnxDft {
178            data: self.data.as_mut_slice(),
179            n: self.n,
180            cols: self.cols,
181            size: 1,
182            _phantom: PhantomData,
183        }
184    }
185}
186
187impl<B: Backend> VecZnxDftToRef<B> for ScalarZnxDft<Vec<u8>, B> {
188    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
189        VecZnxDft {
190            data: self.data.as_slice(),
191            n: self.n,
192            cols: self.cols,
193            size: 1,
194            _phantom: PhantomData,
195        }
196    }
197}
198
199impl<B: Backend> VecZnxDftToMut<B> for ScalarZnxDft<&mut [u8], B> {
200    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
201        VecZnxDft {
202            data: self.data,
203            n: self.n,
204            cols: self.cols,
205            size: 1,
206            _phantom: PhantomData,
207        }
208    }
209}
210
211impl<B: Backend> VecZnxDftToRef<B> for ScalarZnxDft<&mut [u8], B> {
212    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
213        VecZnxDft {
214            data: self.data,
215            n: self.n,
216            cols: self.cols,
217            size: 1,
218            _phantom: PhantomData,
219        }
220    }
221}
222
223impl<B: Backend> VecZnxDftToRef<B> for ScalarZnxDft<&[u8], B> {
224    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
225        VecZnxDft {
226            data: self.data,
227            n: self.n,
228            cols: self.cols,
229            size: 1,
230            _phantom: PhantomData,
231        }
232    }
233}