1use itertools::izip;
  2use rug::{Assign, Float};
  3
  4use crate::{
  5    layouts::{DataMut, DataRef, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut},
  6    reference::znx::{
  7        ZnxNormalizeFinalStepInplace, ZnxNormalizeFirstStepInplace, ZnxNormalizeMiddleStepInplace, ZnxRef, ZnxZero,
  8        get_carry_i128, get_digit_i128, znx_zero_ref,
  9    },
 10};
 11
 12impl<D: DataMut> VecZnx<D> {
 13    pub fn encode_vec_i64(&mut self, base2k: usize, col: usize, k: usize, data: &[i64]) {
 14        let size: usize = k.div_ceil(base2k);
 15
 16        #[cfg(debug_assertions)]
 17        {
 18            let a: VecZnx<&mut [u8]> = self.to_mut();
 19            assert!(
 20                size <= a.size(),
 21                "invalid argument k.div_ceil(base2k)={} > a.size()={}",
 22                size,
 23                a.size()
 24            );
 25            assert!(col < a.cols());
 26            assert!(data.len() == a.n())
 27        }
 28
 29        let mut a: VecZnx<&mut [u8]> = self.to_mut();
 30        let a_size: usize = a.size();
 31
 32        // Zeroes coefficients of the i-th column
 33        for i in 0..a_size {
 34            znx_zero_ref(a.at_mut(col, i));
 35        }
 36
 37        // Copies the data on the correct limb
 38        a.at_mut(col, size - 1).copy_from_slice(data);
 39
 40        let mut carry: Vec<i64> = vec![0i64; a.n()];
 41        let k_rem: usize = (base2k - (k % base2k)) % base2k;
 42
 43        // Normalizes and shift if necessary.
 44        for j in (0..size).rev() {
 45            if j == size - 1 {
 46                ZnxRef::znx_normalize_first_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
 47            } else if j == 0 {
 48                ZnxRef::znx_normalize_final_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
 49            } else {
 50                ZnxRef::znx_normalize_middle_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
 51            }
 52        }
 53    }
 54
 55    pub fn encode_vec_i128(&mut self, base2k: usize, col: usize, k: usize, data: &[i128]) {
 56        let size: usize = k.div_ceil(base2k);
 57
 58        #[cfg(debug_assertions)]
 59        {
 60            let a: VecZnx<&mut [u8]> = self.to_mut();
 61            assert!(
 62                size <= a.size(),
 63                "invalid argument k.div_ceil(base2k)={} > a.size()={}",
 64                size,
 65                a.size()
 66            );
 67            assert!(col < a.cols());
 68            assert!(data.len() == a.n())
 69        }
 70
 71        let mut a: VecZnx<&mut [u8]> = self.to_mut();
 72        let a_size: usize = a.size();
 73
 74        {
 75            let mut carry_i128: Vec<i128> = vec![0i128; a.n()];
 76            carry_i128.copy_from_slice(data);
 77
 78            for j in (0..size).rev() {
 79                for (x, a) in izip!(a.at_mut(col, j).iter_mut(), carry_i128.iter_mut()) {
 80                    let digit: i128 = get_digit_i128(base2k, *a);
 81                    let carry: i128 = get_carry_i128(base2k, *a, digit);
 82                    *x = digit as i64;
 83                    *a = carry;
 84                }
 85            }
 86        }
 87
 88        for j in size..a_size {
 89            ZnxRef::znx_zero(a.at_mut(col, j));
 90        }
 91
 92        let mut carry: Vec<i64> = vec![0i64; a.n()];
 93        let k_rem: usize = (base2k - (k % base2k)) % base2k;
 94
 95        for j in (0..size).rev() {
 96            if j == a_size - 1 {
 97                ZnxRef::znx_normalize_first_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
 98            } else if j == 0 {
 99                ZnxRef::znx_normalize_final_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
100            } else {
101                ZnxRef::znx_normalize_middle_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
102            }
103        }
104    }
105
106    pub fn encode_coeff_i64(&mut self, base2k: usize, col: usize, k: usize, idx: usize, data: i64) {
107        let size: usize = k.div_ceil(base2k);
108
109        #[cfg(debug_assertions)]
110        {
111            let a: VecZnx<&mut [u8]> = self.to_mut();
112            assert!(idx < a.n());
113            assert!(
114                size <= a.size(),
115                "invalid argument k.div_ceil(base2k)={} > a.size()={}",
116                size,
117                a.size()
118            );
119            assert!(col < a.cols());
120        }
121
122        let mut a: VecZnx<&mut [u8]> = self.to_mut();
123        let a_size = a.size();
124
125        for j in 0..a_size {
126            a.at_mut(col, j)[idx] = 0
127        }
128
129        a.at_mut(col, size - 1)[idx] = data;
130
131        let mut carry: Vec<i64> = vec![0i64; 1];
132        let k_rem: usize = (base2k - (k % base2k)) % base2k;
133
134        for j in (0..size).rev() {
135            let slice = &mut a.at_mut(col, j)[idx..idx + 1];
136
137            if j == size - 1 {
138                ZnxRef::znx_normalize_first_step_inplace(base2k, k_rem, slice, &mut carry);
139            } else if j == 0 {
140                ZnxRef::znx_normalize_final_step_inplace(base2k, k_rem, slice, &mut carry);
141            } else {
142                ZnxRef::znx_normalize_middle_step_inplace(base2k, k_rem, slice, &mut carry);
143            }
144        }
145    }
146}
147
148impl<D: DataRef> VecZnx<D> {
149    pub fn decode_vec_i64(&self, base2k: usize, col: usize, k: usize, data: &mut [i64]) {
150        let size: usize = k.div_ceil(base2k);
151        #[cfg(debug_assertions)]
152        {
153            let a: VecZnx<&[u8]> = self.to_ref();
154            assert!(
155                data.len() >= a.n(),
156                "invalid data: data.len()={} < a.n()={}",
157                data.len(),
158                a.n()
159            );
160            assert!(col < a.cols());
161        }
162
163        let a: VecZnx<&[u8]> = self.to_ref();
164        data.copy_from_slice(a.at(col, 0));
165        let rem: usize = base2k - (k % base2k);
166        if k < base2k {
167            let scale = 1 << rem as i64;
168            data.iter_mut().for_each(|x| *x = div_round(*x, scale));
169        } else {
170            (1..size).for_each(|i| {
171                if i == size - 1 && rem != base2k {
172                    let k_rem: usize = (base2k - rem) % base2k;
173                    let scale: i64 = 1 << rem as i64;
174                    izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
175                        *y = (*y << k_rem) + div_round(*x, scale);
176                    });
177                } else {
178                    izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
179                        *y = (*y << base2k) + x;
180                    });
181                }
182            })
183        }
184    }
185
186    pub fn decode_coeff_i64(&self, base2k: usize, col: usize, k: usize, idx: usize) -> i64 {
187        #[cfg(debug_assertions)]
188        {
189            let a: VecZnx<&[u8]> = self.to_ref();
190            assert!(idx < a.n());
191            assert!(col < a.cols())
192        }
193
194        let a: VecZnx<&[u8]> = self.to_ref();
195        let size: usize = k.div_ceil(base2k);
196        let mut res: i64 = 0;
197        let rem: usize = base2k - (k % base2k);
198        (0..size).for_each(|j| {
199            let x: i64 = a.at(col, j)[idx];
200            if j == size - 1 && rem != base2k {
201                let k_rem: usize = (base2k - rem) % base2k;
202                let scale: i64 = 1 << rem as i64;
203                res = (res << k_rem) + div_round(x, scale);
204            } else {
205                res = (res << base2k) + x;
206            }
207        });
208        res
209    }
210
211    pub fn decode_vec_float(&self, base2k: usize, col: usize, data: &mut [Float]) {
212        #[cfg(debug_assertions)]
213        {
214            let a: VecZnx<&[u8]> = self.to_ref();
215            assert!(
216                data.len() >= a.n(),
217                "invalid data: data.len()={} < a.n()={}",
218                data.len(),
219                a.n()
220            );
221            assert!(col < a.cols());
222        }
223
224        let a: VecZnx<&[u8]> = self.to_ref();
225        let size: usize = a.size();
226        let prec: u32 = (base2k * size) as u32;
227
228        // 2^{base2k}
229        let base: Float = Float::with_val(prec, (1u64 << base2k) as f64);
230
231        // y[i] = sum x[j][i] * 2^{-base2k*j}
232        (0..size).for_each(|i| {
233            if i == 0 {
234                izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
235                    y.assign(*x);
236                    *y /= &base;
237                });
238            } else {
239                izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
240                    *y += Float::with_val(prec, *x);
241                    *y /= &base;
242                });
243            }
244        });
245    }
246}
247
248#[inline]
249pub fn div_round(a: i64, b: i64) -> i64 {
250    assert!(b != 0, "division by zero");
251    let div: i64 = a / b;
252    let rem: i64 = a % b;
253    if (2 * rem.abs()) >= b.abs() {
254        div + a.signum() * b.signum()
255    } else {
256        div
257    }
258}