1use std::f64::consts::SQRT_2;
  2
  3use crate::{
  4    api::VecZnxBigAddNormal,
  5    layouts::{
  6        Backend, Module, VecZnx, VecZnxBig, VecZnxBigToMut, VecZnxBigToRef, VecZnxToMut, VecZnxToRef, ZnxView, ZnxViewMut,
  7    },
  8    oep::VecZnxBigAllocBytesImpl,
  9    reference::{
 10        vec_znx::{
 11            vec_znx_add, vec_znx_add_inplace, vec_znx_automorphism, vec_znx_automorphism_inplace, vec_znx_negate,
 12            vec_znx_negate_inplace, vec_znx_normalize, vec_znx_normalize_tmp_bytes, vec_znx_sub, vec_znx_sub_inplace,
 13            vec_znx_sub_negate_inplace,
 14        },
 15        znx::{
 16            ZnxAdd, ZnxAddInplace, ZnxAutomorphism, ZnxCopy, ZnxExtractDigitAddMul, ZnxMulPowerOfTwoInplace, ZnxNegate,
 17            ZnxNegateInplace, ZnxNormalizeDigit, ZnxNormalizeFinalStep, ZnxNormalizeFinalStepInplace, ZnxNormalizeFirstStep,
 18            ZnxNormalizeFirstStepCarryOnly, ZnxNormalizeMiddleStep, ZnxNormalizeMiddleStepCarryOnly,
 19            ZnxNormalizeMiddleStepInplace, ZnxSub, ZnxSubInplace, ZnxSubNegateInplace, ZnxZero, znx_add_normal_f64_ref,
 20        },
 21    },
 22    source::Source,
 23};
 24
 25pub fn vec_znx_big_add<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 26where
 27    BE: Backend<ScalarBig = i64> + ZnxAdd + ZnxCopy + ZnxZero,
 28    R: VecZnxBigToMut<BE>,
 29    A: VecZnxBigToRef<BE>,
 30    B: VecZnxBigToRef<BE>,
 31{
 32    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
 33    let a: VecZnxBig<&[u8], BE> = a.to_ref();
 34    let b: VecZnxBig<&[u8], BE> = b.to_ref();
 35
 36    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
 37        data: res.data,
 38        n: res.n,
 39        cols: res.cols,
 40        size: res.size,
 41        max_size: res.max_size,
 42    };
 43
 44    let a_vznx: VecZnx<&[u8]> = VecZnx {
 45        data: a.data,
 46        n: a.n,
 47        cols: a.cols,
 48        size: a.size,
 49        max_size: a.max_size,
 50    };
 51
 52    let b_vznx: VecZnx<&[u8]> = VecZnx {
 53        data: b.data,
 54        n: b.n,
 55        cols: b.cols,
 56        size: b.size,
 57        max_size: b.max_size,
 58    };
 59
 60    vec_znx_add::<_, _, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col, &b_vznx, b_col);
 61}
 62
 63pub fn vec_znx_big_add_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
 64where
 65    BE: Backend<ScalarBig = i64> + ZnxAddInplace,
 66    R: VecZnxBigToMut<BE>,
 67    A: VecZnxBigToRef<BE>,
 68{
 69    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
 70    let a: VecZnxBig<&[u8], BE> = a.to_ref();
 71
 72    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
 73        data: res.data,
 74        n: res.n,
 75        cols: res.cols,
 76        size: res.size,
 77        max_size: res.max_size,
 78    };
 79
 80    let a_vznx: VecZnx<&[u8]> = VecZnx {
 81        data: a.data,
 82        n: a.n,
 83        cols: a.cols,
 84        size: a.size,
 85        max_size: a.max_size,
 86    };
 87
 88    vec_znx_add_inplace::<_, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col);
 89}
 90
 91pub fn vec_znx_big_add_small<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 92where
 93    BE: Backend<ScalarBig = i64> + ZnxAdd + ZnxCopy + ZnxZero,
 94    R: VecZnxBigToMut<BE>,
 95    A: VecZnxBigToRef<BE>,
 96    B: VecZnxToRef,
 97{
 98    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
 99    let a: VecZnxBig<&[u8], BE> = a.to_ref();
100
101    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
102        data: res.data,
103        n: res.n,
104        cols: res.cols,
105        size: res.size,
106        max_size: res.max_size,
107    };
108
109    let a_vznx: VecZnx<&[u8]> = VecZnx {
110        data: a.data,
111        n: a.n,
112        cols: a.cols,
113        size: a.size,
114        max_size: a.max_size,
115    };
116
117    vec_znx_add::<_, _, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col, b, b_col);
118}
119
120pub fn vec_znx_big_add_small_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
121where
122    BE: Backend<ScalarBig = i64> + ZnxAddInplace,
123    R: VecZnxBigToMut<BE>,
124    A: VecZnxToRef,
125{
126    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
127    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
128        data: res.data,
129        n: res.n,
130        cols: res.cols,
131        size: res.size,
132        max_size: res.max_size,
133    };
134
135    vec_znx_add_inplace::<_, _, BE>(&mut res_vznx, res_col, a, a_col);
136}
137
138pub fn vec_znx_big_automorphism_inplace_tmp_bytes(n: usize) -> usize {
139    n * size_of::<i64>()
140}
141
142pub fn vec_znx_big_automorphism<R, A, BE>(p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
143where
144    BE: Backend<ScalarBig = i64> + ZnxAutomorphism + ZnxZero,
145    R: VecZnxBigToMut<BE>,
146    A: VecZnxBigToRef<BE>,
147{
148    let res: VecZnxBig<&mut [u8], _> = res.to_mut();
149    let a: VecZnxBig<&[u8], _> = a.to_ref();
150
151    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
152        data: res.data,
153        n: res.n,
154        cols: res.cols,
155        size: res.size,
156        max_size: res.max_size,
157    };
158
159    let a_vznx: VecZnx<&[u8]> = VecZnx {
160        data: a.data,
161        n: a.n,
162        cols: a.cols,
163        size: a.size,
164        max_size: a.max_size,
165    };
166
167    vec_znx_automorphism::<_, _, BE>(p, &mut res_vznx, res_col, &a_vznx, a_col);
168}
169
170pub fn vec_znx_big_automorphism_inplace<R, BE>(p: i64, res: &mut R, res_col: usize, tmp: &mut [i64])
171where
172    BE: Backend<ScalarBig = i64> + ZnxAutomorphism + ZnxCopy,
173    R: VecZnxBigToMut<BE>,
174{
175    let res: VecZnxBig<&mut [u8], _> = res.to_mut();
176
177    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
178        data: res.data,
179        n: res.n,
180        cols: res.cols,
181        size: res.size,
182        max_size: res.max_size,
183    };
184
185    vec_znx_automorphism_inplace::<_, BE>(p, &mut res_vznx, res_col, tmp);
186}
187
188pub fn vec_znx_big_negate<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
189where
190    BE: Backend<ScalarBig = i64> + ZnxNegate + ZnxZero,
191    R: VecZnxBigToMut<BE>,
192    A: VecZnxBigToRef<BE>,
193{
194    let res: VecZnxBig<&mut [u8], _> = res.to_mut();
195    let a: VecZnxBig<&[u8], _> = a.to_ref();
196
197    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
198        data: res.data,
199        n: res.n,
200        cols: res.cols,
201        size: res.size,
202        max_size: res.max_size,
203    };
204
205    let a_vznx: VecZnx<&[u8]> = VecZnx {
206        data: a.data,
207        n: a.n,
208        cols: a.cols,
209        size: a.size,
210        max_size: a.max_size,
211    };
212
213    vec_znx_negate::<_, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col);
214}
215
216pub fn vec_znx_big_negate_inplace<R, BE>(res: &mut R, res_col: usize)
217where
218    BE: Backend<ScalarBig = i64> + ZnxNegateInplace,
219    R: VecZnxBigToMut<BE>,
220{
221    let res: VecZnxBig<&mut [u8], _> = res.to_mut();
222
223    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
224        data: res.data,
225        n: res.n,
226        cols: res.cols,
227        size: res.size,
228        max_size: res.max_size,
229    };
230
231    vec_znx_negate_inplace::<_, BE>(&mut res_vznx, res_col);
232}
233
234pub fn vec_znx_big_normalize_tmp_bytes(n: usize) -> usize {
235    vec_znx_normalize_tmp_bytes(n)
236}
237
238#[allow(clippy::too_many_arguments)]
239pub fn vec_znx_big_normalize<R, A, BE>(
240    res: &mut R,
241    res_base2k: usize,
242    res_offset: i64,
243    res_col: usize,
244    a: &A,
245    a_base2k: usize,
246    a_col: usize,
247    carry: &mut [i64],
248) where
249    R: VecZnxToMut,
250    A: VecZnxBigToRef<BE>,
251    BE: Backend<ScalarBig = i64>
252        + ZnxZero
253        + ZnxCopy
254        + ZnxAddInplace
255        + ZnxMulPowerOfTwoInplace
256        + ZnxNormalizeFirstStepCarryOnly
257        + ZnxNormalizeMiddleStepCarryOnly
258        + ZnxNormalizeMiddleStep
259        + ZnxNormalizeFinalStep
260        + ZnxNormalizeFirstStep
261        + ZnxExtractDigitAddMul
262        + ZnxNormalizeDigit
263        + ZnxNormalizeMiddleStepInplace
264        + ZnxNormalizeFinalStepInplace,
265{
266    let a: VecZnxBig<&[u8], _> = a.to_ref();
267    let a_vznx: VecZnx<&[u8]> = VecZnx {
268        data: a.data,
269        n: a.n,
270        cols: a.cols,
271        size: a.size,
272        max_size: a.max_size,
273    };
274
275    vec_znx_normalize::<_, _, BE>(res, res_base2k, res_offset, res_col, &a_vznx, a_base2k, a_col, carry);
276}
277
278pub fn vec_znx_big_add_normal_ref<R, B: Backend<ScalarBig = i64>>(
279    base2k: usize,
280    res: &mut R,
281    res_col: usize,
282    k: usize,
283    sigma: f64,
284    bound: f64,
285    source: &mut Source,
286) where
287    R: VecZnxBigToMut<B>,
288{
289    let mut res: VecZnxBig<&mut [u8], B> = res.to_mut();
290    assert!(
291        (bound.log2().ceil() as i64) < 64,
292        "invalid bound: ceil(log2(bound))={} > 63",
293        (bound.log2().ceil() as i64)
294    );
295
296    let limb: usize = k.div_ceil(base2k) - 1;
297    let scale: f64 = (1 << ((limb + 1) * base2k - k)) as f64;
298    znx_add_normal_f64_ref(res.at_mut(res_col, limb), sigma * scale, bound * scale, source)
299}
300
301pub fn test_vec_znx_big_add_normal<B>(module: &Module<B>)
302where
303    Module<B>: VecZnxBigAddNormal<B>,
304    B: Backend<ScalarBig = i64> + VecZnxBigAllocBytesImpl,
305{
306    let n: usize = module.n();
307    let base2k: usize = 17;
308    let k: usize = 2 * 17;
309    let size: usize = 5;
310    let sigma: f64 = 3.2;
311    let bound: f64 = 6.0 * sigma;
312    let mut source: Source = Source::new([0u8; 32]);
313    let cols: usize = 2;
314    let zero: Vec<i64> = vec![0; n];
315    let k_f64: f64 = (1u64 << k as u64) as f64;
316    let sqrt2: f64 = SQRT_2;
317    (0..cols).for_each(|col_i| {
318        let mut a: VecZnxBig<Vec<u8>, B> = VecZnxBig::alloc(n, cols, size);
319        module.vec_znx_big_add_normal(base2k, &mut a, col_i, k, &mut source, sigma, bound);
320        module.vec_znx_big_add_normal(base2k, &mut a, col_i, k, &mut source, sigma, bound);
321        (0..cols).for_each(|col_j| {
322            if col_j != col_i {
323                (0..size).for_each(|limb_i| {
324                    assert_eq!(a.at(col_j, limb_i), zero);
325                })
326            } else {
327                let std: f64 = a.stats(base2k, col_i).std() * k_f64;
328                assert!((std - sigma * sqrt2).abs() < 0.1, "std={} ~!= {}", std, sigma * sqrt2);
329            }
330        })
331    });
332}
333
334/// R <- A - B
335pub fn vec_znx_big_sub<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
336where
337    BE: Backend<ScalarBig = i64> + ZnxSub + ZnxNegate + ZnxZero + ZnxCopy,
338    R: VecZnxBigToMut<BE>,
339    A: VecZnxBigToRef<BE>,
340    B: VecZnxBigToRef<BE>,
341{
342    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
343    let a: VecZnxBig<&[u8], BE> = a.to_ref();
344    let b: VecZnxBig<&[u8], BE> = b.to_ref();
345
346    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
347        data: res.data,
348        n: res.n,
349        cols: res.cols,
350        size: res.size,
351        max_size: res.max_size,
352    };
353
354    let a_vznx: VecZnx<&[u8]> = VecZnx {
355        data: a.data,
356        n: a.n,
357        cols: a.cols,
358        size: a.size,
359        max_size: a.max_size,
360    };
361
362    let b_vznx: VecZnx<&[u8]> = VecZnx {
363        data: b.data,
364        n: b.n,
365        cols: b.cols,
366        size: b.size,
367        max_size: b.max_size,
368    };
369
370    vec_znx_sub::<_, _, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col, &b_vznx, b_col);
371}
372
373/// R <- A - B
374pub fn vec_znx_big_sub_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
375where
376    BE: Backend<ScalarBig = i64> + ZnxSubInplace,
377    R: VecZnxBigToMut<BE>,
378    A: VecZnxBigToRef<BE>,
379{
380    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
381    let a: VecZnxBig<&[u8], BE> = a.to_ref();
382
383    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
384        data: res.data,
385        n: res.n,
386        cols: res.cols,
387        size: res.size,
388        max_size: res.max_size,
389    };
390
391    let a_vznx: VecZnx<&[u8]> = VecZnx {
392        data: a.data,
393        n: a.n,
394        cols: a.cols,
395        size: a.size,
396        max_size: a.max_size,
397    };
398
399    vec_znx_sub_inplace::<_, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col);
400}
401
402/// R <- B - A
403pub fn vec_znx_big_sub_negate_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
404where
405    BE: Backend<ScalarBig = i64> + ZnxSubNegateInplace + ZnxNegateInplace,
406    R: VecZnxBigToMut<BE>,
407    A: VecZnxBigToRef<BE>,
408{
409    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
410    let a: VecZnxBig<&[u8], BE> = a.to_ref();
411
412    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
413        data: res.data,
414        n: res.n,
415        cols: res.cols,
416        size: res.size,
417        max_size: res.max_size,
418    };
419
420    let a_vznx: VecZnx<&[u8]> = VecZnx {
421        data: a.data,
422        n: a.n,
423        cols: a.cols,
424        size: a.size,
425        max_size: a.max_size,
426    };
427
428    vec_znx_sub_negate_inplace::<_, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col);
429}
430
431/// R <- A - B
432pub fn vec_znx_big_sub_small_a<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
433where
434    BE: Backend<ScalarBig = i64> + ZnxSub + ZnxNegate + ZnxZero + ZnxCopy,
435    R: VecZnxBigToMut<BE>,
436    A: VecZnxToRef,
437    B: VecZnxBigToRef<BE>,
438{
439    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
440    let b: VecZnxBig<&[u8], BE> = b.to_ref();
441
442    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
443        data: res.data,
444        n: res.n,
445        cols: res.cols,
446        size: res.size,
447        max_size: res.max_size,
448    };
449
450    let b_vznx: VecZnx<&[u8]> = VecZnx {
451        data: b.data,
452        n: b.n,
453        cols: b.cols,
454        size: b.size,
455        max_size: b.max_size,
456    };
457
458    vec_znx_sub::<_, _, _, BE>(&mut res_vznx, res_col, a, a_col, &b_vznx, b_col);
459}
460
461/// R <- A - B
462pub fn vec_znx_big_sub_small_b<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
463where
464    BE: Backend<ScalarBig = i64> + ZnxSub + ZnxNegate + ZnxZero + ZnxCopy,
465    R: VecZnxBigToMut<BE>,
466    A: VecZnxBigToRef<BE>,
467    B: VecZnxToRef,
468{
469    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
470    let a: VecZnxBig<&[u8], BE> = a.to_ref();
471
472    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
473        data: res.data,
474        n: res.n,
475        cols: res.cols,
476        size: res.size,
477        max_size: res.max_size,
478    };
479
480    let a_vznx: VecZnx<&[u8]> = VecZnx {
481        data: a.data,
482        n: a.n,
483        cols: a.cols,
484        size: a.size,
485        max_size: a.max_size,
486    };
487
488    vec_znx_sub::<_, _, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col, b, b_col);
489}
490
491///  R <- R - A
492pub fn vec_znx_big_sub_small_a_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
493where
494    BE: Backend<ScalarBig = i64> + ZnxSubInplace,
495    R: VecZnxBigToMut<BE>,
496    A: VecZnxToRef,
497{
498    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
499
500    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
501        data: res.data,
502        n: res.n,
503        cols: res.cols,
504        size: res.size,
505        max_size: res.max_size,
506    };
507
508    vec_znx_sub_inplace::<_, _, BE>(&mut res_vznx, res_col, a, a_col);
509}
510
511/// R <- A - R
512pub fn vec_znx_big_sub_small_b_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
513where
514    BE: Backend<ScalarBig = i64> + ZnxSubNegateInplace + ZnxNegateInplace,
515    R: VecZnxBigToMut<BE>,
516    A: VecZnxToRef,
517{
518    let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
519
520    let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
521        data: res.data,
522        n: res.n,
523        cols: res.cols,
524        size: res.size,
525        max_size: res.max_size,
526    };
527
528    vec_znx_sub_negate_inplace::<_, _, BE>(&mut res_vznx, res_col, a, a_col);
529}