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