1use crate::modulus::barrett::Barrett;
  2use crate::modulus::ONCE;
  3use crate::poly::Poly;
  4use crate::ring::Ring;
  5use std::cmp::min;
  6
  7impl Ring<u64> {
  8
  9    pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(
 10        &self,
 11        polys: &mut Vec<Option<Poly<u64>>>,
 12        log_gap: usize,
 13    ) {
 14        let log_n: usize = self.log_n();
 15        let log_nth_root: usize = log_n + 1;
 16        let nth_root: usize = 1 << log_nth_root;
 17        let log_start: usize = log_n - log_gap;
 18        let mut log_end: usize = log_n;
 19
 20        let mut indices: Vec<usize> = Vec::<usize>::new();
 21
 22        // Retrives non-empty indexes
 23        polys.iter().enumerate().for_each(|(i, poly)| {
 24            if !poly.is_none() {
 25                indices.push(i);
 26            }
 27        });
 28
 29        let gap: usize = max_gap(&indices);
 30
 31        if !ZEROGARBAGE {
 32            if gap > 0 {
 33                log_end -= gap.trailing_zeros() as usize;
 34            }
 35        }
 36
 37        assert!(
 38            log_start < log_end,
 39            "invalid input polys: gap between non None value is smaller than 2^log_gap"
 40        );
 41
 42        let n_inv: Barrett<u64> = self
 43            .modulus
 44            .barrett
 45            .prepare(self.modulus.inv(1 << (log_end - log_start)));
 46
 47        indices.iter().for_each(|i| {
 48            if let Some(poly) = polys[*i].as_mut() {
 49                if !NTT {
 50                    self.ntt_inplace::<true>(poly);
 51                }
 52                self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, poly);
 53            }
 54        });
 55
 56        let x_pow2: Vec<Poly<u64>> = self.gen_x_pow_2::<true, false>(log_n);
 57        let mut tmpa: Poly<u64> = self.new_poly();
 58        let mut tmpb: Poly<u64> = self.new_poly();
 59
 60        for i in log_start..log_end {
 61            let t: usize = 1 << (log_n - 1 - i);
 62
 63            let (polys_lo, polys_hi) = polys.split_at_mut(t);
 64
 65            for j in 0..t {
 66                if let Some(poly_hi) = polys_hi[j].as_mut() {
 67                    self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n - i - 1], poly_hi);
 68
 69                    if let Some(poly_lo) = polys_lo[j].as_mut() {
 70                        self.a_sub_b_into_c::<1, ONCE>(poly_lo, poly_hi, &mut tmpa);
 71                        self.a_add_b_into_b::<ONCE>(poly_hi, poly_lo);
 72                    }
 73                }
 74
 75                if let Some(poly_lo) = polys_lo[j].as_mut() {
 76                    let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root);
 77
 78                    if !polys_hi[j].is_none() {
 79                        self.a_apply_automorphism_add_b_into_b::<ONCE, true>(&tmpa, gal_el, 2 << self.log_n(), poly_lo);
 80                    } else {
 81                        self.a_apply_automorphism_into_b::<true>(poly_lo, gal_el, nth_root, &mut tmpa);
 82                        self.a_add_b_into_b::<ONCE>(&tmpa, poly_lo);
 83                    }
 84                } else if let Some(poly_hi) = polys_hi[j].as_mut() {
 85                    let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root);
 86                    self.a_apply_automorphism_into_b::<true>(poly_hi, gal_el, nth_root, &mut tmpa);
 87                    self.a_sub_b_into_a::<1, ONCE>(&tmpa, poly_hi);
 88                    std::mem::swap(&mut polys_lo[j], &mut polys_hi[j]);
 89                }
 90            }
 91
 92            polys.truncate(t);
 93        }
 94
 95        if !NTT {
 96            if let Some(poly) = polys[0].as_mut() {
 97                self.intt_inplace::<false>(poly);
 98            }
 99        }
100    }
101}
102
103
104// Returns the largest gap between two values in an ordered array of distinct values.
105// Panics if the array is not ordered or values are not distincts.
106fn max_gap(vec: &[usize]) -> usize {
107    let mut gap: usize = usize::MAX;
108    for i in 1..vec.len() {
109        let (l, r) = (vec[i - 1], vec[i]);
110        assert!(r > l, "invalid input vec: not sorted or collision between indices");
111        gap = min(gap, r - l);
112        if gap == 1 {
113            break;
114        }
115    }
116    gap
117}