poulpy/poulpy-hal/src/layouts/encoding.rs
poulpy-hal/src/layouts/encoding.rsRaw
1use itertools::izip;
2use rug::{Assign, Float};
3
4use crate::{
5 layouts::{DataMut, DataRef, VecZnx, VecZnxToMut, VecZnxToRef, Zn, ZnToMut, ZnToRef, ZnxInfos, ZnxView, ZnxViewMut},
6 reference::znx::{
7 ZnxNormalizeFinalStepInplace, ZnxNormalizeFirstStepInplace, ZnxNormalizeMiddleStepInplace, ZnxRef, ZnxZero,
8 get_carry_i128, get_digit_i128, znx_zero_ref,
9 },
10};
11
12impl<D: DataMut> VecZnx<D> {
13 pub fn encode_vec_i64(&mut self, base2k: usize, col: usize, k: usize, data: &[i64]) {
14 let size: usize = k.div_ceil(base2k);
15
16 #[cfg(debug_assertions)]
17 {
18 let a: VecZnx<&mut [u8]> = self.to_mut();
19 assert!(
20 size <= a.size(),
21 "invalid argument k.div_ceil(base2k)={} > a.size()={}",
22 size,
23 a.size()
24 );
25 assert!(col < a.cols());
26 assert!(data.len() == a.n())
27 }
28
29 let mut a: VecZnx<&mut [u8]> = self.to_mut();
30 let a_size: usize = a.size();
31
32 // Zeroes coefficients of the i-th column
33 for i in 0..a_size {
34 znx_zero_ref(a.at_mut(col, i));
35 }
36
37 // Copies the data on the correct limb
38 a.at_mut(col, size - 1).copy_from_slice(data);
39
40 let mut carry: Vec<i64> = vec![0i64; a.n()];
41 let k_rem: usize = (base2k - (k % base2k)) % base2k;
42
43 // Normalizes and shift if necessary.
44 for j in (0..size).rev() {
45 if j == size - 1 {
46 ZnxRef::znx_normalize_first_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
47 } else if j == 0 {
48 ZnxRef::znx_normalize_final_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
49 } else {
50 ZnxRef::znx_normalize_middle_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
51 }
52 }
53 }
54
55 pub fn encode_vec_i128(&mut self, base2k: usize, col: usize, k: usize, data: &[i128]) {
56 let size: usize = k.div_ceil(base2k);
57
58 #[cfg(debug_assertions)]
59 {
60 let a: VecZnx<&mut [u8]> = self.to_mut();
61 assert!(
62 size <= a.size(),
63 "invalid argument k.div_ceil(base2k)={} > a.size()={}",
64 size,
65 a.size()
66 );
67 assert!(col < a.cols());
68 assert!(data.len() == a.n())
69 }
70
71 let mut a: VecZnx<&mut [u8]> = self.to_mut();
72 let a_size: usize = a.size();
73
74 {
75 let mut carry_i128: Vec<i128> = vec![0i128; a.n()];
76 carry_i128.copy_from_slice(data);
77
78 for j in (0..size).rev() {
79 for (x, a) in izip!(a.at_mut(col, j).iter_mut(), carry_i128.iter_mut()) {
80 let digit: i128 = get_digit_i128(base2k, *a);
81 let carry: i128 = get_carry_i128(base2k, *a, digit);
82 *x = digit as i64;
83 *a = carry;
84 }
85 }
86 }
87
88 for j in size..a_size {
89 ZnxRef::znx_zero(a.at_mut(col, j));
90 }
91
92 let mut carry: Vec<i64> = vec![0i64; a.n()];
93 let k_rem: usize = (base2k - (k % base2k)) % base2k;
94
95 for j in (0..size).rev() {
96 if j == a_size - 1 {
97 ZnxRef::znx_normalize_first_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
98 } else if j == 0 {
99 ZnxRef::znx_normalize_final_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
100 } else {
101 ZnxRef::znx_normalize_middle_step_inplace(base2k, k_rem, a.at_mut(col, j), &mut carry);
102 }
103 }
104 }
105
106 pub fn encode_coeff_i64(&mut self, base2k: usize, col: usize, k: usize, idx: usize, data: i64) {
107 let size: usize = k.div_ceil(base2k);
108
109 #[cfg(debug_assertions)]
110 {
111 let a: VecZnx<&mut [u8]> = self.to_mut();
112 assert!(idx < a.n());
113 assert!(
114 size <= a.size(),
115 "invalid argument k.div_ceil(base2k)={} > a.size()={}",
116 size,
117 a.size()
118 );
119 assert!(col < a.cols());
120 }
121
122 let mut a: VecZnx<&mut [u8]> = self.to_mut();
123 let a_size = a.size();
124
125 for j in 0..a_size {
126 a.at_mut(col, j)[idx] = 0
127 }
128
129 a.at_mut(col, size - 1)[idx] = data;
130
131 let mut carry: Vec<i64> = vec![0i64; 1];
132 let k_rem: usize = (base2k - (k % base2k)) % base2k;
133
134 for j in (0..size).rev() {
135 let slice = &mut a.at_mut(col, j)[idx..idx + 1];
136
137 if j == size - 1 {
138 ZnxRef::znx_normalize_first_step_inplace(base2k, k_rem, slice, &mut carry);
139 } else if j == 0 {
140 ZnxRef::znx_normalize_final_step_inplace(base2k, k_rem, slice, &mut carry);
141 } else {
142 ZnxRef::znx_normalize_middle_step_inplace(base2k, k_rem, slice, &mut carry);
143 }
144 }
145 }
146}
147
148impl<D: DataRef> VecZnx<D> {
149 pub fn decode_vec_i64(&self, base2k: usize, col: usize, k: usize, data: &mut [i64]) {
150 let size: usize = k.div_ceil(base2k);
151 #[cfg(debug_assertions)]
152 {
153 let a: VecZnx<&[u8]> = self.to_ref();
154 assert!(
155 data.len() >= a.n(),
156 "invalid data: data.len()={} < a.n()={}",
157 data.len(),
158 a.n()
159 );
160 assert!(col < a.cols());
161 }
162
163 let a: VecZnx<&[u8]> = self.to_ref();
164 data.copy_from_slice(a.at(col, 0));
165 let rem: usize = base2k - (k % base2k);
166 if k < base2k {
167 let scale = 1 << rem as i64;
168 data.iter_mut().for_each(|x| *x = div_round(*x, scale));
169 } else {
170 (1..size).for_each(|i| {
171 if i == size - 1 && rem != base2k {
172 let k_rem: usize = (base2k - rem) % base2k;
173 let scale: i64 = 1 << rem as i64;
174 izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
175 *y = (*y << k_rem) + div_round(*x, scale);
176 });
177 } else {
178 izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
179 *y = (*y << base2k) + x;
180 });
181 }
182 })
183 }
184 }
185
186 pub fn decode_coeff_i64(&self, base2k: usize, col: usize, k: usize, idx: usize) -> i64 {
187 #[cfg(debug_assertions)]
188 {
189 let a: VecZnx<&[u8]> = self.to_ref();
190 assert!(idx < a.n());
191 assert!(col < a.cols())
192 }
193
194 let a: VecZnx<&[u8]> = self.to_ref();
195 let size: usize = k.div_ceil(base2k);
196 let mut res: i64 = 0;
197 let rem: usize = base2k - (k % base2k);
198 (0..size).for_each(|j| {
199 let x: i64 = a.at(col, j)[idx];
200 if j == size - 1 && rem != base2k {
201 let k_rem: usize = (base2k - rem) % base2k;
202 let scale: i64 = 1 << rem as i64;
203 res = (res << k_rem) + div_round(x, scale);
204 } else {
205 res = (res << base2k) + x;
206 }
207 });
208 res
209 }
210
211 pub fn decode_vec_float(&self, base2k: usize, col: usize, data: &mut [Float]) {
212 #[cfg(debug_assertions)]
213 {
214 let a: VecZnx<&[u8]> = self.to_ref();
215 assert!(
216 data.len() >= a.n(),
217 "invalid data: data.len()={} < a.n()={}",
218 data.len(),
219 a.n()
220 );
221 assert!(col < a.cols());
222 }
223
224 let a: VecZnx<&[u8]> = self.to_ref();
225 let size: usize = a.size();
226 let prec: u32 = (base2k * size) as u32;
227
228 // 2^{base2k}
229 let base: Float = Float::with_val(prec, (1u64 << base2k) as f64);
230
231 // y[i] = sum x[j][i] * 2^{-base2k*j}
232 (0..size).for_each(|i| {
233 if i == 0 {
234 izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
235 y.assign(*x);
236 *y /= &base;
237 });
238 } else {
239 izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
240 *y += Float::with_val(prec, *x);
241 *y /= &base;
242 });
243 }
244 });
245 }
246}
247
248impl<D: DataMut> Zn<D> {
249 pub fn encode_i64(&mut self, base2k: usize, k: usize, data: i64) {
250 let size: usize = k.div_ceil(base2k);
251
252 #[cfg(debug_assertions)]
253 {
254 let a: Zn<&mut [u8]> = self.to_mut();
255 assert!(
256 size <= a.size(),
257 "invalid argument k.div_ceil(base2k)={} > a.size()={}",
258 size,
259 a.size()
260 );
261 }
262
263 let mut a: Zn<&mut [u8]> = self.to_mut();
264 let a_size = a.size();
265
266 for j in 0..a_size {
267 a.at_mut(0, j)[0] = 0
268 }
269
270 a.at_mut(0, size - 1)[0] = data;
271
272 let mut carry: Vec<i64> = vec![0i64; 1];
273 let k_rem: usize = (base2k - (k % base2k)) % base2k;
274
275 for j in (0..size).rev() {
276 let slice = &mut a.at_mut(0, j)[..1];
277
278 if j == size - 1 {
279 ZnxRef::znx_normalize_first_step_inplace(base2k, k_rem, slice, &mut carry);
280 } else if j == 0 {
281 ZnxRef::znx_normalize_final_step_inplace(base2k, k_rem, slice, &mut carry);
282 } else {
283 ZnxRef::znx_normalize_middle_step_inplace(base2k, k_rem, slice, &mut carry);
284 }
285 }
286 }
287}
288
289impl<D: DataRef> Zn<D> {
290 pub fn decode_i64(&self, base2k: usize, k: usize) -> i64 {
291 let a: Zn<&[u8]> = self.to_ref();
292 let size: usize = k.div_ceil(base2k);
293 let mut res: i64 = 0;
294 let rem: usize = base2k - (k % base2k);
295 (0..size).for_each(|j| {
296 let x: i64 = a.at(0, j)[0];
297 if j == size - 1 && rem != base2k {
298 let k_rem: usize = (base2k - rem) % base2k;
299 let scale: i64 = 1 << rem as i64;
300 res = (res << k_rem) + div_round(x, scale);
301 } else {
302 res = (res << base2k) + x;
303 }
304 });
305 res
306 }
307
308 pub fn decode_float(&self, base2k: usize) -> Float {
309 let a: Zn<&[u8]> = self.to_ref();
310 let size: usize = a.size();
311 let prec: u32 = (base2k * size) as u32;
312
313 // 2^{base2k}
314 let base: Float = Float::with_val(prec, (1 << base2k) as f64);
315 let mut res: Float = Float::with_val(prec, (1 << base2k) as f64);
316
317 // y[i] = sum x[j][i] * 2^{-base2k*j}
318 (0..size).for_each(|i| {
319 if i == 0 {
320 res.assign(a.at(0, size - i - 1)[0]);
321 res /= &base;
322 } else {
323 res += Float::with_val(prec, a.at(0, size - i - 1)[0]);
324 res /= &base;
325 }
326 });
327
328 res
329 }
330}
331
332#[inline]
333pub fn div_round(a: i64, b: i64) -> i64 {
334 assert!(b != 0, "division by zero");
335 let div: i64 = a / b;
336 let rem: i64 = a % b;
337 if (2 * rem.abs()) >= b.abs() {
338 div + a.signum() * b.signum()
339 } else {
340 div
341 }
342}