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;
  7use std::mem::transmute;
  8
  9
 10impl Ring<u64>{
 11
 12    // Generates a vector storing {X^{2^0}, X^{2^1}, .., X^{2^log_n}}.
 13    pub fn gen_x_pow_2<const NTT: bool, const INV: bool>(&self, log_n: usize) -> Vec<Poly<u64>>{
 14        let mut x_pow: Vec<Poly<u64>> = Vec::<Poly<u64>>::with_capacity(log_n);
 15
 16        (0..log_n).for_each(|i|{
 17            let mut idx: usize = 1<<i;
 18
 19            if INV{
 20                idx = self.n() - idx;
 21            }
 22
 23            x_pow.push(self.new_poly());
 24
 25            if i == 0{
 26                x_pow[i].0[idx] = self.modulus.montgomery.one();
 27                self.ntt_inplace::<false>(&mut x_pow[i]);
 28            }else{
 29                let (left, right) = x_pow.split_at_mut(i);
 30                self.a_mul_b_montgomery_into_c::<ONCE>(&left[i-1], &left[i-1], &mut right[0]);
 31            }
 32        });
 33
 34        if INV{
 35            self.a_neg_into_a::<1, ONCE>(&mut x_pow[0]);
 36        }
 37
 38        if !NTT{
 39            x_pow.iter_mut().for_each(|x| self.intt_inplace::<false>(x));
 40        }
 41
 42        x_pow
 43    }
 44
 45    pub fn pack<'a, const ZEROGARBAGE: bool, const NTT: bool>(&self, polys: &'a mut HashMap<usize, &'a mut Poly<u64>>, log_gap: usize) -> &'a Poly<u64>{
 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 keys: Vec<usize> = polys.keys().copied().collect();
 54        keys.sort();
 55
 56        let mut gap = 0usize;
 57
 58        if keys.len() > 1{
 59            gap = max_pow2_gap(&keys);
 60        }else{
 61            gap = 1<<log_n;
 62        }
 63
 64        let log_gap: usize = gap.log2();
 65
 66        if !ZEROGARBAGE{
 67            if gap > 0 {
 68                log_end -= log_gap;
 69            }
 70        }
 71        
 72        let n_inv: Barrett<u64> = self.modulus.barrett.prepare(self.modulus.inv(1<<(log_end - log_start)));
 73
 74        for (_, poly) in polys.iter_mut() {
 75            if !NTT{
 76                self.ntt_inplace::<true>(poly);
 77            }
 78
 79            self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv,  poly);
 80        }
 81
 82        let x_pow2: Vec<Poly<u64>> = self.gen_x_pow_2::<true, false>(log_n);
 83        let mut tmpa: Poly<u64> = self.new_poly();
 84        let mut tmpb: Poly<u64> = self.new_poly();
 85
 86        for i in log_start..log_end{
 87
 88            let t: usize = 1<<(log_n-1-i);
 89
 90            for j in 0..t{
 91
 92                let option_lo: Option<&&mut Poly<u64>> = polys.get(&i);
 93                let option_hi: Option<&&mut Poly<u64>> = polys.get(&(i+t));
 94                let mut hi_exists: bool = false;
 95
 96                match option_hi{
 97                    Some(hi) =>{
 98
 99                            // Unsafe code is necessary because two mutables references are
100                            // accessed from the map.
101                            unsafe{
102                                self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n-i-1],  transmute(*hi as *const Poly<u64> as *mut Poly<u64>));
103                            }
104
105                            hi_exists = true;
106
107                            match option_lo{
108                                Some(lo) =>{
109
110                                    self.a_sub_b_into_c::<1, ONCE>(lo, hi, &mut tmpa);
111                                    
112                                    // Ensures unsafe blocks are "safe".
113                                    let ptr_hi: *mut Poly<u64> = *hi as *const Poly<u64> as *mut Poly<u64>;
114                                    let ptr_lo: *mut Poly<u64> = *lo as *const Poly<u64> as *mut Poly<u64>;
115                                    assert!(ptr_hi != ptr_lo, "something went seriously wrong");
116
117                                    unsafe{
118                                        self.a_add_b_into_b::<ONCE>(hi, transmute(ptr_lo));
119                                    }
120                                }
121                                None =>{
122                                    unsafe{
123                                        polys.insert(j, transmute(*hi as *const Poly<u64> as *mut Poly<u64>));
124                                    }
125                                },
126                            }
127
128                            polys.remove(&(j+t));
129                        }
130
131                    None =>{},  
132                }
133
134                let option_lo: Option<&&mut Poly<u64>> = polys.get(&i);
135                let option_hi: Option<&&mut Poly<u64>> = polys.get(&(i+t));
136
137                match option_lo{
138                    Some(lo) =>{
139                        
140                        let gal_el: usize = self.galois_element(1<<(i-1), i == 0, log_nth_root);
141
142                        if hi_exists{
143                            self.automorphism::<true>(&tmpa, gal_el, 2<<self.log_n(), &mut tmpb);
144                        }else{
145                            self.automorphism::<true>(*lo, gal_el, nth_root, &mut tmpa);
146                        }
147                        unsafe{
148                            self.a_add_b_into_b::<ONCE>(&tmpa, transmute(*lo as *const Poly<u64> as *mut Poly<u64>));
149                        }
150                    }
151
152                    None =>{
153                        match option_hi{
154                            Some(hi) =>{
155                                let gal_el: usize = self.galois_element(1<<(i-1), i == 0, log_nth_root);
156
157                                self.automorphism::<true>(*hi, gal_el, nth_root, &mut tmpa);
158
159                                unsafe{
160                                    self.a_sub_b_into_a::<1, ONCE>(&tmpa, transmute(*hi as *const Poly<u64> as *mut Poly<u64>))
161                                }
162                            }
163
164                            None =>{}
165                        }
166                    }
167                }
168            }
169        }
170
171        *polys.get(&0).unwrap()
172    }
173}
174
175
176// Returns the largest
177fn max_pow2_gap(vec: &[usize]) -> usize{
178    let mut gap: usize = usize::MAX;
179    for i in 1..vec.len(){
180        let (l, r) = (vec[i-1], vec[i]);
181        assert!(l > r, "invalid input vec: not sorted");
182        gap = min(gap, r-l);
183        if gap == 1{
184            break;
185        }
186    };
187    1 << gap.trailing_zeros()
188}