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