poulpy/poulpy-hal/src/layouts/encoding.rs
poulpy-hal/src/layouts/encoding.rsRaw
1use itertools::izip;
2use rug::{Assign, Float};
3
4use crate::layouts::{
5 DataMut, DataRef, VecZnx, VecZnxToMut, VecZnxToRef, Zn, ZnToMut, ZnToRef, ZnxInfos, ZnxView, ZnxViewMut, ZnxZero,
6};
7
8impl<D: DataMut> VecZnx<D> {
9 pub fn encode_vec_i64(&mut self, basek: usize, col: usize, k: usize, data: &[i64], log_max: usize) {
10 let size: usize = k.div_ceil(basek);
11
12 #[cfg(debug_assertions)]
13 {
14 let a: VecZnx<&mut [u8]> = self.to_mut();
15 assert!(
16 size <= a.size(),
17 "invalid argument k.div_ceil(basek)={} > a.size()={}",
18 size,
19 a.size()
20 );
21 assert!(col < a.cols());
22 assert!(data.len() <= a.n())
23 }
24
25 let data_len: usize = data.len();
26 let mut a: VecZnx<&mut [u8]> = self.to_mut();
27 let k_rem: usize = basek - (k % basek);
28
29 // Zeroes coefficients of the i-th column
30 (0..a.size()).for_each(|i| {
31 a.zero_at(col, i);
32 });
33
34 // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
35 // values on the last limb.
36 // Else we decompose values base2k.
37 if log_max + k_rem < 63 || k_rem == basek {
38 a.at_mut(col, size - 1)[..data_len].copy_from_slice(&data[..data_len]);
39 } else {
40 let mask: i64 = (1 << basek) - 1;
41 let steps: usize = size.min(log_max.div_ceil(basek));
42 (size - steps..size)
43 .rev()
44 .enumerate()
45 .for_each(|(i, i_rev)| {
46 let shift: usize = i * basek;
47 izip!(a.at_mut(col, i_rev).iter_mut(), data.iter()).for_each(|(y, x)| *y = (x >> shift) & mask);
48 })
49 }
50
51 // Case where self.prec % self.k != 0.
52 if k_rem != basek {
53 let steps: usize = size.min(log_max.div_ceil(basek));
54 (size - steps..size).rev().for_each(|i| {
55 a.at_mut(col, i)[..data_len]
56 .iter_mut()
57 .for_each(|x| *x <<= k_rem);
58 })
59 }
60 }
61
62 pub fn encode_coeff_i64(&mut self, basek: usize, col: usize, k: usize, idx: usize, data: i64, log_max: usize) {
63 let size: usize = k.div_ceil(basek);
64
65 #[cfg(debug_assertions)]
66 {
67 let a: VecZnx<&mut [u8]> = self.to_mut();
68 assert!(idx < a.n());
69 assert!(
70 size <= a.size(),
71 "invalid argument k.div_ceil(basek)={} > a.size()={}",
72 size,
73 a.size()
74 );
75 assert!(col < a.cols());
76 }
77
78 let k_rem: usize = basek - (k % basek);
79 let mut a: VecZnx<&mut [u8]> = self.to_mut();
80 (0..a.size()).for_each(|j| a.at_mut(col, j)[idx] = 0);
81
82 // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
83 // values on the last limb.
84 // Else we decompose values base2k.
85 if log_max + k_rem < 63 || k_rem == basek {
86 a.at_mut(col, size - 1)[idx] = data;
87 } else {
88 let mask: i64 = (1 << basek) - 1;
89 let steps: usize = size.min(log_max.div_ceil(basek));
90 (size - steps..size)
91 .rev()
92 .enumerate()
93 .for_each(|(j, j_rev)| {
94 a.at_mut(col, j_rev)[idx] = (data >> (j * basek)) & mask;
95 })
96 }
97
98 // Case where prec % k != 0.
99 if k_rem != basek {
100 let steps: usize = size.min(log_max.div_ceil(basek));
101 (size - steps..size).rev().for_each(|j| {
102 a.at_mut(col, j)[idx] <<= k_rem;
103 })
104 }
105 }
106}
107
108impl<D: DataRef> VecZnx<D> {
109 pub fn decode_vec_i64(&self, basek: usize, col: usize, k: usize, data: &mut [i64]) {
110 let size: usize = k.div_ceil(basek);
111 #[cfg(debug_assertions)]
112 {
113 let a: VecZnx<&[u8]> = self.to_ref();
114 assert!(
115 data.len() >= a.n(),
116 "invalid data: data.len()={} < a.n()={}",
117 data.len(),
118 a.n()
119 );
120 assert!(col < a.cols());
121 }
122
123 let a: VecZnx<&[u8]> = self.to_ref();
124 data.copy_from_slice(a.at(col, 0));
125 let rem: usize = basek - (k % basek);
126 if k < basek {
127 data.iter_mut().for_each(|x| *x >>= rem);
128 } else {
129 (1..size).for_each(|i| {
130 if i == size - 1 && rem != basek {
131 let k_rem: usize = basek - rem;
132 izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
133 *y = (*y << k_rem) + (x >> rem);
134 });
135 } else {
136 izip!(a.at(col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
137 *y = (*y << basek) + x;
138 });
139 }
140 })
141 }
142 }
143
144 pub fn decode_coeff_i64(&self, basek: usize, col: usize, k: usize, idx: usize) -> i64 {
145 #[cfg(debug_assertions)]
146 {
147 let a: VecZnx<&[u8]> = self.to_ref();
148 assert!(idx < a.n());
149 assert!(col < a.cols())
150 }
151
152 let a: VecZnx<&[u8]> = self.to_ref();
153 let size: usize = k.div_ceil(basek);
154 let mut res: i64 = 0;
155 let rem: usize = basek - (k % basek);
156 (0..size).for_each(|j| {
157 let x: i64 = a.at(col, j)[idx];
158 if j == size - 1 && rem != basek {
159 let k_rem: usize = basek - rem;
160 res = (res << k_rem) + (x >> rem);
161 } else {
162 res = (res << basek) + x;
163 }
164 });
165 res
166 }
167
168 pub fn decode_vec_float(&self, basek: usize, col: usize, data: &mut [Float]) {
169 #[cfg(debug_assertions)]
170 {
171 let a: VecZnx<&[u8]> = self.to_ref();
172 assert!(
173 data.len() >= a.n(),
174 "invalid data: data.len()={} < a.n()={}",
175 data.len(),
176 a.n()
177 );
178 assert!(col < a.cols());
179 }
180
181 let a: VecZnx<&[u8]> = self.to_ref();
182 let size: usize = a.size();
183 let prec: u32 = (basek * size) as u32;
184
185 // 2^{basek}
186 let base = Float::with_val(prec, (1 << basek) as f64);
187
188 // y[i] = sum x[j][i] * 2^{-basek*j}
189 (0..size).for_each(|i| {
190 if i == 0 {
191 izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
192 y.assign(*x);
193 *y /= &base;
194 });
195 } else {
196 izip!(a.at(col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
197 *y += Float::with_val(prec, *x);
198 *y /= &base;
199 });
200 }
201 });
202 }
203}
204
205impl<D: DataMut> Zn<D> {
206 pub fn encode_i64(&mut self, basek: usize, k: usize, data: i64, log_max: usize) {
207 let size: usize = k.div_ceil(basek);
208
209 #[cfg(debug_assertions)]
210 {
211 let a: Zn<&mut [u8]> = self.to_mut();
212 assert!(
213 size <= a.size(),
214 "invalid argument k.div_ceil(basek)={} > a.size()={}",
215 size,
216 a.size()
217 );
218 }
219
220 let k_rem: usize = basek - (k % basek);
221 let mut a: Zn<&mut [u8]> = self.to_mut();
222 (0..a.size()).for_each(|j| a.at_mut(0, j)[0] = 0);
223
224 // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
225 // values on the last limb.
226 // Else we decompose values base2k.
227 if log_max + k_rem < 63 || k_rem == basek {
228 a.at_mut(0, size - 1)[0] = data;
229 } else {
230 let mask: i64 = (1 << basek) - 1;
231 let steps: usize = size.min(log_max.div_ceil(basek));
232 (size - steps..size)
233 .rev()
234 .enumerate()
235 .for_each(|(j, j_rev)| {
236 a.at_mut(0, j_rev)[0] = (data >> (j * basek)) & mask;
237 })
238 }
239
240 // Case where prec % k != 0.
241 if k_rem != basek {
242 let steps: usize = size.min(log_max.div_ceil(basek));
243 (size - steps..size).rev().for_each(|j| {
244 a.at_mut(0, j)[0] <<= k_rem;
245 })
246 }
247 }
248}
249
250impl<D: DataRef> Zn<D> {
251 pub fn decode_i64(&self, basek: usize, k: usize) -> i64 {
252 let a: Zn<&[u8]> = self.to_ref();
253 let size: usize = k.div_ceil(basek);
254 let mut res: i64 = 0;
255 let rem: usize = basek - (k % basek);
256 (0..size).for_each(|j| {
257 let x: i64 = a.at(0, j)[0];
258 if j == size - 1 && rem != basek {
259 let k_rem: usize = basek - rem;
260 res = (res << k_rem) + (x >> rem);
261 } else {
262 res = (res << basek) + x;
263 }
264 });
265 res
266 }
267
268 pub fn decode_float(&self, basek: usize) -> Float {
269 let a: Zn<&[u8]> = self.to_ref();
270 let size: usize = a.size();
271 let prec: u32 = (basek * size) as u32;
272
273 // 2^{basek}
274 let base: Float = Float::with_val(prec, (1 << basek) as f64);
275 let mut res: Float = Float::with_val(prec, (1 << basek) as f64);
276
277 // y[i] = sum x[j][i] * 2^{-basek*j}
278 (0..size).for_each(|i| {
279 if i == 0 {
280 res.assign(a.at(0, size - i - 1)[0]);
281 res /= &base;
282 } else {
283 res += Float::with_val(prec, a.at(0, size - i - 1)[0]);
284 res /= &base;
285 }
286 });
287
288 res
289 }
290}