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