1use itertools::izip;
  2use rug::{Assign, Float};
  3
  4use crate::hal::{
  5    api::{ZnxInfos, ZnxView, ZnxViewMut, ZnxZero},
  6    layouts::{DataMut, DataRef, VecZnx, VecZnxToMut, VecZnxToRef},
  7};
  8
  9impl<D: DataMut> VecZnx<D> {
 10    pub fn encode_vec_i64(&mut self, basek: usize, col: usize, k: usize, data: &[i64], log_max: usize) {
 11        let size: usize = k.div_ceil(basek);
 12
 13        #[cfg(debug_assertions)]
 14        {
 15            let a: VecZnx<&mut [u8]> = self.to_mut();
 16            assert!(
 17                size <= a.size(),
 18                "invalid argument k: k.div_ceil(basek)={} > a.size()={}",
 19                size,
 20                a.size()
 21            );
 22            assert!(col < a.cols());
 23            assert!(data.len() <= a.n())
 24        }
 25
 26        let data_len: usize = data.len();
 27        let mut a: VecZnx<&mut [u8]> = self.to_mut();
 28        let k_rem: usize = basek - (k % basek);
 29
 30        // Zeroes coefficients of the i-th column
 31        (0..a.size()).for_each(|i| {
 32            a.zero_at(col, i);
 33        });
 34
 35        // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
 36        // values on the last limb.
 37        // Else we decompose values base2k.
 38        if log_max + k_rem < 63 || k_rem == basek {
 39            a.at_mut(col, size - 1)[..data_len].copy_from_slice(&data[..data_len]);
 40        } else {
 41            let mask: i64 = (1 << basek) - 1;
 42            let steps: usize = size.min(log_max.div_ceil(basek));
 43            (size - steps..size)
 44                .rev()
 45                .enumerate()
 46                .for_each(|(i, i_rev)| {
 47                    let shift: usize = i * basek;
 48                    izip!(a.at_mut(col, i_rev).iter_mut(), data.iter()).for_each(|(y, x)| *y = (x >> shift) & mask);
 49                })
 50        }
 51
 52        // Case where self.prec % self.k != 0.
 53        if k_rem != basek {
 54            let steps: usize = size.min(log_max.div_ceil(basek));
 55            (size - steps..size).rev().for_each(|i| {
 56                a.at_mut(col, i)[..data_len]
 57                    .iter_mut()
 58                    .for_each(|x| *x <<= k_rem);
 59            })
 60        }
 61    }
 62
 63    pub fn encode_coeff_i64(&mut self, basek: usize, col: usize, k: usize, idx: usize, data: i64, log_max: usize) {
 64        let size: usize = k.div_ceil(basek);
 65
 66        #[cfg(debug_assertions)]
 67        {
 68            let a: VecZnx<&mut [u8]> = self.to_mut();
 69            assert!(idx < a.n());
 70            assert!(
 71                size <= a.size(),
 72                "invalid argument k: k.div_ceil(basek)={} > a.size()={}",
 73                size,
 74                a.size()
 75            );
 76            assert!(col < a.cols());
 77        }
 78
 79        let k_rem: usize = basek - (k % basek);
 80        let mut a: VecZnx<&mut [u8]> = self.to_mut();
 81        (0..a.size()).for_each(|j| a.at_mut(col, j)[idx] = 0);
 82
 83        // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
 84        // values on the last limb.
 85        // Else we decompose values base2k.
 86        if log_max + k_rem < 63 || k_rem == basek {
 87            a.at_mut(col, size - 1)[idx] = data;
 88        } else {
 89            let mask: i64 = (1 << basek) - 1;
 90            let steps: usize = size.min(log_max.div_ceil(basek));
 91            (size - steps..size)
 92                .rev()
 93                .enumerate()
 94                .for_each(|(j, j_rev)| {
 95                    a.at_mut(col, j_rev)[idx] = (data >> (j * basek)) & mask;
 96                })
 97        }
 98
 99        // Case where prec % k != 0.
100        if k_rem != basek {
101            let steps: usize = size.min(log_max.div_ceil(basek));
102            (size - steps..size).rev().for_each(|j| {
103                a.at_mut(col, j)[idx] <<= k_rem;
104            })
105        }
106    }
107}
108
109impl<D: DataRef> VecZnx<D> {
110    pub fn decode_vec_i64(&self, basek: usize, col: usize, k: usize, data: &mut [i64]) {
111        let size: usize = k.div_ceil(basek);
112        #[cfg(debug_assertions)]
113        {
114            let a: VecZnx<&[u8]> = self.to_ref();
115            assert!(
116                data.len() >= a.n(),
117                "invalid data: data.len()={} < a.n()={}",
118                data.len(),
119                a.n()
120            );
121            assert!(col < a.cols());
122        }
123
124        let a: VecZnx<&[u8]> = self.to_ref();
125        data.copy_from_slice(a.at(col, 0));
126        let rem: usize = basek - (k % basek);
127        if k < basek {
128            data.iter_mut().for_each(|x| *x >>= rem);
129        } else {
130            (1..size).for_each(|i| {
131                if i == size - 1 && rem != basek {
132                    let k_rem: usize = basek - rem;
133                    izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
134                        *y = (*y << k_rem) + (x >> rem);
135                    });
136                } else {
137                    izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
138                        *y = (*y << basek) + x;
139                    });
140                }
141            })
142        }
143    }
144
145    pub fn decode_coeff_i64(&self, basek: usize, col: usize, k: usize, idx: usize) -> i64 {
146        #[cfg(debug_assertions)]
147        {
148            let a: VecZnx<&[u8]> = self.to_ref();
149            assert!(idx < a.n());
150            assert!(col < a.cols())
151        }
152
153        let a: VecZnx<&[u8]> = self.to_ref();
154        let size: usize = k.div_ceil(basek);
155        let mut res: i64 = 0;
156        let rem: usize = basek - (k % basek);
157        (0..size).for_each(|j| {
158            let x: i64 = a.at(col, j)[idx];
159            if j == size - 1 && rem != basek {
160                let k_rem: usize = basek - rem;
161                res = (res << k_rem) + (x >> rem);
162            } else {
163                res = (res << basek) + x;
164            }
165        });
166        res
167    }
168
169    pub fn decode_vec_float(&self, basek: usize, col: usize, data: &mut [Float]) {
170        #[cfg(debug_assertions)]
171        {
172            let a: VecZnx<&[u8]> = self.to_ref();
173            assert!(
174                data.len() >= a.n(),
175                "invalid data: data.len()={} < a.n()={}",
176                data.len(),
177                a.n()
178            );
179            assert!(col < a.cols());
180        }
181
182        let a: VecZnx<&[u8]> = self.to_ref();
183        let size: usize = a.size();
184        let prec: u32 = (basek * size) as u32;
185
186        // 2^{basek}
187        let base = Float::with_val(prec, (1 << basek) as f64);
188
189        // y[i] = sum x[j][i] * 2^{-basek*j}
190        (0..size).for_each(|i| {
191            if i == 0 {
192                izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
193                    y.assign(*x);
194                    *y /= &base;
195                });
196            } else {
197                izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
198                    *y += Float::with_val(prec, *x);
199                    *y /= &base;
200                });
201            }
202        });
203    }
204}