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    // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
141    // values on the last limb.
142    // Else we decompose values base2k.
143    if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
144        (0..a.cols()).for_each(|i| unsafe {
145            znx_zero_i64_ref(size as u64, a.at_mut(i).as_mut_ptr());
146        });
147        a.at_mut(cols - 1)[..size].copy_from_slice(&data[..size]);
148    } else {
149        let mask: i64 = (1 << log_base2k) - 1;
150        let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
151
152        (0..steps).for_each(|i| unsafe {
153            znx_zero_i64_ref(size as u64, a.at_mut(i).as_mut_ptr());
154        });
155
156        (cols - steps..cols)
157            .rev()
158            .enumerate()
159            .for_each(|(i, i_rev)| {
160                let shift: usize = i * log_base2k;
161                izip!(a.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
162                    .for_each(|(y, x)| *y = (x >> shift) & mask);
163            })
164    }
165
166    // Case where self.prec % self.k != 0.
167    if log_k_rem != log_base2k {
168        let cols = a.cols();
169        let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
170        (cols - steps..cols).rev().for_each(|i| {
171            a.at_mut(i)[..size]
172                .iter_mut()
173                .for_each(|x| *x <<= log_k_rem);
174        })
175    }
176}
177
178fn decode_vec_i64<T: VecZnxCommon>(a: &T, log_base2k: usize, log_k: usize, data: &mut [i64]) {
179    let cols: usize = (log_k + log_base2k - 1) / log_base2k;
180    debug_assert!(
181        data.len() >= a.n(),
182        "invalid data: data.len()={} < a.n()={}",
183        data.len(),
184        a.n()
185    );
186    data.copy_from_slice(a.at(0));
187    let rem: usize = log_base2k - (log_k % log_base2k);
188    (1..cols).for_each(|i| {
189        if i == cols - 1 && rem != log_base2k {
190            let k_rem: usize = log_base2k - rem;
191            izip!(a.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
192                *y = (*y << k_rem) + (x >> rem);
193            });
194        } else {
195            izip!(a.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
196                *y = (*y << log_base2k) + x;
197            });
198        }
199    })
200}
201
202fn decode_vec_float<T: VecZnxCommon>(a: &T, log_base2k: usize, data: &mut [Float]) {
203    let cols: usize = a.cols();
204    debug_assert!(
205        data.len() >= a.n(),
206        "invalid data: data.len()={} < a.n()={}",
207        data.len(),
208        a.n()
209    );
210
211    let prec: u32 = (log_base2k * cols) as u32;
212
213    // 2^{log_base2k}
214    let base = Float::with_val(prec, (1 << log_base2k) as f64);
215
216    // y[i] = sum x[j][i] * 2^{-log_base2k*j}
217    (0..cols).for_each(|i| {
218        if i == 0 {
219            izip!(a.at(cols - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
220                y.assign(*x);
221                *y /= &base;
222            });
223        } else {
224            izip!(a.at(cols - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
225                *y += Float::with_val(prec, *x);
226                *y /= &base;
227            });
228        }
229    });
230}
231
232fn encode_coeff_i64<T: VecZnxCommon>(
233    a: &mut T,
234    log_base2k: usize,
235    log_k: usize,
236    i: usize,
237    value: i64,
238    log_max: usize,
239) {
240    debug_assert!(i < a.n());
241    let cols: usize = (log_k + log_base2k - 1) / log_base2k;
242    debug_assert!(
243        cols <= a.cols(),
244        "invalid argument log_k: (log_k + a.log_base2k - 1)/a.log_base2k={} > a.cols()={}",
245        cols,
246        a.cols()
247    );
248    let log_k_rem: usize = log_base2k - (log_k % log_base2k);
249    let cols = a.cols();
250
251    // If 2^{log_base2k} * 2^{log_k_rem} < 2^{63}-1, then we can simply copy
252    // values on the last limb.
253    // Else we decompose values base2k.
254    if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
255        (0..cols - 1).for_each(|j| a.at_mut(j)[i] = 0);
256
257        a.at_mut(a.cols() - 1)[i] = value;
258    } else {
259        let mask: i64 = (1 << log_base2k) - 1;
260        let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
261
262        (0..cols - steps).for_each(|j| a.at_mut(j)[i] = 0);
263
264        (cols - steps..cols)
265            .rev()
266            .enumerate()
267            .for_each(|(j, j_rev)| {
268                a.at_mut(j_rev)[i] = (value >> (j * log_base2k)) & mask;
269            })
270    }
271
272    // Case where prec % k != 0.
273    if log_k_rem != log_base2k {
274        let cols = a.cols();
275        let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
276        (cols - steps..cols).rev().for_each(|j| {
277            a.at_mut(j)[i] <<= log_k_rem;
278        })
279    }
280}
281
282fn decode_coeff_i64<T: VecZnxCommon>(a: &T, log_base2k: usize, log_k: usize, i: usize) -> i64 {
283    let cols: usize = (log_k + log_base2k - 1) / log_base2k;
284    debug_assert!(i < a.n());
285    let data: &[i64] = a.raw();
286    let mut res: i64 = data[i];
287    let rem: usize = log_base2k - (log_k % log_base2k);
288    (1..cols).for_each(|i| {
289        let x = data[i * a.n()];
290        if i == cols - 1 && rem != log_base2k {
291            let k_rem: usize = log_base2k - rem;
292            res = (res << k_rem) + (x >> rem);
293        } else {
294            res = (res << log_base2k) + x;
295        }
296    });
297    res
298}
299
300#[cfg(test)]
301mod tests {
302    use crate::{Encoding, VecZnx};
303    use itertools::izip;
304    use sampling::source::Source;
305
306    #[test]
307    fn test_set_get_i64_lo_norm() {
308        let n: usize = 8;
309        let log_base2k: usize = 17;
310        let cols: usize = 5;
311        let log_k: usize = cols * log_base2k - 5;
312        let mut a: VecZnx = VecZnx::new(n, cols);
313        let mut have: Vec<i64> = vec![i64::default(); n];
314        have.iter_mut()
315            .enumerate()
316            .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
317        a.encode_vec_i64(log_base2k, log_k, &have, 10);
318        let mut want = vec![i64::default(); n];
319        a.decode_vec_i64(log_base2k, log_k, &mut want);
320        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
321    }
322
323    #[test]
324    fn test_set_get_i64_hi_norm() {
325        let n: usize = 8;
326        let log_base2k: usize = 17;
327        let cols: usize = 5;
328        let log_k: usize = cols * log_base2k - 5;
329        let mut a: VecZnx = VecZnx::new(n, cols);
330        let mut have: Vec<i64> = vec![i64::default(); n];
331        let mut source = Source::new([1; 32]);
332        have.iter_mut().for_each(|x| {
333            *x = source
334                .next_u64n(u64::MAX, u64::MAX)
335                .wrapping_sub(u64::MAX / 2 + 1) as i64;
336        });
337        a.encode_vec_i64(log_base2k, log_k, &have, 63);
338        //(0..a.cols()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
339        let mut want = vec![i64::default(); n];
340        //(0..a.cols()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
341        a.decode_vec_i64(log_base2k, log_k, &mut want);
342        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
343    }
344}