1use crate::ffi::vec_znx_dft;
 2use crate::znx_base::{GetZnxBase, ZnxAlloc, ZnxBase, ZnxInfos, ZnxLayout, ZnxSliceSize, ZnxZero};
 3use crate::{Backend, FFT64, Module, VecZnxBig};
 4use std::marker::PhantomData;
 5
 6const VEC_ZNX_DFT_ROWS: usize = 1;
 7
 8pub struct VecZnxDft<B: Backend> {
 9    inner: ZnxBase,
10    pub _marker: PhantomData<B>,
11}
12
13impl<B: Backend> GetZnxBase for VecZnxDft<B> {
14    fn znx(&self) -> &ZnxBase {
15        &self.inner
16    }
17
18    fn znx_mut(&mut self) -> &mut ZnxBase {
19        &mut self.inner
20    }
21}
22
23impl<B: Backend> ZnxInfos for VecZnxDft<B> {}
24
25impl<B: Backend> ZnxAlloc<B> for VecZnxDft<B> {
26    type Scalar = u8;
27
28    fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> Self {
29        debug_assert_eq!(bytes.len(), Self::bytes_of(module, _rows, cols, size));
30        Self {
31            inner: ZnxBase::from_bytes_borrow(module.n(), VEC_ZNX_DFT_ROWS, cols, size, bytes),
32            _marker: PhantomData,
33        }
34    }
35
36    fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, size: usize) -> usize {
37        debug_assert_eq!(
38            _rows, VEC_ZNX_DFT_ROWS,
39            "rows != {} not supported for VecZnxDft",
40            VEC_ZNX_DFT_ROWS
41        );
42        unsafe { vec_znx_dft::bytes_of_vec_znx_dft(module.ptr, size as u64) as usize * cols }
43    }
44}
45
46impl ZnxLayout for VecZnxDft<FFT64> {
47    type Scalar = f64;
48}
49
50impl ZnxZero for VecZnxDft<FFT64> {}
51
52impl ZnxSliceSize for VecZnxDft<FFT64> {
53    fn sl(&self) -> usize {
54        self.n()
55    }
56}
57
58impl VecZnxDft<FFT64> {
59    pub fn print(&self, n: usize, col: usize) {
60        (0..self.size()).for_each(|i| println!("{}: {:?}", i, &self.at(col, i)[..n]));
61    }
62}
63
64impl<B: Backend> VecZnxDft<B> {
65    /// Cast a [VecZnxDft] into a [VecZnxBig].
66    /// The returned [VecZnxBig] shares the backing array
67    /// with the original [VecZnxDft].
68    pub fn alias_as_vec_znx_big(&mut self) -> VecZnxBig<B> {
69        assert!(
70            self.data().len() == 0,
71            "cannot alias VecZnxDft into VecZnxBig if it owns the data"
72        );
73        VecZnxBig::<B> {
74            inner: ZnxBase {
75                data: Vec::new(),
76                ptr: self.ptr(),
77                n: self.n(),
78                rows: self.rows(),
79                cols: self.cols(),
80                size: self.size(),
81            },
82            _marker: PhantomData,
83        }
84    }
85}