1use bytemuck::cast_slice_mut;
  2
  3use crate::{
  4    layouts::{
  5        Backend, Data, VecZnx, VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef, VecZnxToRef, ZnxInfos,
  6        ZnxView, ZnxViewMut,
  7    },
  8    reference::{
  9        fft64::reim::{
 10            ReimAdd, ReimAddInplace, ReimCopy, ReimDFTExecute, ReimFFTTable, ReimFromZnx, ReimIFFTTable, ReimNegate,
 11            ReimNegateInplace, ReimSub, ReimSubInplace, ReimSubNegateInplace, ReimToZnx, ReimToZnxInplace, ReimZero,
 12        },
 13        znx::ZnxZero,
 14    },
 15};
 16
 17pub fn vec_znx_dft_add<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 18where
 19    BE: Backend<ScalarPrep = f64> + ReimAdd + ReimCopy + ReimZero,
 20    R: VecZnxDftToMut<BE>,
 21    A: VecZnxDftToRef<BE>,
 22    B: VecZnxDftToRef<BE>,
 23{
 24    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
 25    let a: VecZnxDft<&[u8], BE> = a.to_ref();
 26    let b: VecZnxDft<&[u8], BE> = b.to_ref();
 27
 28    #[cfg(debug_assertions)]
 29    {
 30        assert_eq!(a.n(), res.n());
 31        assert_eq!(b.n(), res.n());
 32    }
 33
 34    let res_size: usize = res.size();
 35    let a_size: usize = a.size();
 36    let b_size: usize = b.size();
 37
 38    if a_size <= b_size {
 39        let sum_size: usize = a_size.min(res_size);
 40        let cpy_size: usize = b_size.min(res_size);
 41
 42        for j in 0..sum_size {
 43            BE::reim_add(res.at_mut(res_col, j), a.at(a_col, j), b.at(b_col, j));
 44        }
 45
 46        for j in sum_size..cpy_size {
 47            BE::reim_copy(res.at_mut(res_col, j), b.at(b_col, j));
 48        }
 49
 50        for j in cpy_size..res_size {
 51            BE::reim_zero(res.at_mut(res_col, j));
 52        }
 53    } else {
 54        let sum_size: usize = b_size.min(res_size);
 55        let cpy_size: usize = a_size.min(res_size);
 56
 57        for j in 0..sum_size {
 58            BE::reim_add(res.at_mut(res_col, j), a.at(a_col, j), b.at(b_col, j));
 59        }
 60
 61        for j in sum_size..cpy_size {
 62            BE::reim_copy(res.at_mut(res_col, j), a.at(a_col, j));
 63        }
 64
 65        for j in cpy_size..res_size {
 66            BE::reim_zero(res.at_mut(res_col, j));
 67        }
 68    }
 69}
 70
 71pub fn vec_znx_dft_add_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
 72where
 73    BE: Backend<ScalarPrep = f64> + ReimAddInplace,
 74    R: VecZnxDftToMut<BE>,
 75    A: VecZnxDftToRef<BE>,
 76{
 77    let a: VecZnxDft<&[u8], BE> = a.to_ref();
 78    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
 79
 80    #[cfg(debug_assertions)]
 81    {
 82        assert_eq!(a.n(), res.n());
 83    }
 84
 85    let res_size: usize = res.size();
 86    let a_size: usize = a.size();
 87
 88    let sum_size: usize = a_size.min(res_size);
 89
 90    for j in 0..sum_size {
 91        BE::reim_add_inplace(res.at_mut(res_col, j), a.at(a_col, j));
 92    }
 93}
 94
 95pub fn vec_znx_dft_copy<R, A, BE>(step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
 96where
 97    BE: Backend<ScalarPrep = f64> + ReimCopy + ReimZero,
 98    R: VecZnxDftToMut<BE>,
 99    A: VecZnxDftToRef<BE>,
100{
101    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
102    let a: VecZnxDft<&[u8], BE> = a.to_ref();
103
104    #[cfg(debug_assertions)]
105    {
106        assert_eq!(res.n(), a.n())
107    }
108
109    let steps: usize = a.size().div_ceil(step);
110    let min_steps: usize = res.size().min(steps);
111
112    (0..min_steps).for_each(|j| {
113        let limb: usize = offset + j * step;
114        if limb < a.size() {
115            BE::reim_copy(res.at_mut(res_col, j), a.at(a_col, limb));
116        }
117    });
118    (min_steps..res.size()).for_each(|j| {
119        BE::reim_zero(res.at_mut(res_col, j));
120    })
121}
122
123pub fn vec_znx_dft_apply<R, A, BE>(
124    table: &ReimFFTTable<f64>,
125    step: usize,
126    offset: usize,
127    res: &mut R,
128    res_col: usize,
129    a: &A,
130    a_col: usize,
131) where
132    BE: Backend<ScalarPrep = f64> + ReimDFTExecute<ReimFFTTable<f64>, f64> + ReimFromZnx + ReimZero,
133    R: VecZnxDftToMut<BE>,
134    A: VecZnxToRef,
135{
136    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
137    let a: VecZnx<&[u8]> = a.to_ref();
138
139    #[cfg(debug_assertions)]
140    {
141        assert!(step > 0);
142        assert_eq!(table.m() << 1, res.n());
143        assert_eq!(a.n(), res.n());
144    }
145
146    let a_size: usize = a.size();
147    let res_size: usize = res.size();
148
149    let steps: usize = a_size.div_ceil(step);
150    let min_steps: usize = res_size.min(steps);
151
152    for j in 0..min_steps {
153        let limb = offset + j * step;
154        if limb < a_size {
155            BE::reim_from_znx(res.at_mut(res_col, j), a.at(a_col, limb));
156            BE::reim_dft_execute(table, res.at_mut(res_col, j));
157        }
158    }
159
160    (min_steps..res.size()).for_each(|j| {
161        BE::reim_zero(res.at_mut(res_col, j));
162    });
163}
164
165pub fn vec_znx_idft_apply<R, A, BE>(table: &ReimIFFTTable<f64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
166where
167    BE: Backend<ScalarPrep = f64, ScalarBig = i64>
168        + ReimDFTExecute<ReimIFFTTable<f64>, f64>
169        + ReimCopy
170        + ReimToZnxInplace
171        + ZnxZero,
172    R: VecZnxBigToMut<BE>,
173    A: VecZnxDftToRef<BE>,
174{
175    let mut res: VecZnxBig<&mut [u8], BE> = res.to_mut();
176    let a: VecZnxDft<&[u8], BE> = a.to_ref();
177
178    #[cfg(debug_assertions)]
179    {
180        assert_eq!(table.m() << 1, res.n());
181        assert_eq!(a.n(), res.n());
182    }
183
184    let res_size: usize = res.size();
185    let min_size: usize = res_size.min(a.size());
186
187    let divisor: f64 = table.m() as f64;
188
189    for j in 0..min_size {
190        let res_slice_f64: &mut [f64] = cast_slice_mut(res.at_mut(res_col, j));
191        BE::reim_copy(res_slice_f64, a.at(a_col, j));
192        BE::reim_dft_execute(table, res_slice_f64);
193        BE::reim_to_znx_inplace(res_slice_f64, divisor);
194    }
195
196    for j in min_size..res_size {
197        BE::znx_zero(res.at_mut(res_col, j));
198    }
199}
200
201pub fn vec_znx_idft_apply_tmpa<R, A, BE>(table: &ReimIFFTTable<f64>, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
202where
203    BE: Backend<ScalarPrep = f64, ScalarBig = i64> + ReimDFTExecute<ReimIFFTTable<f64>, f64> + ReimToZnx + ZnxZero,
204    R: VecZnxBigToMut<BE>,
205    A: VecZnxDftToMut<BE>,
206{
207    let mut res: VecZnxBig<&mut [u8], BE> = res.to_mut();
208    let mut a: VecZnxDft<&mut [u8], BE> = a.to_mut();
209
210    #[cfg(debug_assertions)]
211    {
212        assert_eq!(table.m() << 1, res.n());
213        assert_eq!(a.n(), res.n());
214    }
215
216    let res_size = res.size();
217    let min_size: usize = res_size.min(a.size());
218
219    let divisor: f64 = table.m() as f64;
220
221    for j in 0..min_size {
222        BE::reim_dft_execute(table, a.at_mut(a_col, j));
223        BE::reim_to_znx(res.at_mut(res_col, j), divisor, a.at(a_col, j));
224    }
225
226    for j in min_size..res_size {
227        BE::znx_zero(res.at_mut(res_col, j));
228    }
229}
230
231pub fn vec_znx_idft_apply_consume<D: Data, BE>(table: &ReimIFFTTable<f64>, mut res: VecZnxDft<D, BE>) -> VecZnxBig<D, BE>
232where
233    BE: Backend<ScalarPrep = f64, ScalarBig = i64> + ReimDFTExecute<ReimIFFTTable<f64>, f64> + ReimToZnxInplace,
234    VecZnxDft<D, BE>: VecZnxDftToMut<BE>,
235{
236    {
237        let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
238
239        #[cfg(debug_assertions)]
240        {
241            assert_eq!(table.m() << 1, res.n());
242        }
243
244        let divisor: f64 = table.m() as f64;
245
246        for i in 0..res.cols() {
247            for j in 0..res.size() {
248                BE::reim_dft_execute(table, res.at_mut(i, j));
249                BE::reim_to_znx_inplace(res.at_mut(i, j), divisor);
250            }
251        }
252    }
253
254    res.into_big()
255}
256
257pub fn vec_znx_dft_sub<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
258where
259    BE: Backend<ScalarPrep = f64> + ReimSub + ReimNegate + ReimZero + ReimCopy,
260    R: VecZnxDftToMut<BE>,
261    A: VecZnxDftToRef<BE>,
262    B: VecZnxDftToRef<BE>,
263{
264    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
265    let a: VecZnxDft<&[u8], BE> = a.to_ref();
266    let b: VecZnxDft<&[u8], BE> = b.to_ref();
267
268    #[cfg(debug_assertions)]
269    {
270        assert_eq!(a.n(), res.n());
271        assert_eq!(b.n(), res.n());
272    }
273
274    let res_size: usize = res.size();
275    let a_size: usize = a.size();
276    let b_size: usize = b.size();
277
278    if a_size <= b_size {
279        let sum_size: usize = a_size.min(res_size);
280        let cpy_size: usize = b_size.min(res_size);
281
282        for j in 0..sum_size {
283            BE::reim_sub(res.at_mut(res_col, j), a.at(a_col, j), b.at(b_col, j));
284        }
285
286        for j in sum_size..cpy_size {
287            BE::reim_negate(res.at_mut(res_col, j), b.at(b_col, j));
288        }
289
290        for j in cpy_size..res_size {
291            BE::reim_zero(res.at_mut(res_col, j));
292        }
293    } else {
294        let sum_size: usize = b_size.min(res_size);
295        let cpy_size: usize = a_size.min(res_size);
296
297        for j in 0..sum_size {
298            BE::reim_sub(res.at_mut(res_col, j), a.at(a_col, j), b.at(b_col, j));
299        }
300
301        for j in sum_size..cpy_size {
302            BE::reim_copy(res.at_mut(res_col, j), a.at(a_col, j));
303        }
304
305        for j in cpy_size..res_size {
306            BE::reim_zero(res.at_mut(res_col, j));
307        }
308    }
309}
310
311pub fn vec_znx_dft_sub_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
312where
313    BE: Backend<ScalarPrep = f64> + ReimSubInplace,
314    R: VecZnxDftToMut<BE>,
315    A: VecZnxDftToRef<BE>,
316{
317    let a: VecZnxDft<&[u8], BE> = a.to_ref();
318    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
319
320    #[cfg(debug_assertions)]
321    {
322        assert_eq!(a.n(), res.n());
323    }
324
325    let res_size: usize = res.size();
326    let a_size: usize = a.size();
327
328    let sum_size: usize = a_size.min(res_size);
329
330    for j in 0..sum_size {
331        BE::reim_sub_inplace(res.at_mut(res_col, j), a.at(a_col, j));
332    }
333}
334
335pub fn vec_znx_dft_sub_negate_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
336where
337    BE: Backend<ScalarPrep = f64> + ReimSubNegateInplace + ReimNegateInplace,
338    R: VecZnxDftToMut<BE>,
339    A: VecZnxDftToRef<BE>,
340{
341    let a: VecZnxDft<&[u8], BE> = a.to_ref();
342    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
343
344    #[cfg(debug_assertions)]
345    {
346        assert_eq!(a.n(), res.n());
347    }
348
349    let res_size: usize = res.size();
350    let a_size: usize = a.size();
351
352    let sum_size: usize = a_size.min(res_size);
353
354    for j in 0..sum_size {
355        BE::reim_sub_negate_inplace(res.at_mut(res_col, j), a.at(a_col, j));
356    }
357
358    for j in sum_size..res_size {
359        BE::reim_negate_inplace(res.at_mut(res_col, j));
360    }
361}
362
363pub fn vec_znx_dft_zero<R, BE>(res: &mut R)
364where
365    R: VecZnxDftToMut<BE>,
366    BE: Backend<ScalarPrep = f64> + ReimZero,
367{
368    BE::reim_zero(res.to_mut().raw_mut());
369}