1use poulpy_hal::{
  2    layouts::{
  3        Backend, Data, Module, Scratch, VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftOwned, VecZnxDftToMut, VecZnxDftToRef,
  4        VecZnxToRef,
  5    },
  6    oep::{
  7        VecZnxDftAddImpl, VecZnxDftAddInplaceImpl, VecZnxDftAddScaledInplaceImpl, VecZnxDftAllocBytesImpl, VecZnxDftAllocImpl,
  8        VecZnxDftApplyImpl, VecZnxDftCopyImpl, VecZnxDftFromBytesImpl, VecZnxDftSubImpl, VecZnxDftSubInplaceImpl,
  9        VecZnxDftSubNegateInplaceImpl, VecZnxDftZeroImpl, VecZnxIdftApplyConsumeImpl, VecZnxIdftApplyImpl,
 10        VecZnxIdftApplyTmpAImpl, VecZnxIdftApplyTmpBytesImpl,
 11    },
 12    reference::fft64::vec_znx_dft::{
 13        vec_znx_dft_add, vec_znx_dft_add_inplace, vec_znx_dft_add_scaled_inplace, vec_znx_dft_apply, vec_znx_dft_copy,
 14        vec_znx_dft_sub, vec_znx_dft_sub_inplace, vec_znx_dft_sub_negate_inplace, vec_znx_dft_zero, vec_znx_idft_apply,
 15        vec_znx_idft_apply_consume, vec_znx_idft_apply_tmpa,
 16    },
 17};
 18
 19use crate::{FFT64Ref, module::FFT64ModuleHandle};
 20
 21unsafe impl VecZnxDftFromBytesImpl<Self> for FFT64Ref {
 22    fn vec_znx_dft_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<Self> {
 23        VecZnxDft::<Vec<u8>, Self>::from_bytes(n, cols, size, bytes)
 24    }
 25}
 26
 27unsafe impl VecZnxDftAllocBytesImpl<Self> for FFT64Ref {
 28    fn vec_znx_dft_bytes_of_impl(n: usize, cols: usize, size: usize) -> usize {
 29        Self::layout_prep_word_count() * n * cols * size * size_of::<<FFT64Ref as Backend>::ScalarPrep>()
 30    }
 31}
 32
 33unsafe impl VecZnxDftAllocImpl<Self> for FFT64Ref {
 34    fn vec_znx_dft_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxDftOwned<Self> {
 35        VecZnxDftOwned::alloc(n, cols, size)
 36    }
 37}
 38
 39unsafe impl VecZnxIdftApplyTmpBytesImpl<Self> for FFT64Ref {
 40    fn vec_znx_idft_apply_tmp_bytes_impl(_module: &Module<Self>) -> usize {
 41        0
 42    }
 43}
 44
 45unsafe impl VecZnxIdftApplyImpl<Self> for FFT64Ref {
 46    fn vec_znx_idft_apply_impl<R, A>(
 47        module: &Module<Self>,
 48        res: &mut R,
 49        res_col: usize,
 50        a: &A,
 51        a_col: usize,
 52        _scratch: &mut Scratch<Self>,
 53    ) where
 54        R: VecZnxBigToMut<Self>,
 55        A: VecZnxDftToRef<Self>,
 56    {
 57        vec_znx_idft_apply(module.get_ifft_table(), res, res_col, a, a_col);
 58    }
 59}
 60
 61unsafe impl VecZnxIdftApplyTmpAImpl<Self> for FFT64Ref {
 62    fn vec_znx_idft_apply_tmpa_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
 63    where
 64        R: VecZnxBigToMut<Self>,
 65        A: VecZnxDftToMut<Self>,
 66    {
 67        vec_znx_idft_apply_tmpa(module.get_ifft_table(), res, res_col, a, a_col);
 68    }
 69}
 70
 71unsafe impl VecZnxIdftApplyConsumeImpl<Self> for FFT64Ref {
 72    fn vec_znx_idft_apply_consume_impl<D: Data>(module: &Module<Self>, res: VecZnxDft<D, FFT64Ref>) -> VecZnxBig<D, FFT64Ref>
 73    where
 74        VecZnxDft<D, FFT64Ref>: VecZnxDftToMut<Self>,
 75    {
 76        vec_znx_idft_apply_consume(module.get_ifft_table(), res)
 77    }
 78}
 79
 80unsafe impl VecZnxDftApplyImpl<Self> for FFT64Ref {
 81    fn vec_znx_dft_apply_impl<R, A>(
 82        module: &Module<Self>,
 83        step: usize,
 84        offset: usize,
 85        res: &mut R,
 86        res_col: usize,
 87        a: &A,
 88        a_col: usize,
 89    ) where
 90        R: VecZnxDftToMut<Self>,
 91        A: VecZnxToRef,
 92    {
 93        vec_znx_dft_apply(module.get_fft_table(), step, offset, res, res_col, a, a_col);
 94    }
 95}
 96
 97unsafe impl VecZnxDftAddImpl<Self> for FFT64Ref {
 98    fn vec_znx_dft_add_impl<R, A, B>(
 99        _module: &Module<Self>,
100        res: &mut R,
101        res_col: usize,
102        a: &A,
103        a_col: usize,
104        b: &B,
105        b_col: usize,
106    ) where
107        R: VecZnxDftToMut<Self>,
108        A: VecZnxDftToRef<Self>,
109        B: VecZnxDftToRef<Self>,
110    {
111        vec_znx_dft_add(res, res_col, a, a_col, b, b_col);
112    }
113}
114
115unsafe impl VecZnxDftAddScaledInplaceImpl<Self> for FFT64Ref {
116    fn vec_znx_dft_add_scaled_inplace_impl<R, A>(
117        _module: &Module<Self>,
118        res: &mut R,
119        res_col: usize,
120        a: &A,
121        a_col: usize,
122        a_scale: i64,
123    ) where
124        R: VecZnxDftToMut<Self>,
125        A: VecZnxDftToRef<Self>,
126    {
127        vec_znx_dft_add_scaled_inplace(res, res_col, a, a_col, a_scale);
128    }
129}
130
131unsafe impl VecZnxDftAddInplaceImpl<Self> for FFT64Ref {
132    fn vec_znx_dft_add_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
133    where
134        R: VecZnxDftToMut<Self>,
135        A: VecZnxDftToRef<Self>,
136    {
137        vec_znx_dft_add_inplace(res, res_col, a, a_col);
138    }
139}
140
141unsafe impl VecZnxDftSubImpl<Self> for FFT64Ref {
142    fn vec_znx_dft_sub_impl<R, A, B>(
143        _module: &Module<Self>,
144        res: &mut R,
145        res_col: usize,
146        a: &A,
147        a_col: usize,
148        b: &B,
149        b_col: usize,
150    ) where
151        R: VecZnxDftToMut<Self>,
152        A: VecZnxDftToRef<Self>,
153        B: VecZnxDftToRef<Self>,
154    {
155        vec_znx_dft_sub(res, res_col, a, a_col, b, b_col);
156    }
157}
158
159unsafe impl VecZnxDftSubInplaceImpl<Self> for FFT64Ref {
160    fn vec_znx_dft_sub_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
161    where
162        R: VecZnxDftToMut<Self>,
163        A: VecZnxDftToRef<Self>,
164    {
165        vec_znx_dft_sub_inplace(res, res_col, a, a_col);
166    }
167}
168
169unsafe impl VecZnxDftSubNegateInplaceImpl<Self> for FFT64Ref {
170    fn vec_znx_dft_sub_negate_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
171    where
172        R: VecZnxDftToMut<Self>,
173        A: VecZnxDftToRef<Self>,
174    {
175        vec_znx_dft_sub_negate_inplace(res, res_col, a, a_col);
176    }
177}
178
179unsafe impl VecZnxDftCopyImpl<Self> for FFT64Ref {
180    fn vec_znx_dft_copy_impl<R, A>(
181        _module: &Module<Self>,
182        step: usize,
183        offset: usize,
184        res: &mut R,
185        res_col: usize,
186        a: &A,
187        a_col: usize,
188    ) where
189        R: VecZnxDftToMut<Self>,
190        A: VecZnxDftToRef<Self>,
191    {
192        vec_znx_dft_copy(step, offset, res, res_col, a, a_col);
193    }
194}
195
196unsafe impl VecZnxDftZeroImpl<Self> for FFT64Ref {
197    fn vec_znx_dft_zero_impl<R>(_module: &Module<Self>, res: &mut R, res_col: usize)
198    where
199        R: VecZnxDftToMut<Self>,
200    {
201        vec_znx_dft_zero(res, res_col);
202    }
203}