1use crate::modulus::barrett::Barrett;
  2use crate::modulus::montgomery::Montgomery;
  3use crate::modulus::ONCE;
  4use crate::poly::Poly;
  5use crate::ring::Ring;
  6use std::cmp::min;
  7use std::rc::Rc;
  8
  9impl Ring<u64> {
 10
 11    pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(
 12        &self,
 13        polys: &mut Vec<Option<Poly<u64>>>,
 14        log_gap: usize,
 15    ) {
 16        let log_n: usize = self.log_n();
 17        let log_nth_root: usize = log_n + 1;
 18        let nth_root: usize = 1 << log_nth_root;
 19        let log_start: usize = log_n - log_gap;
 20        let mut log_end: usize = log_n;
 21
 22        let mut indices: Vec<usize> = Vec::<usize>::new();
 23
 24        // Retrives non-empty indexes
 25        polys.iter().enumerate().for_each(|(i, poly)| {
 26            if !poly.is_none() {
 27                indices.push(i);
 28            }
 29        });
 30
 31        let gap: usize = max_gap(&indices);
 32
 33        if !ZEROGARBAGE {
 34            if gap > 0 {
 35                log_end -= gap.trailing_zeros() as usize;
 36            }
 37        }
 38
 39        assert!(
 40            log_start < log_end,
 41            "invalid input polys: gap between non None value is smaller than 2^log_gap"
 42        );
 43
 44        let n_inv: Barrett<u64> = self
 45            .modulus
 46            .barrett
 47            .prepare(self.modulus.inv(1 << (log_end - log_start)));
 48
 49        indices.iter().for_each(|i| {
 50            if let Some(poly) = polys[*i].as_mut() {
 51                if !NTT {
 52                    self.ntt_inplace::<true>(poly);
 53                }
 54                self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, poly);
 55            }
 56        });
 57
 58        let x_pow2: Vec<Poly<Montgomery<u64>>> = self.gen_x_pow_2::<true, false>(log_n);
 59        let mut tmpa: Poly<u64> = self.new_poly();
 60        let mut tmpb: Poly<u64> = self.new_poly();
 61
 62        for i in log_start..log_end {
 63            let t: usize = 1 << (log_n - 1 - i);
 64
 65            let (polys_lo, polys_hi) = polys.split_at_mut(t);
 66
 67            for j in 0..t {
 68                if let Some(poly_hi) = polys_hi[j].as_mut() {
 69                    self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n - i - 1], poly_hi);
 70
 71                    if let Some(poly_lo) = polys_lo[j].as_mut() {
 72                        self.a_sub_b_into_c::<1, ONCE>(poly_lo, poly_hi, &mut tmpa);
 73                        self.a_add_b_into_b::<ONCE>(poly_hi, poly_lo);
 74                    }
 75                }
 76
 77                if let Some(poly_lo) = polys_lo[j].as_mut() {
 78                    let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root);
 79
 80                    if !polys_hi[j].is_none() {
 81                        self.a_apply_automorphism_add_b_into_b::<ONCE, true>(&tmpa, gal_el, 2 << self.log_n(), poly_lo);
 82                    } else {
 83                        self.a_apply_automorphism_into_b::<true>(poly_lo, gal_el, nth_root, &mut tmpa);
 84                        self.a_add_b_into_b::<ONCE>(&tmpa, poly_lo);
 85                    }
 86                } else if let Some(poly_hi) = polys_hi[j].as_mut() {
 87                    let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root);
 88                    self.a_apply_automorphism_into_b::<true>(poly_hi, gal_el, nth_root, &mut tmpa);
 89                    self.a_sub_b_into_a::<1, ONCE>(&tmpa, poly_hi);
 90                    std::mem::swap(&mut polys_lo[j], &mut polys_hi[j]);
 91                }
 92            }
 93
 94            polys.truncate(t);
 95        }
 96
 97        if !NTT {
 98            if let Some(poly) = polys[0].as_mut() {
 99                self.intt_inplace::<false>(poly);
100            }
101        }
102    }
103}
104
105
106// Returns the largest gap between two values in an ordered array of distinct values.
107// Panics if the array is not ordered or values are not distincts.
108fn max_gap(vec: &[usize]) -> usize {
109    let mut gap: usize = usize::MAX;
110    for i in 1..vec.len() {
111        let (l, r) = (vec[i - 1], vec[i]);
112        assert!(r > l, "invalid input vec: not sorted or collision between indices");
113        gap = min(gap, r - l);
114        if gap == 1 {
115            break;
116        }
117    }
118    gap
119}
120
121
122pub struct StreamRepacker{
123    accumulators: Vec<Accumulator>,
124    buf0: Poly<u64>,
125    buf1: Poly<u64>,
126    buf_auto: Poly<u64>,
127    x_pow_2: Vec<Poly<Montgomery<u64>>>,
128    n_inv: Barrett<u64>,
129    counter: usize,
130}
131
132pub struct Accumulator{
133    buf: [Option<Rc<Poly<u64>>>; 2],
134    control: bool,
135}
136
137impl Accumulator{
138    pub fn new(r: &Ring<u64>) -> Self{
139        Self { buf: [Some(Rc::new(r.new_poly())), None], control: false }
140    }
141}
142
143impl StreamRepacker{
144    pub fn new(r: &Ring<u64>) -> Self{
145
146        let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
147
148        (0..r.log_n()).for_each(|_|
149            accumulators.push(Accumulator::new(r))
150        );
151
152        Self{
153            accumulators: accumulators,
154            buf0: r.new_poly(),
155            buf1: r.new_poly(),
156            buf_auto: r.new_poly(),
157            x_pow_2: r.gen_x_pow_2::<true, false>(r.log_n()),
158            n_inv: r.modulus.barrett.prepare(r.modulus.inv(r.n() as u64)),
159            counter:0,
160        }
161    }
162
163    fn merge_ab(&mut self, r: &Ring<u64>, a: &Poly<u64>, b: &Poly<u64>, i: usize) -> &Poly<u64>{
164
165        let tmp_a: &mut Poly<u64> = &mut self.buf0;
166        let tmp_b: &mut Poly<u64> = &mut self.buf1;
167
168        r.a_mul_b_montgomery_into_c::<ONCE>(a, &self.x_pow_2[r.log_n()-i-1], tmp_a);
169        r.a_sub_b_into_c::<1, ONCE>(a, tmp_a, tmp_b);
170        r.a_add_b_into_b::<ONCE>(a, tmp_a);
171
172        if i == 0{
173            r.a_mul_b_scalar_barrett_into_a::<ONCE>(&self.n_inv, tmp_a);
174            r.a_mul_b_scalar_barrett_into_a::<ONCE>(&self.n_inv, tmp_b);
175        }
176
177        let log_nth_root = r.log_n()+1;
178        let nth_root = 1<<log_nth_root;
179            
180        let gal_el: usize = r.galois_element((1 << i) >> 1, i == 0, log_nth_root);
181
182        r.a_apply_automorphism_add_b_into_b::<ONCE, true>(tmp_b, gal_el, nth_root, tmp_a);
183
184        tmp_a
185    }
186
187    fn merge_a(&mut self, r: &Ring<u64>, a: &Poly<u64>, i: usize) -> &Poly<u64>{
188
189        let tmp_a: &mut Poly<u64> = &mut self.buf0;
190 
191        let log_nth_root = r.log_n()+1;
192        let nth_root = 1<<log_nth_root;
193        let gal_el: usize = r.galois_element((1 << i) >> 1, i == 0, log_nth_root);
194
195        if i == 0{
196            r.a_mul_b_scalar_barrett_into_a::<ONCE>(&self.n_inv, tmp_a);
197            r.a_apply_automorphism_into_b::<true>(tmp_a, gal_el, nth_root, &mut self.buf_auto)
198            r.a_add_b_into_c::<ONCE>(&self.buf_auto, a, tmp_a);
199        }else{
200            r.a_apply_automorphism_into_b::<true>(a, gal_el, nth_root, tmp_a);
201            r.a_add_b_into_b::<ONCE>(a, tmp_a);
202        }
203
204        tmp_a
205    }
206
207    fn merge_b(&mut self, r: &Ring<u64>, b: &Poly<u64>, i: usize) -> &Poly<u64>{
208
209        let tmp_a: &mut Poly<u64> = &mut self.buf0;
210        let tmp_b: &mut Poly<u64> = &mut self.buf1;
211
212        let log_nth_root = r.log_n()+1;
213        let nth_root = 1<<log_nth_root;
214        let gal_el: usize = r.galois_element((1 << i) >> 1, i == 0, log_nth_root);
215
216        if i == 0{
217            r.a_mul_b_scalar_barrett_into_c::<ONCE>(&self.n_inv, b, tmp_b);
218            r.a_mul_b_montgomery_into_a::<ONCE>(&self.x_pow_2[r.log_n()-i-1], tmp_b);
219        }else{
220            r.a_mul_b_montgomery_into_c::<ONCE>(b, &self.x_pow_2[r.log_n()-i-1], tmp_b);
221        }
222
223        r.a_apply_automorphism_into_b::<true>(tmp_b, gal_el, nth_root, &mut self.buf_auto);
224        r.a_sub_b_into_a::<1, ONCE>(&self.buf_auto, tmp_b);
225
226        tmp_b
227    }
228}