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