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 data.iter_mut().for_each(|x| *x >>= rem);
168 } else {
169 (1..size).for_each(|i| {
170 if i == size - 1 && rem != base2k {
171 let k_rem: usize = (base2k - rem) % base2k;
172 izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
173 *y = (*y << k_rem) + (x >> rem);
174 });
175 } else {
176 izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
177 *y = (*y << base2k) + x;
178 });
179 }
180 })
181 }
182 }
183
184 pub fn decode_coeff_i64(&self, base2k: usize, col: usize, k: usize, idx: usize) -> i64 {
185 #[cfg(debug_assertions)]
186 {
187 let a: VecZnx<&[u8]> = self.to_ref();
188 assert!(idx < a.n());
189 assert!(col < a.cols())
190 }
191
192 let a: VecZnx<&[u8]> = self.to_ref();
193 let size: usize = k.div_ceil(base2k);
194 let mut res: i64 = 0;
195 let rem: usize = base2k - (k % base2k);
196 (0..size).for_each(|j| {
197 let x: i64 = a.at(col, j)[idx];
198 if j == size - 1 && rem != base2k {
199 let k_rem: usize = (base2k - rem) % base2k;
200 res = (res << k_rem) + (x >> rem);
201 } else {
202 res = (res << base2k) + x;
203 }
204 });
205 res
206 }
207
208 pub fn decode_vec_float(&self, base2k: usize, col: usize, data: &mut [Float]) {
209 #[cfg(debug_assertions)]
210 {
211 let a: VecZnx<&[u8]> = self.to_ref();
212 assert!(
213 data.len() >= a.n(),
214 "invalid data: data.len()={} < a.n()={}",
215 data.len(),
216 a.n()
217 );
218 assert!(col < a.cols());
219 }
220
221 let a: VecZnx<&[u8]> = self.to_ref();
222 let size: usize = a.size();
223 let prec: u32 = (base2k * size) as u32;
224
225 // 2^{base2k}
226 let base: Float = Float::with_val(prec, (1u64 << base2k) as f64);
227
228 // y[i] = sum x[j][i] * 2^{-base2k*j}
229 (0..size).for_each(|i| {
230 if i == 0 {
231 izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
232 y.assign(*x);
233 *y /= &base;
234 });
235 } else {
236 izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
237 *y += Float::with_val(prec, *x);
238 *y /= &base;
239 });
240 }
241 });
242 }
243}
244
245impl<D: DataMut> Zn<D> {
246 pub fn encode_i64(&mut self, base2k: usize, k: usize, data: i64) {
247 let size: usize = k.div_ceil(base2k);
248
249 #[cfg(debug_assertions)]
250 {
251 let a: Zn<&mut [u8]> = self.to_mut();
252 assert!(
253 size <= a.size(),
254 "invalid argument k.div_ceil(base2k)={} > a.size()={}",
255 size,
256 a.size()
257 );
258 }
259
260 let mut a: Zn<&mut [u8]> = self.to_mut();
261 let a_size = a.size();
262
263 for j in 0..a_size {
264 a.at_mut(0, j)[0] = 0
265 }
266
267 a.at_mut(0, size - 1)[0] = data;
268
269 let mut carry: Vec<i64> = vec![0i64; 1];
270 let k_rem: usize = (base2k - (k % base2k)) % base2k;
271
272 for j in (0..size).rev() {
273 let slice = &mut a.at_mut(0, j)[..1];
274
275 if j == size - 1 {
276 ZnxRef::znx_normalize_first_step_inplace(base2k, k_rem, slice, &mut carry);
277 } else if j == 0 {
278 ZnxRef::znx_normalize_final_step_inplace(base2k, k_rem, slice, &mut carry);
279 } else {
280 ZnxRef::znx_normalize_middle_step_inplace(base2k, k_rem, slice, &mut carry);
281 }
282 }
283 }
284}
285
286impl<D: DataRef> Zn<D> {
287 pub fn decode_i64(&self, base2k: usize, k: usize) -> i64 {
288 let a: Zn<&[u8]> = self.to_ref();
289 let size: usize = k.div_ceil(base2k);
290 let mut res: i64 = 0;
291 let rem: usize = base2k - (k % base2k);
292 (0..size).for_each(|j| {
293 let x: i64 = a.at(0, j)[0];
294 if j == size - 1 && rem != base2k {
295 let k_rem: usize = (base2k - rem) % base2k;
296 res = (res << k_rem) + (x >> rem);
297 } else {
298 res = (res << base2k) + x;
299 }
300 });
301 res
302 }
303
304 pub fn decode_float(&self, base2k: usize) -> Float {
305 let a: Zn<&[u8]> = self.to_ref();
306 let size: usize = a.size();
307 let prec: u32 = (base2k * size) as u32;
308
309 // 2^{base2k}
310 let base: Float = Float::with_val(prec, (1 << base2k) as f64);
311 let mut res: Float = Float::with_val(prec, (1 << base2k) as f64);
312
313 // y[i] = sum x[j][i] * 2^{-base2k*j}
314 (0..size).for_each(|i| {
315 if i == 0 {
316 res.assign(a.at(0, size - i - 1)[0]);
317 res /= &base;
318 } else {
319 res += Float::with_val(prec, a.at(0, size - i - 1)[0]);
320 res /= &base;
321 }
322 });
323
324 res
325 }
326}