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, VecZnxApi};
  3use itertools::izip;
  4use std::cmp::min;
  5
  6pub trait Encoding {
  7    /// encode a vector of i64 on the receiver.
  8    ///
  9    /// # Arguments
 10    ///
 11    /// * `log_base2k`: base two logarithm decomposition of the receiver.
 12    /// * `log_k`: base two logarithm of the scaling of the data.
 13    /// * `data`: data to encode on the receiver.
 14    /// * `log_max`: base two logarithm of the infinity norm of the input data.
 15    fn encode_vec_i64(&mut self, log_base2k: usize, log_k: usize, data: &[i64], log_max: usize);
 16
 17    /// decode a vector of i64 from the receiver.
 18    ///
 19    /// # Arguments
 20    ///
 21    /// * `log_base2k`: base two logarithm decomposition of the receiver.
 22    /// * `log_k`: base two logarithm of the scaling of the data.
 23    /// * `data`: data to decode from the receiver.
 24    fn decode_vec_i64(&self, log_base2k: usize, log_k: usize, data: &mut [i64]);
 25
 26    /// encodes a single i64 on the receiver at the given index.
 27    ///
 28    /// # Arguments
 29    ///
 30    /// * `log_base2k`: base two logarithm decomposition of the receiver.
 31    /// * `log_k`: base two logarithm of the scaling of the data.
 32    /// * `i`: index of the coefficient on which to encode the data.
 33    /// * `data`: data to encode on the receiver.
 34    /// * `log_max`: base two logarithm of the infinity norm of the input data.
 35    fn encode_coeff_i64(
 36        &mut self,
 37        log_base2k: usize,
 38        log_k: usize,
 39        i: usize,
 40        data: i64,
 41        log_max: usize,
 42    );
 43
 44    /// decode a single of i64 from the receiver at the given index.
 45    ///
 46    /// # Arguments
 47    ///
 48    /// * `log_base2k`: base two logarithm decomposition of the receiver.
 49    /// * `log_k`: base two logarithm of the scaling of the data.
 50    /// * `i`: index of the coefficient to decode.
 51    /// * `data`: data to decode from the receiver.
 52    fn decode_coeff_i64(&self, log_base2k: usize, log_k: usize, i: usize) -> i64;
 53}
 54
 55impl Encoding for VecZnx {
 56    fn encode_vec_i64(&mut self, log_base2k: usize, log_k: usize, data: &[i64], log_max: usize) {
 57        
 58        let cols: usize = (log_k + log_base2k - 1) / log_base2k;
 59
 60        assert!(cols <= self.cols(), "invalid argument log_k: (log_k + self.log_base2k - 1)/self.log_base2k={} > self.cols()={}", cols, self.cols());
 61
 62        let size: usize = min(data.len(), self.n());
 63        let log_k_rem: usize = log_base2k - (log_k % log_base2k);
 64
 65        // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
 66        // values on the last limb.
 67        // Else we decompose values base2k.
 68        if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
 69            (0..self.cols()).for_each(|i| unsafe {
 70                znx_zero_i64_ref(size as u64, self.at_mut(i).as_mut_ptr());
 71            });
 72            self.at_mut(cols - 1)[..size].copy_from_slice(&data[..size]);
 73        } else {
 74            let mask: i64 = (1 << log_base2k) - 1;
 75            let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
 76
 77            (0..steps).for_each(|i| unsafe {
 78                znx_zero_i64_ref(size as u64, self.at_mut(i).as_mut_ptr());
 79            });
 80
 81            (cols - steps..cols)
 82                .rev()
 83                .enumerate()
 84                .for_each(|(i, i_rev)| {
 85                    let shift: usize = i * log_base2k;
 86                    izip!(self.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
 87                        .for_each(|(y, x)| *y = (x >> shift) & mask);
 88                })
 89        }
 90
 91        // Case where self.prec % self.k != 0.
 92        if log_k_rem != log_base2k {
 93            let cols = self.cols();
 94            let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
 95            (cols - steps..cols).rev().for_each(|i| {
 96                self.at_mut(i)[..size]
 97                    .iter_mut()
 98                    .for_each(|x| *x <<= log_k_rem);
 99            })
100        }
101    }
102
103    fn decode_vec_i64(&self, log_base2k: usize, log_k: usize, data: &mut [i64]) {
104        let cols: usize = (log_k + log_base2k - 1) / log_base2k;
105        assert!(
106            data.len() >= self.n,
107            "invalid data: data.len()={} < self.n()={}",
108            data.len(),
109            self.n
110        );
111        data.copy_from_slice(self.at(0));
112        let rem: usize = log_base2k - (log_k % log_base2k);
113        (1..cols).for_each(|i| {
114            if i == cols - 1 && rem != log_base2k {
115                let k_rem: usize = log_base2k - rem;
116                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
117                    *y = (*y << k_rem) + (x >> rem);
118                });
119            } else {
120                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
121                    *y = (*y << log_base2k) + x;
122                });
123            }
124        })
125    }
126
127    fn encode_coeff_i64(
128        &mut self,
129        log_base2k: usize,
130        log_k: usize,
131        i: usize,
132        value: i64,
133        log_max: usize,
134    ) {
135        assert!(i < self.n());
136        let cols: usize = (log_k + log_base2k - 1) / log_base2k;
137        assert!(cols <= self.cols(), "invalid argument log_k: (log_k + self.log_base2k - 1)/self.log_base2k={} > self.cols()={}", cols, self.cols());
138        let log_k_rem: usize = log_base2k - (log_k % log_base2k);
139        let cols = self.cols();
140
141        // If 2^{log_base2k} * 2^{log_k_rem} < 2^{63}-1, then we can simply copy
142        // values on the last limb.
143        // Else we decompose values base2k.
144        if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
145            (0..cols - 1).for_each(|j| self.at_mut(j)[i] = 0);
146
147            self.at_mut(self.cols() - 1)[i] = value;
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..cols - steps).for_each(|j| self.at_mut(j)[i] = 0);
153
154            (cols - steps..cols)
155                .rev()
156                .enumerate()
157                .for_each(|(j, j_rev)| {
158                    self.at_mut(j_rev)[i] = (value >> (j * log_base2k)) & mask;
159                })
160        }
161
162        // Case where self.prec % self.k != 0.
163        if log_k_rem != log_base2k {
164            let cols = self.cols();
165            let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
166            (cols - steps..cols).rev().for_each(|j| {
167                self.at_mut(j)[i] <<= log_k_rem;
168            })
169        }
170    }
171
172    fn decode_coeff_i64(&self, log_base2k: usize, log_k: usize, i: usize) -> i64 {
173        let cols: usize = (log_k + log_base2k - 1) / log_base2k;
174        assert!(i < self.n());
175        let mut res: i64 = self.data[i];
176        let rem: usize = log_base2k - (log_k % log_base2k);
177        (1..cols).for_each(|i| {
178            let x = self.data[i * self.n];
179            if i == cols - 1 && rem != log_base2k {
180                let k_rem: usize = log_base2k - rem;
181                res = (res << k_rem) + (x >> rem);
182            } else {
183                res = (res << log_base2k) + x;
184            }
185        });
186        res
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use crate::{Encoding, VecZnx};
193    use itertools::izip;
194    use sampling::source::Source;
195
196    #[test]
197    fn test_set_get_i64_lo_norm() {
198        let n: usize = 8;
199        let log_base2k: usize = 17;
200        let cols: usize = 5;
201        let log_k: usize = cols * log_base2k - 5;
202        let mut a: VecZnx = VecZnx::new(n, cols);
203        let mut have: Vec<i64> = vec![i64::default(); n];
204        have.iter_mut()
205            .enumerate()
206            .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
207        a.encode_vec_i64(log_base2k, log_k, &have, 10);
208        let mut want = vec![i64::default(); n];
209        a.decode_vec_i64(log_base2k, log_k, &mut want);
210        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
211    }
212
213    #[test]
214    fn test_set_get_i64_hi_norm() {
215        let n: usize = 8;
216        let log_base2k: usize = 17;
217        let cols: usize = 5;
218        let log_k: usize = cols * log_base2k - 5;
219        let mut a: VecZnx = VecZnx::new(n, cols);
220        let mut have: Vec<i64> = vec![i64::default(); n];
221        let mut source = Source::new([1; 32]);
222        have.iter_mut().for_each(|x| {
223            *x = source
224                .next_u64n(u64::MAX, u64::MAX)
225                .wrapping_sub(u64::MAX / 2 + 1) as i64;
226        });
227        a.encode_vec_i64(log_base2k, log_k, &have, 63);
228        //(0..a.cols()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
229        let mut want = vec![i64::default(); n];
230        //(0..a.cols()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
231        a.decode_vec_i64(log_base2k, log_k, &mut want);
232        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
233    }
234}