1use crate::ffi::{vec_znx_big, vec_znx_dft};
  2use crate::vec_znx_dft::bytes_of_vec_znx_dft;
  3use crate::znx_base::ZnxInfos;
  4use crate::{Backend, Scratch, VecZnxBigToMut, VecZnxDftOwned, VecZnxDftToMut, VecZnxDftToRef, VecZnxToRef, ZnxSliceSize};
  5use crate::{FFT64, Module, ZnxView, ZnxViewMut, ZnxZero};
  6use std::cmp::min;
  7
  8pub trait VecZnxDftAlloc<B: Backend> {
  9    /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
 10    fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B>;
 11
 12    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
 13    ///
 14    /// Behavior: takes ownership of the backing array.
 15    ///
 16    /// # Arguments
 17    ///
 18    /// * `cols`: the number of cols of the [VecZnxDft].
 19    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
 20    ///
 21    /// # Panics
 22    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
 23    fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B>;
 24
 25    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
 26    ///
 27    /// # Arguments
 28    ///
 29    /// * `cols`: the number of cols of the [VecZnxDft].
 30    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
 31    ///
 32    /// # Panics
 33    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
 34    fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize;
 35}
 36
 37pub trait VecZnxDftOps<B: Backend> {
 38    /// Returns the minimum number of bytes necessary to allocate
 39    /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
 40    fn vec_znx_idft_tmp_bytes(&self) -> usize;
 41
 42    /// b <- IDFT(a), uses a as scratch space.
 43    fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_cols: usize)
 44    where
 45        R: VecZnxBigToMut<B>,
 46        A: VecZnxDftToMut<B>;
 47
 48    fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
 49    where
 50        R: VecZnxBigToMut<B>,
 51        A: VecZnxDftToRef<B>;
 52
 53    fn vec_znx_dft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
 54    where
 55        R: VecZnxDftToMut<B>,
 56        A: VecZnxToRef;
 57}
 58
 59impl<B: Backend> VecZnxDftAlloc<B> for Module<B> {
 60    fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B> {
 61        VecZnxDftOwned::new(&self, cols, size)
 62    }
 63
 64    fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B> {
 65        VecZnxDftOwned::new_from_bytes(self, cols, size, bytes)
 66    }
 67
 68    fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize {
 69        bytes_of_vec_znx_dft(self, cols, size)
 70    }
 71}
 72
 73impl VecZnxDftOps<FFT64> for Module<FFT64> {
 74    fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
 75    where
 76        R: VecZnxBigToMut<FFT64>,
 77        A: VecZnxDftToMut<FFT64>,
 78    {
 79        let mut res_mut = res.to_mut();
 80        let mut a_mut = a.to_mut();
 81
 82        let min_size: usize = min(res_mut.size(), a_mut.size());
 83
 84        unsafe {
 85            (0..min_size).for_each(|j| {
 86                vec_znx_dft::vec_znx_idft_tmp_a(
 87                    self.ptr,
 88                    res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
 89                    1 as u64,
 90                    a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
 91                    1 as u64,
 92                )
 93            });
 94            (min_size..res_mut.size()).for_each(|j| {
 95                res_mut.zero_at(res_col, j);
 96            })
 97        }
 98    }
 99
100    fn vec_znx_idft_tmp_bytes(&self) -> usize {
101        unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
102    }
103
104    /// b <- DFT(a)
105    ///
106    /// # Panics
107    /// If b.cols < a_cols
108    fn vec_znx_dft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
109    where
110        R: VecZnxDftToMut<FFT64>,
111        A: VecZnxToRef,
112    {
113        let mut res_mut = res.to_mut();
114        let a_ref = a.to_ref();
115
116        let min_size: usize = min(res_mut.size(), a_ref.size());
117
118        unsafe {
119            (0..min_size).for_each(|j| {
120                vec_znx_dft::vec_znx_dft(
121                    self.ptr,
122                    res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
123                    1 as u64,
124                    a_ref.at_ptr(a_col, j),
125                    1 as u64,
126                    a_ref.sl() as u64,
127                )
128            });
129            (min_size..res_mut.size()).for_each(|j| {
130                res_mut.zero_at(res_col, j);
131            });
132        }
133    }
134
135    // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
136    fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
137    where
138        R: VecZnxBigToMut<FFT64>,
139        A: VecZnxDftToRef<FFT64>,
140    {
141        let mut res_mut = res.to_mut();
142        let a_ref = a.to_ref();
143
144        let (tmp_bytes, _) = scratch.tmp_slice(self.vec_znx_idft_tmp_bytes());
145
146        let min_size: usize = min(res_mut.size(), a_ref.size());
147
148        unsafe {
149            (0..min_size).for_each(|j| {
150                vec_znx_dft::vec_znx_idft(
151                    self.ptr,
152                    res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
153                    1 as u64,
154                    a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
155                    1 as u64,
156                    tmp_bytes.as_mut_ptr(),
157                )
158            });
159            (min_size..res_mut.size()).for_each(|j| {
160                res_mut.zero_at(res_col, j);
161            });
162        }
163    }
164}