1use crate::ffi::vec_znx_dft;
 2use crate::znx_base::{GetZnxBase, ZnxAlloc, ZnxBase, ZnxInfos, ZnxLayout, ZnxSliceSize};
 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        Self {
30            inner: ZnxBase::from_bytes_borrow(module.n(), VEC_ZNX_DFT_ROWS, cols, size, bytes),
31            _marker: PhantomData,
32        }
33    }
34
35    fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, size: usize) -> usize {
36        debug_assert_eq!(
37            _rows, VEC_ZNX_DFT_ROWS,
38            "rows != {} not supported for VecZnxDft",
39            VEC_ZNX_DFT_ROWS
40        );
41        unsafe { vec_znx_dft::bytes_of_vec_znx_dft(module.ptr, size as u64) as usize * cols }
42    }
43}
44
45impl ZnxLayout for VecZnxDft<FFT64> {
46    type Scalar = f64;
47}
48
49impl ZnxSliceSize for VecZnxDft<FFT64> {
50    fn sl(&self) -> usize {
51        self.n()
52    }
53}
54
55impl VecZnxDft<FFT64> {
56    pub fn print(&self, n: usize) {
57        (0..self.size()).for_each(|i| println!("{}: {:?}", i, &self.at_limb(i)[..n]));
58    }
59}
60
61impl<B: Backend> VecZnxDft<B> {
62    /// Cast a [VecZnxDft] into a [VecZnxBig].
63    /// The returned [VecZnxBig] shares the backing array
64    /// with the original [VecZnxDft].
65    pub fn alias_as_vec_znx_big(&mut self) -> VecZnxBig<B> {
66        VecZnxBig::<B> {
67            inner: ZnxBase {
68                data: Vec::new(),
69                ptr: self.ptr(),
70                n: self.n(),
71                rows: self.rows(),
72                cols: self.cols(),
73                size: self.size(),
74            },
75            _marker: PhantomData,
76        }
77    }
78}