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
 95/// res = res + a * 2^{a_scale * base2k}.
 96pub fn vec_znx_dft_add_scaled_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, a_scale: i64)
 97where
 98    BE: Backend<ScalarPrep = f64> + ReimAddInplace,
 99    R: VecZnxDftToMut<BE>,
100    A: VecZnxDftToRef<BE>,
101{
102    let a: VecZnxDft<&[u8], BE> = a.to_ref();
103    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
104
105    #[cfg(debug_assertions)]
106    {
107        assert_eq!(a.n(), res.n());
108    }
109
110    let res_size: usize = res.size();
111    let a_size: usize = a.size();
112
113    if a_scale > 0 {
114        let shift: usize = (a_scale as usize).min(a_size);
115        let sum_size: usize = a_size.min(res_size).saturating_sub(shift);
116        for j in 0..sum_size {
117            BE::reim_add_inplace(res.at_mut(res_col, j), a.at(a_col, j + shift));
118        }
119    } else if a_scale < 0 {
120        let shift: usize = (a_scale.unsigned_abs() as usize).min(res_size);
121        let sum_size: usize = a_size.min(res_size).saturating_sub(shift);
122        for j in 0..sum_size {
123            BE::reim_add_inplace(res.at_mut(res_col, j + shift), a.at(a_col, j));
124        }
125    } else {
126        let sum_size: usize = a_size.min(res_size);
127        for j in 0..sum_size {
128            BE::reim_add_inplace(res.at_mut(res_col, j), a.at(a_col, j));
129        }
130    }
131}
132
133pub fn vec_znx_dft_copy<R, A, BE>(step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
134where
135    BE: Backend<ScalarPrep = f64> + ReimCopy + ReimZero,
136    R: VecZnxDftToMut<BE>,
137    A: VecZnxDftToRef<BE>,
138{
139    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
140    let a: VecZnxDft<&[u8], BE> = a.to_ref();
141
142    #[cfg(debug_assertions)]
143    {
144        assert_eq!(res.n(), a.n())
145    }
146
147    let steps: usize = a.size().div_ceil(step);
148    let min_steps: usize = res.size().min(steps);
149
150    (0..min_steps).for_each(|j| {
151        let limb: usize = offset + j * step;
152        if limb < a.size() {
153            BE::reim_copy(res.at_mut(res_col, j), a.at(a_col, limb));
154        }
155    });
156    (min_steps..res.size()).for_each(|j| {
157        BE::reim_zero(res.at_mut(res_col, j));
158    })
159}
160
161pub fn vec_znx_dft_apply<R, A, BE>(
162    table: &ReimFFTTable<f64>,
163    step: usize,
164    offset: usize,
165    res: &mut R,
166    res_col: usize,
167    a: &A,
168    a_col: usize,
169) where
170    BE: Backend<ScalarPrep = f64> + ReimDFTExecute<ReimFFTTable<f64>, f64> + ReimFromZnx + ReimZero,
171    R: VecZnxDftToMut<BE>,
172    A: VecZnxToRef,
173{
174    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
175    let a: VecZnx<&[u8]> = a.to_ref();
176
177    #[cfg(debug_assertions)]
178    {
179        assert!(step > 0);
180        assert_eq!(table.m() << 1, res.n());
181        assert_eq!(a.n(), res.n());
182    }
183
184    let a_size: usize = a.size();
185    let res_size: usize = res.size();
186
187    let steps: usize = a_size.div_ceil(step);
188    let min_steps: usize = res_size.min(steps);
189
190    for j in 0..min_steps {
191        let limb = offset + j * step;
192        if limb < a_size {
193            BE::reim_from_znx(res.at_mut(res_col, j), a.at(a_col, limb));
194            BE::reim_dft_execute(table, res.at_mut(res_col, j));
195        }
196    }
197
198    (min_steps..res.size()).for_each(|j| {
199        BE::reim_zero(res.at_mut(res_col, j));
200    });
201}
202
203pub fn vec_znx_idft_apply<R, A, BE>(table: &ReimIFFTTable<f64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
204where
205    BE: Backend<ScalarPrep = f64, ScalarBig = i64>
206        + ReimDFTExecute<ReimIFFTTable<f64>, f64>
207        + ReimCopy
208        + ReimToZnxInplace
209        + ZnxZero,
210    R: VecZnxBigToMut<BE>,
211    A: VecZnxDftToRef<BE>,
212{
213    let mut res: VecZnxBig<&mut [u8], BE> = res.to_mut();
214    let a: VecZnxDft<&[u8], BE> = a.to_ref();
215
216    #[cfg(debug_assertions)]
217    {
218        assert_eq!(table.m() << 1, res.n());
219        assert_eq!(a.n(), res.n());
220    }
221
222    let res_size: usize = res.size();
223    let min_size: usize = res_size.min(a.size());
224
225    let divisor: f64 = table.m() as f64;
226
227    for j in 0..min_size {
228        let res_slice_f64: &mut [f64] = cast_slice_mut(res.at_mut(res_col, j));
229        BE::reim_copy(res_slice_f64, a.at(a_col, j));
230        BE::reim_dft_execute(table, res_slice_f64);
231        BE::reim_to_znx_inplace(res_slice_f64, divisor);
232    }
233
234    for j in min_size..res_size {
235        BE::znx_zero(res.at_mut(res_col, j));
236    }
237}
238
239pub fn vec_znx_idft_apply_tmpa<R, A, BE>(table: &ReimIFFTTable<f64>, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
240where
241    BE: Backend<ScalarPrep = f64, ScalarBig = i64> + ReimDFTExecute<ReimIFFTTable<f64>, f64> + ReimToZnx + ZnxZero,
242    R: VecZnxBigToMut<BE>,
243    A: VecZnxDftToMut<BE>,
244{
245    let mut res: VecZnxBig<&mut [u8], BE> = res.to_mut();
246    let mut a: VecZnxDft<&mut [u8], BE> = a.to_mut();
247
248    #[cfg(debug_assertions)]
249    {
250        assert_eq!(table.m() << 1, res.n());
251        assert_eq!(a.n(), res.n());
252    }
253
254    let res_size = res.size();
255    let min_size: usize = res_size.min(a.size());
256
257    let divisor: f64 = table.m() as f64;
258
259    for j in 0..min_size {
260        BE::reim_dft_execute(table, a.at_mut(a_col, j));
261        BE::reim_to_znx(res.at_mut(res_col, j), divisor, a.at(a_col, j));
262    }
263
264    for j in min_size..res_size {
265        BE::znx_zero(res.at_mut(res_col, j));
266    }
267}
268
269pub fn vec_znx_idft_apply_consume<D: Data, BE>(table: &ReimIFFTTable<f64>, mut res: VecZnxDft<D, BE>) -> VecZnxBig<D, BE>
270where
271    BE: Backend<ScalarPrep = f64, ScalarBig = i64> + ReimDFTExecute<ReimIFFTTable<f64>, f64> + ReimToZnxInplace,
272    VecZnxDft<D, BE>: VecZnxDftToMut<BE>,
273{
274    {
275        let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
276
277        #[cfg(debug_assertions)]
278        {
279            assert_eq!(table.m() << 1, res.n());
280        }
281
282        let divisor: f64 = table.m() as f64;
283
284        for i in 0..res.cols() {
285            for j in 0..res.size() {
286                BE::reim_dft_execute(table, res.at_mut(i, j));
287                BE::reim_to_znx_inplace(res.at_mut(i, j), divisor);
288            }
289        }
290    }
291
292    res.into_big()
293}
294
295pub 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)
296where
297    BE: Backend<ScalarPrep = f64> + ReimSub + ReimNegate + ReimZero + ReimCopy,
298    R: VecZnxDftToMut<BE>,
299    A: VecZnxDftToRef<BE>,
300    B: VecZnxDftToRef<BE>,
301{
302    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
303    let a: VecZnxDft<&[u8], BE> = a.to_ref();
304    let b: VecZnxDft<&[u8], BE> = b.to_ref();
305
306    #[cfg(debug_assertions)]
307    {
308        assert_eq!(a.n(), res.n());
309        assert_eq!(b.n(), res.n());
310    }
311
312    let res_size: usize = res.size();
313    let a_size: usize = a.size();
314    let b_size: usize = b.size();
315
316    if a_size <= b_size {
317        let sum_size: usize = a_size.min(res_size);
318        let cpy_size: usize = b_size.min(res_size);
319
320        for j in 0..sum_size {
321            BE::reim_sub(res.at_mut(res_col, j), a.at(a_col, j), b.at(b_col, j));
322        }
323
324        for j in sum_size..cpy_size {
325            BE::reim_negate(res.at_mut(res_col, j), b.at(b_col, j));
326        }
327
328        for j in cpy_size..res_size {
329            BE::reim_zero(res.at_mut(res_col, j));
330        }
331    } else {
332        let sum_size: usize = b_size.min(res_size);
333        let cpy_size: usize = a_size.min(res_size);
334
335        for j in 0..sum_size {
336            BE::reim_sub(res.at_mut(res_col, j), a.at(a_col, j), b.at(b_col, j));
337        }
338
339        for j in sum_size..cpy_size {
340            BE::reim_copy(res.at_mut(res_col, j), a.at(a_col, j));
341        }
342
343        for j in cpy_size..res_size {
344            BE::reim_zero(res.at_mut(res_col, j));
345        }
346    }
347}
348
349pub fn vec_znx_dft_sub_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
350where
351    BE: Backend<ScalarPrep = f64> + ReimSubInplace,
352    R: VecZnxDftToMut<BE>,
353    A: VecZnxDftToRef<BE>,
354{
355    let a: VecZnxDft<&[u8], BE> = a.to_ref();
356    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
357
358    #[cfg(debug_assertions)]
359    {
360        assert_eq!(a.n(), res.n());
361    }
362
363    let res_size: usize = res.size();
364    let a_size: usize = a.size();
365
366    let sum_size: usize = a_size.min(res_size);
367
368    for j in 0..sum_size {
369        BE::reim_sub_inplace(res.at_mut(res_col, j), a.at(a_col, j));
370    }
371}
372
373pub fn vec_znx_dft_sub_negate_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
374where
375    BE: Backend<ScalarPrep = f64> + ReimSubNegateInplace + ReimNegateInplace,
376    R: VecZnxDftToMut<BE>,
377    A: VecZnxDftToRef<BE>,
378{
379    let a: VecZnxDft<&[u8], BE> = a.to_ref();
380    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
381
382    #[cfg(debug_assertions)]
383    {
384        assert_eq!(a.n(), res.n());
385    }
386
387    let res_size: usize = res.size();
388    let a_size: usize = a.size();
389
390    let sum_size: usize = a_size.min(res_size);
391
392    for j in 0..sum_size {
393        BE::reim_sub_negate_inplace(res.at_mut(res_col, j), a.at(a_col, j));
394    }
395
396    for j in sum_size..res_size {
397        BE::reim_negate_inplace(res.at_mut(res_col, j));
398    }
399}
400
401pub fn vec_znx_dft_zero<R, BE>(res: &mut R)
402where
403    R: VecZnxDftToMut<BE>,
404    BE: Backend<ScalarPrep = f64> + ReimZero,
405{
406    BE::reim_zero(res.to_mut().raw_mut());
407}