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