1use crate::modulus::barrett::Barrett;
  2use crate::modulus::{WordOps, ONCE};
  3use crate::poly::Poly;
  4use crate::ring::Ring;
  5use std::cmp::min;
  6use std::collections::HashSet;
  7use std::mem::transmute;
  8
  9impl Ring<u64> {
 10    // Generates a vector storing {X^{2^0}, X^{2^1}, .., X^{2^log_n}}.
 11    pub fn gen_x_pow_2<const NTT: bool, const INV: bool>(&self, log_n: usize) -> Vec<Poly<u64>> {
 12        let mut x_pow: Vec<Poly<u64>> = Vec::<Poly<u64>>::with_capacity(log_n);
 13
 14        (0..log_n).for_each(|i| {
 15            let mut idx: usize = 1 << i;
 16
 17            if INV {
 18                idx = self.n() - idx;
 19            }
 20
 21            x_pow.push(self.new_poly());
 22
 23            if i == 0 {
 24                x_pow[i].0[idx] = self.modulus.montgomery.one();
 25                self.ntt_inplace::<false>(&mut x_pow[i]);
 26            } else {
 27                let (left, right) = x_pow.split_at_mut(i);
 28                self.a_mul_b_montgomery_into_c::<ONCE>(&left[i - 1], &left[i - 1], &mut right[0]);
 29            }
 30        });
 31
 32        if INV {
 33            self.a_neg_into_a::<1, ONCE>(&mut x_pow[0]);
 34        }
 35
 36        if !NTT {
 37            x_pow.iter_mut().for_each(|x| self.intt_inplace::<false>(x));
 38        }
 39
 40        x_pow
 41    }
 42
 43    pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(
 44        &self,
 45        polys: &mut [Option<Poly<u64>>],
 46        log_gap: usize,
 47    ) {
 48        let log_n: usize = self.log_n();
 49        let log_nth_root: usize = log_n + 1;
 50        let nth_root: usize = 1 << log_nth_root;
 51        let log_start: usize = log_n - log_gap;
 52        let mut log_end: usize = log_n;
 53
 54        let mut indices: Vec<usize> = Vec::<usize>::new();
 55
 56        // Retrives non-empty indexes
 57        polys.iter().enumerate().for_each(|(i, poly)| {
 58            if Some(poly) != None {
 59                indices.push(i);
 60            }
 61        });
 62
 63        let gap: usize = max_gap(&indices);
 64
 65        let set: HashSet<_> = indices.into_iter().collect();
 66
 67        let max_pow2_gap_divisor: usize = 1 << gap.trailing_zeros();
 68
 69        if !ZEROGARBAGE {
 70            if gap > 0 {
 71                log_end -= max_pow2_gap_divisor;
 72            }
 73        }
 74
 75        let n_inv: Barrett<u64> = self
 76            .modulus
 77            .barrett
 78            .prepare(self.modulus.inv(1 << (log_end - log_start)));
 79
 80        set.iter().for_each(|i| {
 81            if let Some(poly) = polys[*i].as_mut() {
 82                if !NTT {
 83                    self.ntt_inplace::<true>(poly);
 84                }
 85                self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, poly);
 86            }
 87        });
 88
 89        let x_pow2: Vec<Poly<u64>> = self.gen_x_pow_2::<true, false>(log_n);
 90        let mut tmpa: Poly<u64> = self.new_poly();
 91        let mut tmpb: Poly<u64> = self.new_poly();
 92
 93        for i in log_start..log_end {
 94            let t: usize = 1 << (log_n - 1 - i);
 95
 96            let (polys_lo, polys_hi) = polys.split_at_mut(t);
 97
 98            for j in 0..t {
 99                if let Some(poly_hi) = polys_hi[j].as_mut() {
100                    self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n - i - 1], poly_hi);
101
102                    if let Some(poly_lo) = polys_lo[j].as_mut() {
103                        self.a_sub_b_into_c::<1, ONCE>(poly_lo, poly_hi, &mut tmpa);
104                        self.a_add_b_into_b::<ONCE>(poly_hi, poly_lo);
105                    } else {
106                        std::mem::swap(&mut polys_lo[j], &mut polys_hi[j]);
107                    }
108                }
109
110                if let Some(poly_lo) = polys_lo[j].as_mut() {
111                    let gal_el: usize = self.galois_element(1 << (i - 1), i == 0, log_nth_root);
112
113                    if !polys_hi[j].is_none() {
114                        self.automorphism::<true>(&tmpa, gal_el, 2 << self.log_n(), &mut tmpb);
115                    } else {
116                        self.automorphism::<true>(poly_lo, gal_el, nth_root, &mut tmpa);
117                    }
118
119                    self.a_add_b_into_b::<ONCE>(&tmpa, poly_lo);
120                } else if let Some(poly_hi) = polys_hi[j].as_mut() {
121                    let gal_el: usize = self.galois_element(1 << (i - 1), i == 0, log_nth_root);
122
123                    self.automorphism::<true>(poly_hi, gal_el, nth_root, &mut tmpa);
124                    self.a_sub_b_into_a::<1, ONCE>(&tmpa, poly_hi)
125                }
126            }
127        }
128    }
129}
130
131// Returns the largest gap.
132fn max_gap(vec: &[usize]) -> usize {
133    let mut gap: usize = usize::MAX;
134    for i in 1..vec.len() {
135        let (l, r) = (vec[i - 1], vec[i]);
136        assert!(l > r, "invalid input vec: not sorted");
137        gap = min(gap, r - l);
138        if gap == 1 {
139            break;
140        }
141    }
142    gap
143}