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