1use crate::ffi::znx::znx_zero_i64_ref;
  2use crate::{Infos, VecZnx, VecZnxLayout};
  3use itertools::izip;
  4use rug::{Assign, Float};
  5use std::cmp::min;
  6
  7pub trait Encoding {
  8    /// encode a vector of i64 on the receiver.
  9    ///
 10    /// # Arguments
 11    ///
 12    /// * `col_i`: the index of the poly where to encode the data.
 13    /// * `log_base2k`: base two negative logarithm decomposition of the receiver.
 14    /// * `log_k`: base two negative logarithm of the scaling of the data.
 15    /// * `data`: data to encode on the receiver.
 16    /// * `log_max`: base two logarithm of the infinity norm of the input data.
 17    fn encode_vec_i64(&mut self, col_i: usize, log_base2k: usize, log_k: usize, data: &[i64], log_max: usize);
 18
 19    /// decode a vector of i64 from the receiver.
 20    ///
 21    /// # Arguments
 22    ///
 23    /// * `col_i`: the index of the poly where to encode the data.
 24    /// * `log_base2k`: base two negative logarithm decomposition of the receiver.
 25    /// * `log_k`: base two logarithm of the scaling of the data.
 26    /// * `data`: data to decode from the receiver.
 27    fn decode_vec_i64(&self, col_i: usize, log_base2k: usize, log_k: usize, data: &mut [i64]);
 28
 29    /// decode a vector of Float from the receiver.
 30    ///
 31    /// # Arguments
 32    /// * `col_i`: the index of the poly where to encode the data.
 33    /// * `log_base2k`: base two negative logarithm decomposition of the receiver.
 34    /// * `data`: data to decode from the receiver.
 35    fn decode_vec_float(&self, col_i: usize, log_base2k: usize, data: &mut [Float]);
 36
 37    /// encodes a single i64 on the receiver at the given index.
 38    ///
 39    /// # Arguments
 40    ///
 41    /// * `col_i`: the index of the poly where to encode the data.
 42    /// * `log_base2k`: base two negative logarithm decomposition of the receiver.
 43    /// * `log_k`: base two negative logarithm of the scaling of the data.
 44    /// * `i`: index of the coefficient on which to encode the data.
 45    /// * `data`: data to encode on the receiver.
 46    /// * `log_max`: base two logarithm of the infinity norm of the input data.
 47    fn encode_coeff_i64(&mut self, col_i: usize, log_base2k: usize, log_k: usize, i: usize, data: i64, log_max: usize);
 48
 49    /// decode a single of i64 from the receiver at the given index.
 50    ///
 51    /// # Arguments
 52    ///
 53    /// * `col_i`: the index of the poly where to encode the data.
 54    /// * `log_base2k`: base two negative logarithm decomposition of the receiver.
 55    /// * `log_k`: base two negative logarithm of the scaling of the data.
 56    /// * `i`: index of the coefficient to decode.
 57    /// * `data`: data to decode from the receiver.
 58    fn decode_coeff_i64(&self, col_i: usize, log_base2k: usize, log_k: usize, i: usize) -> i64;
 59}
 60
 61impl Encoding for VecZnx {
 62    fn encode_vec_i64(&mut self, col_i: usize, log_base2k: usize, log_k: usize, data: &[i64], log_max: usize) {
 63        encode_vec_i64(self, col_i, log_base2k, log_k, data, log_max)
 64    }
 65
 66    fn decode_vec_i64(&self, col_i: usize, log_base2k: usize, log_k: usize, data: &mut [i64]) {
 67        decode_vec_i64(self, col_i, log_base2k, log_k, data)
 68    }
 69
 70    fn decode_vec_float(&self, col_i: usize, log_base2k: usize, data: &mut [Float]) {
 71        decode_vec_float(self, col_i, log_base2k, data)
 72    }
 73
 74    fn encode_coeff_i64(&mut self, col_i: usize, log_base2k: usize, log_k: usize, i: usize, value: i64, log_max: usize) {
 75        encode_coeff_i64(self, col_i, log_base2k, log_k, i, value, log_max)
 76    }
 77
 78    fn decode_coeff_i64(&self, col_i: usize, log_base2k: usize, log_k: usize, i: usize) -> i64 {
 79        decode_coeff_i64(self, col_i, log_base2k, log_k, i)
 80    }
 81}
 82
 83fn encode_vec_i64(a: &mut VecZnx, col_i: usize, log_base2k: usize, log_k: usize, data: &[i64], log_max: usize) {
 84    let limbs: usize = (log_k + log_base2k - 1) / log_base2k;
 85
 86    #[cfg(debug_assertions)]
 87    {
 88        assert!(
 89            limbs <= a.limbs(),
 90            "invalid argument log_k: (log_k + a.log_base2k - 1)/a.log_base2k={} > a.limbs()={}",
 91            limbs,
 92            a.limbs()
 93        );
 94        assert!(col_i < a.cols());
 95        assert!(data.len() <= a.n())
 96    }
 97
 98    let data_len: usize = data.len();
 99    let log_k_rem: usize = log_base2k - (log_k % log_base2k);
100
101    // Zeroes coefficients of the i-th column
102    (0..a.limbs()).for_each(|i| unsafe {
103        znx_zero_i64_ref(a.n() as u64, a.at_mut_ptr(col_i, i));
104    });
105
106    // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
107    // values on the last limb.
108    // Else we decompose values base2k.
109    if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
110        a.at_poly_mut(col_i, limbs - 1)[..data_len].copy_from_slice(&data[..data_len]);
111    } else {
112        let mask: i64 = (1 << log_base2k) - 1;
113        let steps: usize = min(limbs, (log_max + log_base2k - 1) / log_base2k);
114        (limbs - steps..limbs)
115            .rev()
116            .enumerate()
117            .for_each(|(i, i_rev)| {
118                let shift: usize = i * log_base2k;
119                izip!(a.at_poly_mut(col_i, i_rev).iter_mut(), data.iter()).for_each(|(y, x)| *y = (x >> shift) & mask);
120            })
121    }
122
123    // Case where self.prec % self.k != 0.
124    if log_k_rem != log_base2k {
125        let steps: usize = min(limbs, (log_max + log_base2k - 1) / log_base2k);
126        (limbs - steps..limbs).rev().for_each(|i| {
127            a.at_poly_mut(col_i, i)[..data_len]
128                .iter_mut()
129                .for_each(|x| *x <<= log_k_rem);
130        })
131    }
132}
133
134fn decode_vec_i64(a: &VecZnx, col_i: usize, log_base2k: usize, log_k: usize, data: &mut [i64]) {
135    let limbs: usize = (log_k + log_base2k - 1) / log_base2k;
136    #[cfg(debug_assertions)]
137    {
138        assert!(
139            data.len() >= a.n(),
140            "invalid data: data.len()={} < a.n()={}",
141            data.len(),
142            a.n()
143        );
144        assert!(col_i < a.cols());
145    }
146    data.copy_from_slice(a.at_poly(col_i, 0));
147    let rem: usize = log_base2k - (log_k % log_base2k);
148    (1..limbs).for_each(|i| {
149        if i == limbs - 1 && rem != log_base2k {
150            let k_rem: usize = log_base2k - rem;
151            izip!(a.at_poly(col_i, i).iter(), data.iter_mut()).for_each(|(x, y)| {
152                *y = (*y << k_rem) + (x >> rem);
153            });
154        } else {
155            izip!(a.at_poly(col_i, i).iter(), data.iter_mut()).for_each(|(x, y)| {
156                *y = (*y << log_base2k) + x;
157            });
158        }
159    })
160}
161
162fn decode_vec_float(a: &VecZnx, col_i: usize, log_base2k: usize, data: &mut [Float]) {
163    let limbs: usize = a.limbs();
164    #[cfg(debug_assertions)]
165    {
166        assert!(
167            data.len() >= a.n(),
168            "invalid data: data.len()={} < a.n()={}",
169            data.len(),
170            a.n()
171        );
172        assert!(col_i < a.cols());
173    }
174
175    let prec: u32 = (log_base2k * limbs) as u32;
176
177    // 2^{log_base2k}
178    let base = Float::with_val(prec, (1 << log_base2k) as f64);
179
180    // y[i] = sum x[j][i] * 2^{-log_base2k*j}
181    (0..limbs).for_each(|i| {
182        if i == 0 {
183            izip!(a.at_poly(col_i, limbs - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
184                y.assign(*x);
185                *y /= &base;
186            });
187        } else {
188            izip!(a.at_poly(col_i, limbs - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
189                *y += Float::with_val(prec, *x);
190                *y /= &base;
191            });
192        }
193    });
194}
195
196fn encode_coeff_i64(a: &mut VecZnx, col_i: usize, log_base2k: usize, log_k: usize, i: usize, value: i64, log_max: usize) {
197    let limbs: usize = (log_k + log_base2k - 1) / log_base2k;
198
199    #[cfg(debug_assertions)]
200    {
201        assert!(i < a.n());
202        assert!(
203            limbs <= a.limbs(),
204            "invalid argument log_k: (log_k + a.log_base2k - 1)/a.log_base2k={} > a.limbs()={}",
205            limbs,
206            a.limbs()
207        );
208        assert!(col_i < a.cols());
209    }
210
211    let log_k_rem: usize = log_base2k - (log_k % log_base2k);
212    (0..a.limbs()).for_each(|j| a.at_poly_mut(col_i, j)[i] = 0);
213
214    // If 2^{log_base2k} * 2^{log_k_rem} < 2^{63}-1, then we can simply copy
215    // values on the last limb.
216    // Else we decompose values base2k.
217    if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
218        a.at_poly_mut(col_i, limbs - 1)[i] = value;
219    } else {
220        let mask: i64 = (1 << log_base2k) - 1;
221        let steps: usize = min(limbs, (log_max + log_base2k - 1) / log_base2k);
222        (limbs - steps..limbs)
223            .rev()
224            .enumerate()
225            .for_each(|(j, j_rev)| {
226                a.at_poly_mut(col_i, j_rev)[i] = (value >> (j * log_base2k)) & mask;
227            })
228    }
229
230    // Case where prec % k != 0.
231    if log_k_rem != log_base2k {
232        let steps: usize = min(limbs, (log_max + log_base2k - 1) / log_base2k);
233        (limbs - steps..limbs).rev().for_each(|j| {
234            a.at_poly_mut(col_i, j)[i] <<= log_k_rem;
235        })
236    }
237}
238
239fn decode_coeff_i64(a: &VecZnx, col_i: usize, log_base2k: usize, log_k: usize, i: usize) -> i64 {
240    #[cfg(debug_assertions)]
241    {
242        assert!(i < a.n());
243        assert!(col_i < a.cols())
244    }
245
246    let cols: usize = (log_k + log_base2k - 1) / log_base2k;
247    let data: &[i64] = a.raw();
248    let mut res: i64 = data[i];
249    let rem: usize = log_base2k - (log_k % log_base2k);
250    let slice_size: usize = a.n() * a.limbs();
251    (1..cols).for_each(|i| {
252        let x = data[i * slice_size];
253        if i == cols - 1 && rem != log_base2k {
254            let k_rem: usize = log_base2k - rem;
255            res = (res << k_rem) + (x >> rem);
256        } else {
257            res = (res << log_base2k) + x;
258        }
259    });
260    res
261}
262
263#[cfg(test)]
264mod tests {
265    use crate::{Encoding, Infos, VecZnx, VecZnxLayout};
266    use itertools::izip;
267    use sampling::source::Source;
268
269    #[test]
270    fn test_set_get_i64_lo_norm() {
271        let n: usize = 8;
272        let log_base2k: usize = 17;
273        let cols: usize = 5;
274        let log_k: usize = cols * log_base2k - 5;
275        let mut a: VecZnx = VecZnx::new(n, 2, cols);
276        let mut source: Source = Source::new([0u8; 32]);
277        let raw: &mut [i64] = a.raw_mut();
278        raw.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
279        (0..a.cols()).for_each(|col_i| {
280            let mut have: Vec<i64> = vec![i64::default(); n];
281            have.iter_mut()
282                .for_each(|x| *x = (source.next_i64() << 56) >> 56);
283            a.encode_vec_i64(col_i, log_base2k, log_k, &have, 10);
284            let mut want: Vec<i64> = vec![i64::default(); n];
285            a.decode_vec_i64(col_i, log_base2k, log_k, &mut want);
286            izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
287        });
288    }
289
290    #[test]
291    fn test_set_get_i64_hi_norm() {
292        let n: usize = 8;
293        let log_base2k: usize = 17;
294        let cols: usize = 5;
295        let log_k: usize = cols * log_base2k - 5;
296        let mut a: VecZnx = VecZnx::new(n, 2, cols);
297        let mut source = Source::new([0u8; 32]);
298        let raw: &mut [i64] = a.raw_mut();
299        raw.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
300        (0..a.cols()).for_each(|col_i| {
301            let mut have: Vec<i64> = vec![i64::default(); n];
302            have.iter_mut().for_each(|x| *x = source.next_i64());
303            a.encode_vec_i64(col_i, log_base2k, log_k, &have, 64);
304            let mut want = vec![i64::default(); n];
305            a.decode_vec_i64(col_i, log_base2k, log_k, &mut want);
306            izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
307        })
308    }
309}