1use std::collections::HashMap;
 2use crate::poly::Poly;
 3use crate::ring::Ring;
 4use crate::modulus::{ONCE, WordOps};
 5use crate::modulus::barrett::Barrett;
 6use std::cmp::min;
 7
 8
 9impl Ring<u64>{
10
11    // Generates a vector storing {X^{2^0}, X^{2^1}, .., X^{2^log_n}}.
12    pub fn gen_x_pow_2<const NTT: bool, const INV: bool>(&self, log_n: usize) -> Vec<Poly<u64>>{
13        let mut x_pow: Vec<Poly<u64>> = Vec::<Poly<u64>>::with_capacity(log_n);
14
15        (0..log_n).for_each(|i|{
16            let mut idx: usize = 1<<i;
17
18            if INV{
19                idx = self.n() - idx;
20            }
21
22            x_pow.push(self.new_poly());
23
24            if i == 0{
25                x_pow[i].0[idx] = self.modulus.montgomery.one();
26                self.ntt_inplace::<false>(&mut x_pow[i]);
27            }else{
28                let (left, right) = x_pow.split_at_mut(i);
29                self.a_mul_b_montgomery_into_c::<ONCE>(&left[i-1], &left[i-1], &mut right[0]);
30            }
31        });
32
33        if INV{
34            self.a_neg_into_a::<1, ONCE>(&mut x_pow[0]);
35        }
36
37        if !NTT{
38            x_pow.iter_mut().for_each(|x| self.intt_inplace::<false>(x));
39        }
40
41        x_pow
42    }
43
44    pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(&self, polys: &mut HashMap<usize, Poly<u64>>, log_gap: usize) -> Poly<u64>{
45        
46        let log_n = self.log_n();
47        let log_start = log_n - log_gap;
48        let mut log_end = log_n;
49
50        let mut keys: Vec<usize> = polys.keys().copied().collect();
51        keys.sort();
52
53        let mut gap = 0usize;
54
55        if keys.len() > 1{
56            gap = max_pow2_gap(&keys);
57        }else{
58            gap = 1<<log_n;
59        }
60
61        let log_gap: usize = gap.log2();
62
63        if !ZEROGARBAGE{
64            if gap > 0 {
65                log_end -= log_gap;
66            }
67        }
68        
69        let n_inv: Barrett<u64> = self.modulus.barrett.prepare(self.modulus.inv(1<<(log_end - log_start)));
70
71        for (_, poly) in polys.iter_mut() {
72            if !NTT{
73                self.ntt_inplace::<true>(poly);
74            }
75
76            self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv,  poly);
77        }
78
79        Poly::<u64>::default()
80    }
81}
82
83// Returns the largest
84fn max_pow2_gap(vec: &[usize]) -> usize{
85    let mut gap: usize = usize::MAX;
86    for i in 1..vec.len(){
87        let (l, r) = (vec[i-1], vec[i]);
88        assert!(l > r, "invalid input vec: not sorted");
89        gap = min(gap, r-l);
90        if gap == 1{
91            break;
92        }
93    };
94    1 << gap.trailing_zeros()
95}