poulpy/base2k/src/encoding.rs
base2k/src/encoding.rsRaw
1use crate::ffi::znx::znx_zero_i64_ref;
2use crate::{Infos, VecZnx, VecZnxApi};
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 let cols: usize = (log_k + log_base2k - 1) / log_base2k;
66
67 assert!(cols <= self.cols(), "invalid argument log_k: (log_k + self.log_base2k - 1)/self.log_base2k={} > self.cols()={}", cols, self.cols());
68
69 let size: usize = min(data.len(), self.n());
70 let log_k_rem: usize = log_base2k - (log_k % log_base2k);
71
72 // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
73 // values on the last limb.
74 // Else we decompose values base2k.
75 if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
76 (0..self.cols()).for_each(|i| unsafe {
77 znx_zero_i64_ref(size as u64, self.at_mut(i).as_mut_ptr());
78 });
79 self.at_mut(cols - 1)[..size].copy_from_slice(&data[..size]);
80 } else {
81 let mask: i64 = (1 << log_base2k) - 1;
82 let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
83
84 (0..steps).for_each(|i| unsafe {
85 znx_zero_i64_ref(size as u64, self.at_mut(i).as_mut_ptr());
86 });
87
88 (cols - steps..cols)
89 .rev()
90 .enumerate()
91 .for_each(|(i, i_rev)| {
92 let shift: usize = i * log_base2k;
93 izip!(self.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
94 .for_each(|(y, x)| *y = (x >> shift) & mask);
95 })
96 }
97
98 // Case where self.prec % self.k != 0.
99 if log_k_rem != log_base2k {
100 let cols = self.cols();
101 let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
102 (cols - steps..cols).rev().for_each(|i| {
103 self.at_mut(i)[..size]
104 .iter_mut()
105 .for_each(|x| *x <<= log_k_rem);
106 })
107 }
108 }
109
110 fn decode_vec_i64(&self, log_base2k: usize, log_k: usize, data: &mut [i64]) {
111 let cols: usize = (log_k + log_base2k - 1) / log_base2k;
112 assert!(
113 data.len() >= self.n,
114 "invalid data: data.len()={} < self.n()={}",
115 data.len(),
116 self.n
117 );
118 data.copy_from_slice(self.at(0));
119 let rem: usize = log_base2k - (log_k % log_base2k);
120 (1..cols).for_each(|i| {
121 if i == cols - 1 && rem != log_base2k {
122 let k_rem: usize = log_base2k - rem;
123 izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
124 *y = (*y << k_rem) + (x >> rem);
125 });
126 } else {
127 izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
128 *y = (*y << log_base2k) + x;
129 });
130 }
131 })
132 }
133
134 fn decode_vec_float(&self, log_base2k: usize, data: &mut [Float]) {
135 let cols: usize = self.cols();
136 assert!(
137 data.len() >= self.n(),
138 "invalid data: data.len()={} < self.n()={}",
139 data.len(),
140 self.n()
141 );
142
143 let prec: u32 = (log_base2k * cols) as u32;
144
145 // 2^{log_base2k}
146 let base = Float::with_val(prec, (1 << log_base2k) as f64);
147
148 // y[i] = sum x[j][i] * 2^{-log_base2k*j}
149 (0..cols).for_each(|i| {
150 if i == 0 {
151 izip!(self.at(cols - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
152 y.assign(*x);
153 *y /= &base;
154 });
155 } else {
156 izip!(self.at(cols - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
157 *y += Float::with_val(prec, *x);
158 *y /= &base;
159 });
160 }
161 });
162 }
163
164 fn encode_coeff_i64(
165 &mut self,
166 log_base2k: usize,
167 log_k: usize,
168 i: usize,
169 value: i64,
170 log_max: usize,
171 ) {
172 assert!(i < self.n());
173 let cols: usize = (log_k + log_base2k - 1) / log_base2k;
174 assert!(cols <= self.cols(), "invalid argument log_k: (log_k + self.log_base2k - 1)/self.log_base2k={} > self.cols()={}", cols, self.cols());
175 let log_k_rem: usize = log_base2k - (log_k % log_base2k);
176 let cols = self.cols();
177
178 // If 2^{log_base2k} * 2^{log_k_rem} < 2^{63}-1, then we can simply copy
179 // values on the last limb.
180 // Else we decompose values base2k.
181 if log_max + log_k_rem < 63 || log_k_rem == log_base2k {
182 (0..cols - 1).for_each(|j| self.at_mut(j)[i] = 0);
183
184 self.at_mut(self.cols() - 1)[i] = value;
185 } else {
186 let mask: i64 = (1 << log_base2k) - 1;
187 let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
188
189 (0..cols - steps).for_each(|j| self.at_mut(j)[i] = 0);
190
191 (cols - steps..cols)
192 .rev()
193 .enumerate()
194 .for_each(|(j, j_rev)| {
195 self.at_mut(j_rev)[i] = (value >> (j * log_base2k)) & mask;
196 })
197 }
198
199 // Case where self.prec % self.k != 0.
200 if log_k_rem != log_base2k {
201 let cols = self.cols();
202 let steps: usize = min(cols, (log_max + log_base2k - 1) / log_base2k);
203 (cols - steps..cols).rev().for_each(|j| {
204 self.at_mut(j)[i] <<= log_k_rem;
205 })
206 }
207 }
208
209 fn decode_coeff_i64(&self, log_base2k: usize, log_k: usize, i: usize) -> i64 {
210 let cols: usize = (log_k + log_base2k - 1) / log_base2k;
211 assert!(i < self.n());
212 let mut res: i64 = self.data[i];
213 let rem: usize = log_base2k - (log_k % log_base2k);
214 (1..cols).for_each(|i| {
215 let x = self.data[i * self.n];
216 if i == cols - 1 && rem != log_base2k {
217 let k_rem: usize = log_base2k - rem;
218 res = (res << k_rem) + (x >> rem);
219 } else {
220 res = (res << log_base2k) + x;
221 }
222 });
223 res
224 }
225}
226
227#[cfg(test)]
228mod tests {
229 use crate::{Encoding, VecZnx};
230 use itertools::izip;
231 use sampling::source::Source;
232
233 #[test]
234 fn test_set_get_i64_lo_norm() {
235 let n: usize = 8;
236 let log_base2k: usize = 17;
237 let cols: usize = 5;
238 let log_k: usize = cols * log_base2k - 5;
239 let mut a: VecZnx = VecZnx::new(n, cols);
240 let mut have: Vec<i64> = vec![i64::default(); n];
241 have.iter_mut()
242 .enumerate()
243 .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
244 a.encode_vec_i64(log_base2k, log_k, &have, 10);
245 let mut want = vec![i64::default(); n];
246 a.decode_vec_i64(log_base2k, log_k, &mut want);
247 izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
248 }
249
250 #[test]
251 fn test_set_get_i64_hi_norm() {
252 let n: usize = 8;
253 let log_base2k: usize = 17;
254 let cols: usize = 5;
255 let log_k: usize = cols * log_base2k - 5;
256 let mut a: VecZnx = VecZnx::new(n, cols);
257 let mut have: Vec<i64> = vec![i64::default(); n];
258 let mut source = Source::new([1; 32]);
259 have.iter_mut().for_each(|x| {
260 *x = source
261 .next_u64n(u64::MAX, u64::MAX)
262 .wrapping_sub(u64::MAX / 2 + 1) as i64;
263 });
264 a.encode_vec_i64(log_base2k, log_k, &have, 63);
265 //(0..a.cols()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
266 let mut want = vec![i64::default(); n];
267 //(0..a.cols()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
268 a.decode_vec_i64(log_base2k, log_k, &mut want);
269 izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
270 }
271}