arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/base2k/src/encoding.rs
base2k/src/encoding.rsRaw
  1use crate::ffi::znx::znx_zero_i64_ref;
  2use crate::{Infos, VecZnx};
  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(&mut self, log_base2k: usize, log_k: usize, i: usize, data: i64, log_max: usize);
 44
 45    /// decode a single of i64 from the receiver at the given index.
 46    ///
 47    /// # Arguments
 48    ///
 49    /// * `log_base2k`: base two logarithm decomposition of the receiver.
 50    /// * `log_k`: base two logarithm of the scaling of the data.
 51    /// * `i`: index of the coefficient to decode.
 52    /// * `data`: data to decode from the receiver.
 53    fn decode_coeff_i64(&self, log_base2k: usize, log_k: usize, i: usize) -> i64;
 54}
 55
 56impl Encoding for VecZnx {
 57    fn encode_vec_i64(&mut self, log_base2k: usize, log_k: usize, data: &[i64], log_max: usize) {
 58        encode_vec_i64(self, log_base2k, log_k, data, log_max)
 59    }
 60
 61    fn decode_vec_i64(&self, log_base2k: usize, log_k: usize, data: &mut [i64]) {
 62        decode_vec_i64(self, log_base2k, log_k, data)
 63    }
 64
 65    fn decode_vec_float(&self, log_base2k: usize, data: &mut [Float]) {
 66        decode_vec_float(self, log_base2k, data)
 67    }
 68
 69    fn encode_coeff_i64(&mut self, log_base2k: usize, log_k: usize, i: usize, value: i64, log_max: usize) {
 70        encode_coeff_i64(self, log_base2k, log_k, i, value, log_max)
 71    }
 72
 73    fn decode_coeff_i64(&self, log_base2k: usize, log_k: usize, i: usize) -> i64 {
 74        decode_coeff_i64(self, log_base2k, log_k, i)
 75    }
 76}
 77
 78fn encode_vec_i64(a: &mut VecZnx, log_base2k: usize, log_k: usize, data: &[i64], log_max: usize) {
 79    let cols: usize = (log_k + log_base2k - 1) / log_base2k;
 80
 81    debug_assert!(
 82        cols <= a.cols(),
 83        "invalid argument log_k: (log_k + a.log_base2k - 1)/a.log_base2k={} > a.cols()={}",
 84        cols,
 85        a.cols()
 86    );
 87
 88    let size: usize = min(data.len(), a.n());
 89    let log_k_rem: usize = log_base2k - (log_k % log_base2k);
 90
 91    (0..a.cols()).for_each(|i| unsafe {
 92        znx_zero_i64_ref(size as u64, a.at_mut(i).as_mut_ptr());
 93    });
 94
 95    // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
 96    // values on the last limb.
 97    // Else we decompose values base2k.
 98    if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
 99        a.at_mut(cols - 1)[..size].copy_from_slice(&data[..size]);
100    } else {
101        let mask: i64 = (1 << log_base2k) - 1;
102        let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
103        (cols - steps..cols)
104            .rev()
105            .enumerate()
106            .for_each(|(i, i_rev)| {
107                let shift: usize = i * log_base2k;
108                izip!(a.at_mut(i_rev)[..size].iter_mut(), data[..size].iter()).for_each(|(y, x)| *y = (x >> shift) & mask);
109            })
110    }
111
112    // Case where self.prec % self.k != 0.
113    if log_k_rem != log_base2k {
114        let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
115        (cols - steps..cols).rev().for_each(|i| {
116            a.at_mut(i)[..size]
117                .iter_mut()
118                .for_each(|x| *x <<= log_k_rem);
119        })
120    }
121}
122
123fn decode_vec_i64(a: &VecZnx, log_base2k: usize, log_k: usize, data: &mut [i64]) {
124    let cols: usize = (log_k + log_base2k - 1) / log_base2k;
125    debug_assert!(
126        data.len() >= a.n(),
127        "invalid data: data.len()={} < a.n()={}",
128        data.len(),
129        a.n()
130    );
131    data.copy_from_slice(a.at(0));
132    let rem: usize = log_base2k - (log_k % log_base2k);
133    (1..cols).for_each(|i| {
134        if i == cols - 1 && rem != log_base2k {
135            let k_rem: usize = log_base2k - rem;
136            izip!(a.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
137                *y = (*y << k_rem) + (x >> rem);
138            });
139        } else {
140            izip!(a.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
141                *y = (*y << log_base2k) + x;
142            });
143        }
144    })
145}
146
147fn decode_vec_float(a: &VecZnx, log_base2k: usize, data: &mut [Float]) {
148    let cols: usize = a.cols();
149    debug_assert!(
150        data.len() >= a.n(),
151        "invalid data: data.len()={} < a.n()={}",
152        data.len(),
153        a.n()
154    );
155
156    let prec: u32 = (log_base2k * cols) as u32;
157
158    // 2^{log_base2k}
159    let base = Float::with_val(prec, (1 << log_base2k) as f64);
160
161    // y[i] = sum x[j][i] * 2^{-log_base2k*j}
162    (0..cols).for_each(|i| {
163        if i == 0 {
164            izip!(a.at(cols - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
165                y.assign(*x);
166                *y /= &base;
167            });
168        } else {
169            izip!(a.at(cols - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
170                *y += Float::with_val(prec, *x);
171                *y /= &base;
172            });
173        }
174    });
175}
176
177fn encode_coeff_i64(a: &mut VecZnx, log_base2k: usize, log_k: usize, i: usize, value: i64, log_max: usize) {
178    debug_assert!(i < a.n());
179    let cols: usize = (log_k + log_base2k - 1) / log_base2k;
180    debug_assert!(
181        cols <= a.cols(),
182        "invalid argument log_k: (log_k + a.log_base2k - 1)/a.log_base2k={} > a.cols()={}",
183        cols,
184        a.cols()
185    );
186    let log_k_rem: usize = log_base2k - (log_k % log_base2k);
187    (0..a.cols()).for_each(|j| a.at_mut(j)[i] = 0);
188
189    // If 2^{log_base2k} * 2^{log_k_rem} < 2^{63}-1, then we can simply copy
190    // values on the last limb.
191    // Else we decompose values base2k.
192    if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
193        a.at_mut(cols - 1)[i] = value;
194    } else {
195        let mask: i64 = (1 << log_base2k) - 1;
196        let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
197        (cols - steps..cols)
198            .rev()
199            .enumerate()
200            .for_each(|(j, j_rev)| {
201                a.at_mut(j_rev)[i] = (value >> (j * log_base2k)) & mask;
202            })
203    }
204
205    // Case where prec % k != 0.
206    if log_k_rem != log_base2k {
207        let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
208        (cols - steps..cols).rev().for_each(|j| {
209            a.at_mut(j)[i] <<= log_k_rem;
210        })
211    }
212}
213
214fn decode_coeff_i64(a: &VecZnx, log_base2k: usize, log_k: usize, i: usize) -> i64 {
215    let cols: usize = (log_k + log_base2k - 1) / log_base2k;
216    debug_assert!(i < a.n());
217    let data: &[i64] = a.raw();
218    let mut res: i64 = data[i];
219    let rem: usize = log_base2k - (log_k % log_base2k);
220    (1..cols).for_each(|i| {
221        let x = data[i * a.n()];
222        if i == cols - 1 && rem != log_base2k {
223            let k_rem: usize = log_base2k - rem;
224            res = (res << k_rem) + (x >> rem);
225        } else {
226            res = (res << log_base2k) + x;
227        }
228    });
229    res
230}
231
232#[cfg(test)]
233mod tests {
234    use crate::{Encoding, VecZnx};
235    use itertools::izip;
236    use sampling::source::Source;
237
238    #[test]
239    fn test_set_get_i64_lo_norm() {
240        let n: usize = 8;
241        let log_base2k: usize = 17;
242        let cols: usize = 5;
243        let log_k: usize = cols * log_base2k - 5;
244        let mut a: VecZnx = VecZnx::new(n, cols);
245        let mut have: Vec<i64> = vec![i64::default(); n];
246        have.iter_mut()
247            .enumerate()
248            .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
249        a.encode_vec_i64(log_base2k, log_k, &have, 10);
250        let mut want = vec![i64::default(); n];
251        a.decode_vec_i64(log_base2k, log_k, &mut want);
252        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
253    }
254
255    #[test]
256    fn test_set_get_i64_hi_norm() {
257        let n: usize = 8;
258        let log_base2k: usize = 17;
259        let cols: usize = 5;
260        let log_k: usize = cols * log_base2k - 5;
261        let mut a: VecZnx = VecZnx::new(n, cols);
262        let mut have: Vec<i64> = vec![i64::default(); n];
263        let mut source = Source::new([1; 32]);
264        have.iter_mut().for_each(|x| {
265            *x = source
266                .next_u64n(u64::MAX, u64::MAX)
267                .wrapping_sub(u64::MAX / 2 + 1) as i64;
268        });
269        a.encode_vec_i64(log_base2k, log_k, &have, 63);
270        //(0..a.cols()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
271        let mut want = vec![i64::default(); n];
272        //(0..a.cols()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
273        a.decode_vec_i64(log_base2k, log_k, &mut want);
274        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
275    }
276}