arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/base2k/src/scalar_znx_dft.rs
base2k/src/scalar_znx_dft.rsRaw
 1use std::marker::PhantomData;
 2
 3use crate::ffi::svp;
 4use crate::znx_base::ZnxInfos;
 5use crate::{Backend, DataView, DataViewMut, FFT64, Module, ZnxView, alloc_aligned};
 6
 7pub const SCALAR_ZNX_DFT_ROWS: usize = 1;
 8pub const SCALAR_ZNX_DFT_SIZE: usize = 1;
 9
10pub struct ScalarZnxDft<D, B> {
11    data: D,
12    n: usize,
13    cols: usize,
14    _phantom: PhantomData<B>,
15}
16
17impl<D, B> 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    fn sl(&self) -> usize {
35        self.n()
36    }
37}
38
39impl<D, B> DataView for ScalarZnxDft<D, B> {
40    type D = D;
41    fn data(&self) -> &Self::D {
42        &self.data
43    }
44}
45
46impl<D, B> DataViewMut for ScalarZnxDft<D, B> {
47    fn data_mut(&mut self) -> &mut Self::D {
48        &mut self.data
49    }
50}
51
52impl<D: AsRef<[u8]>> ZnxView for ScalarZnxDft<D, FFT64> {
53    type Scalar = f64;
54}
55
56impl<D: From<Vec<u8>>, B: Backend> ScalarZnxDft<D, B> {
57    pub(crate) fn bytes_of(module: &Module<B>, cols: usize) -> usize {
58        unsafe { svp::bytes_of_svp_ppol(module.ptr) as usize * cols }
59    }
60
61    pub(crate) fn new(module: &Module<B>, cols: usize) -> Self {
62        let data = alloc_aligned::<u8>(Self::bytes_of(module, cols));
63        Self {
64            data: data.into(),
65            n: module.n(),
66            cols,
67            _phantom: PhantomData,
68        }
69    }
70
71    pub(crate) fn new_from_bytes(module: &Module<B>, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
72        let data: Vec<u8> = bytes.into();
73        assert!(data.len() == Self::bytes_of(module, cols));
74        Self {
75            data: data.into(),
76            n: module.n(),
77            cols,
78            _phantom: PhantomData,
79        }
80    }
81
82    // fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, _size: usize, bytes: &mut [u8]) -> Self {
83    //     debug_assert_eq!(bytes.len(), Self::bytes_of(module, _rows, cols, _size));
84    //     Self {
85    //         inner: ZnxBase::from_bytes_borrow(
86    //             module.n(),
87    //             SCALAR_ZNX_DFT_ROWS,
88    //             cols,
89    //             SCALAR_ZNX_DFT_SIZE,
90    //             bytes,
91    //         ),
92    //         _phantom: PhantomData,
93    //     }
94    // }
95}
96
97pub type ScalarZnxDftOwned<B> = ScalarZnxDft<Vec<u8>, B>;