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 /// 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 /// * `log_base2k`: base two negative logarithm decomposition of the receiver.
26 /// * `log_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, log_base2k: usize, log_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 /// * `log_base2k`: base two negative logarithm decomposition of the receiver.
40 /// * `log_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, log_base2k: usize, log_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 /// * `log_base2k`: 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, log_base2k: 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 /// * `log_base2k`: base two negative logarithm decomposition of the receiver.
58 /// * `log_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, log_base2k: usize, log_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, log_base2k: usize, log_k: usize, data: &[i64], log_max: usize) {
66 encode_vec_i64(self, col_i, log_base2k, log_k, data, log_max)
67 }
68
69 fn encode_coeff_i64(&mut self, col_i: usize, log_base2k: usize, log_k: usize, i: usize, value: i64, log_max: usize) {
70 encode_coeff_i64(self, col_i, log_base2k, log_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, log_base2k: usize, log_k: usize, data: &mut [i64]) {
76 decode_vec_i64(self, col_i, log_base2k, log_k, data)
77 }
78
79 fn decode_vec_float(&self, col_i: usize, log_base2k: usize, data: &mut [Float]) {
80 decode_vec_float(self, col_i, log_base2k, data)
81 }
82
83 fn decode_coeff_i64(&self, col_i: usize, log_base2k: usize, log_k: usize, i: usize) -> i64 {
84 decode_coeff_i64(self, col_i, log_base2k, log_k, i)
85 }
86}
87
88fn encode_vec_i64<D: AsMut<[u8]> + AsRef<[u8]>>(
89 a: &mut VecZnx<D>,
90 col_i: usize,
91 log_base2k: usize,
92 log_k: usize,
93 data: &[i64],
94 log_max: usize,
95) {
96 let size: usize = (log_k + log_base2k - 1) / log_base2k;
97
98 #[cfg(debug_assertions)]
99 {
100 assert!(
101 size <= a.size(),
102 "invalid argument log_k: (log_k + a.log_base2k - 1)/a.log_base2k={} > 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 log_k_rem: usize = log_base2k - (log_k % log_base2k);
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^{log_base2k} * 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 + log_k_rem < 63 || log_k_rem == log_base2k {
122 a.at_mut(col_i, size - 1)[..data_len].copy_from_slice(&data[..data_len]);
123 } else {
124 let mask: i64 = (1 << log_base2k) - 1;
125 let steps: usize = min(size, (log_max + log_base2k - 1) / log_base2k);
126 (size - steps..size)
127 .rev()
128 .enumerate()
129 .for_each(|(i, i_rev)| {
130 let shift: usize = i * log_base2k;
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 log_k_rem != log_base2k {
137 let steps: usize = min(size, (log_max + log_base2k - 1) / log_base2k);
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 <<= log_k_rem);
142 })
143 }
144}
145
146fn decode_vec_i64<D: AsRef<[u8]>>(a: &VecZnx<D>, col_i: usize, log_base2k: usize, log_k: usize, data: &mut [i64]) {
147 let size: usize = (log_k + log_base2k - 1) / log_base2k;
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 = log_base2k - (log_k % log_base2k);
160 (1..size).for_each(|i| {
161 if i == size - 1 && rem != log_base2k {
162 let k_rem: usize = log_base2k - 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 << log_base2k) + x;
169 });
170 }
171 })
172}
173
174fn decode_vec_float<D: AsRef<[u8]>>(a: &VecZnx<D>, col_i: usize, log_base2k: 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 = (log_base2k * size) as u32;
188
189 // 2^{log_base2k}
190 let base = Float::with_val(prec, (1 << log_base2k) as f64);
191
192 // y[i] = sum x[j][i] * 2^{-log_base2k*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 log_base2k: usize,
212 log_k: usize,
213 i: usize,
214 value: i64,
215 log_max: usize,
216) {
217 let size: usize = (log_k + log_base2k - 1) / log_base2k;
218
219 #[cfg(debug_assertions)]
220 {
221 assert!(i < a.n());
222 assert!(
223 size <= a.size(),
224 "invalid argument log_k: (log_k + a.log_base2k - 1)/a.log_base2k={} > a.size()={}",
225 size,
226 a.size()
227 );
228 assert!(col_i < a.cols());
229 }
230
231 let log_k_rem: usize = log_base2k - (log_k % log_base2k);
232 (0..a.size()).for_each(|j| a.at_mut(col_i, j)[i] = 0);
233
234 // If 2^{log_base2k} * 2^{log_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 + log_k_rem < 63 || log_k_rem == log_base2k {
238 a.at_mut(col_i, size - 1)[i] = value;
239 } else {
240 let mask: i64 = (1 << log_base2k) - 1;
241 let steps: usize = min(size, (log_max + log_base2k - 1) / log_base2k);
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 * log_base2k)) & mask;
247 })
248 }
249
250 // Case where prec % k != 0.
251 if log_k_rem != log_base2k {
252 let steps: usize = min(size, (log_max + log_base2k - 1) / log_base2k);
253 (size - steps..size).rev().for_each(|j| {
254 a.at_mut(col_i, j)[i] <<= log_k_rem;
255 })
256 }
257}
258
259fn decode_coeff_i64<D: AsRef<[u8]>>(a: &VecZnx<D>, col_i: usize, log_base2k: usize, log_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 = (log_k + log_base2k - 1) / log_base2k;
267 let data: &[i64] = a.raw();
268 let mut res: i64 = data[i];
269 let rem: usize = log_base2k - (log_k % log_base2k);
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 != log_base2k {
274 let k_rem: usize = log_base2k - rem;
275 res = (res << k_rem) + (x >> rem);
276 } else {
277 res = (res << log_base2k) + 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 log_base2k: usize = 17;
296 let size: usize = 5;
297 let log_k: usize = size * log_base2k - 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, log_base2k, log_k, &have, 10);
307 let mut want: Vec<i64> = vec![i64::default(); n];
308 a.decode_vec_i64(col_i, log_base2k, log_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 log_base2k: usize = 17;
318 let size: usize = 5;
319 let log_k: usize = size * log_base2k - 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, log_base2k, log_k, &have, 64);
328 let mut want = vec![i64::default(); n];
329 a.decode_vec_i64(col_i, log_base2k, log_k, &mut want);
330 izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
331 })
332 }
333}