poulpy/spqlios/src/poly.rs
spqlios/src/poly.rsRaw
1use crate::{
2 cast_mut_u8_to_mut_i64_slice, znx_automorphism_i64, znx_automorphism_inplace_i64,
3 znx_normalize, znx_zero_i64_ref,
4};
5use itertools::izip;
6use rand_distr::{Distribution, Normal};
7use sampling::source::Source;
8use std::cmp::min;
9
10pub struct Poly {
11 pub n: usize,
12 pub log_base2k: usize,
13 pub prec: usize,
14 pub data: Vec<i64>,
15}
16
17impl Poly {
18 pub fn new(n: usize, log_base2k: usize, prec: usize) -> Self {
19 Self {
20 n: n,
21 log_base2k: log_base2k,
22 prec: prec,
23 data: vec![i64::default(); Self::buffer_size(n, log_base2k, prec)],
24 }
25 }
26
27 pub fn buffer_size(n: usize, log_base2k: usize, prec: usize) -> usize {
28 n * ((prec + log_base2k - 1) / log_base2k)
29 }
30
31 pub fn from_buffer(&mut self, n: usize, log_base2k: usize, prec: usize, buf: &[i64]) {
32 let size = Self::buffer_size(n, log_base2k, prec);
33 assert!(
34 buf.len() >= size,
35 "invalid buffer: buf.len()={} < self.buffer_size(n={}, k={}, prec={})={}",
36 buf.len(),
37 n,
38 log_base2k,
39 prec,
40 size
41 );
42 self.n = n;
43 self.log_base2k = log_base2k;
44 self.prec = prec;
45 self.data = Vec::from(&buf[..size])
46 }
47
48 pub fn log_n(&self) -> u64 {
49 (u64::BITS - (self.n - 1).leading_zeros()) as _
50 }
51
52 pub fn n(&self) -> usize {
53 self.n
54 }
55
56 pub fn prec(&self) -> usize {
57 self.prec
58 }
59
60 pub fn limbs(&self) -> usize {
61 self.data.len() / self.n
62 }
63
64 pub fn as_ptr(&self) -> *const i64 {
65 self.data.as_ptr()
66 }
67
68 pub fn as_mut_ptr(&mut self) -> *mut i64 {
69 self.data.as_mut_ptr()
70 }
71
72 pub fn at(&self, i: usize) -> &[i64] {
73 &self.data[i * self.n..(i + 1) * self.n]
74 }
75
76 pub fn at_ptr(&self, i: usize) -> *const i64 {
77 &self.data[i * self.n] as *const i64
78 }
79
80 pub fn at_mut_ptr(&mut self, i: usize) -> *mut i64 {
81 &mut self.data[i * self.n] as *mut i64
82 }
83
84 pub fn at_mut(&mut self, i: usize) -> &mut [i64] {
85 &mut self.data[i * self.n..(i + 1) * self.n]
86 }
87
88 pub fn set_i64(&mut self, data: &[i64], log_max: usize) {
89 let size: usize = min(data.len(), self.n());
90 let k_rem: usize = self.log_base2k - (self.prec % self.log_base2k);
91
92 // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
93 // values on the last limb.
94 // Else we decompose values base2k.
95 if log_max + k_rem < 63 || k_rem == self.log_base2k {
96 self.at_mut(self.limbs() - 1).copy_from_slice(&data[..size]);
97 } else {
98 let mask: i64 = (1 << self.log_base2k) - 1;
99 let limbs = self.limbs();
100 let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
101 (limbs - steps..limbs)
102 .rev()
103 .enumerate()
104 .for_each(|(i, i_rev)| {
105 let shift: usize = i * self.log_base2k;
106 izip!(self.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
107 .for_each(|(y, x)| *y = (x >> shift) & mask);
108 })
109 }
110
111 // Case where self.prec % self.k != 0.
112 if k_rem != self.log_base2k {
113 let limbs = self.limbs();
114 let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
115 (limbs - steps..limbs).rev().for_each(|i| {
116 self.at_mut(i)[..size].iter_mut().for_each(|x| *x <<= k_rem);
117 })
118 }
119 }
120
121 pub fn normalize(&mut self, carry: &mut [u8]) {
122 assert!(
123 carry.len() >= self.n * 8,
124 "invalid carry: carry.len()={} < self.n()={}",
125 carry.len(),
126 self.n()
127 );
128
129 let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
130
131 unsafe {
132 znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
133 (0..self.limbs()).rev().for_each(|i| {
134 znx_normalize(
135 self.n as u64,
136 self.log_base2k as u64,
137 self.at_mut_ptr(i),
138 carry_i64.as_mut_ptr(),
139 self.at_mut_ptr(i),
140 carry_i64.as_mut_ptr(),
141 )
142 });
143 }
144 }
145
146 pub fn get_i64(&self, data: &mut [i64]) {
147 assert!(
148 data.len() >= self.n,
149 "invalid data: data.len()={} < self.n()={}",
150 data.len(),
151 self.n
152 );
153 data.copy_from_slice(self.at(0));
154 let rem: usize = self.log_base2k - (self.prec % self.log_base2k);
155 (1..self.limbs()).for_each(|i| {
156 if i == self.limbs() - 1 && rem != self.log_base2k {
157 let k_rem: usize = self.log_base2k - rem;
158 izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
159 *y = (*y << k_rem) + (x >> rem);
160 });
161 } else {
162 izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
163 *y = (*y << self.log_base2k) + x;
164 });
165 }
166 })
167 }
168
169 pub fn automorphism_inplace(&mut self, gal_el: i64) {
170 unsafe {
171 (0..self.limbs()).for_each(|i| {
172 znx_automorphism_inplace_i64(self.n as u64, gal_el, self.at_mut_ptr(i))
173 })
174 }
175 }
176 pub fn automorphism(&mut self, gal_el: i64, a: &mut Poly) {
177 unsafe {
178 (0..self.limbs()).for_each(|i| {
179 znx_automorphism_i64(self.n as u64, gal_el, a.at_mut_ptr(i), self.at_ptr(i))
180 })
181 }
182 }
183
184 pub fn fill_uniform(&mut self, source: &mut Source) {
185 let mut base2k: u64 = 1 << self.log_base2k;
186 let mut mask: u64 = base2k - 1;
187 let mut base2k_half: i64 = (base2k >> 1) as i64;
188
189 let size: usize = self.n() * (self.limbs() - 1);
190
191 self.data[..size]
192 .iter_mut()
193 .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
194
195 let log_base2k_rem: usize = self.prec % self.log_base2k;
196
197 if log_base2k_rem != 0 {
198 base2k = 1 << log_base2k_rem;
199 mask = (base2k - 1) << (self.log_base2k - log_base2k_rem);
200 base2k_half = ((mask >> 1) + 1) as i64;
201 }
202
203 self.data[size..]
204 .iter_mut()
205 .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
206 }
207
208 pub fn add_dist_f64<T: Distribution<f64>>(&mut self, source: &mut Source, dist: T, bound: f64) {
209 let log_base2k_rem: usize = self.prec % self.log_base2k;
210
211 if log_base2k_rem != 0 {
212 self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
213 let mut dist_f64: f64 = dist.sample(source);
214 while dist_f64.abs() > bound {
215 dist_f64 = dist.sample(source)
216 }
217 *a += (dist_f64.round() as i64) << log_base2k_rem
218 });
219 } else {
220 self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
221 let mut dist_f64: f64 = dist.sample(source);
222 while dist_f64.abs() > bound {
223 dist_f64 = dist.sample(source)
224 }
225 *a += dist_f64.round() as i64
226 });
227 }
228 }
229
230 pub fn add_normal(&mut self, source: &mut Source, sigma: f64, bound: f64) {
231 self.add_dist_f64(source, Normal::new(0.0, sigma).unwrap(), bound);
232 }
233
234 pub fn trunc_pow2(&mut self, k: usize) {
235 if k == 0 {
236 return;
237 }
238
239 assert!(
240 k <= self.prec,
241 "invalid argument k: k={} > self.prec()={}",
242 k,
243 self.prec()
244 );
245
246 self.prec -= k;
247 self.data
248 .truncate((self.limbs() - k / self.log_base2k) * self.n());
249
250 let k_rem: usize = k % self.log_base2k;
251
252 if k_rem != 0 {
253 let mask: i64 = ((1 << (self.log_base2k - k_rem - 1)) - 1) << k_rem;
254 self.at_mut(self.limbs() - 1)
255 .iter_mut()
256 .for_each(|x: &mut i64| *x &= mask)
257 }
258 }
259}
260
261#[cfg(test)]
262mod tests {
263 use crate::poly::Poly;
264 use itertools::izip;
265 use sampling::source::Source;
266
267 #[test]
268 fn test_set_get_i64_lo_norm() {
269 let n: usize = 32;
270 let k: usize = 19;
271 let prec: usize = 128;
272 let mut a: Poly = Poly::new(n, k, prec);
273 let mut have: Vec<i64> = vec![i64::default(); n];
274 have.iter_mut()
275 .enumerate()
276 .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
277 a.set_i64(&have, 10);
278 let mut want = vec![i64::default(); n];
279 a.get_i64(&mut want);
280 izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
281 }
282
283 #[test]
284 fn test_set_get_i64_hi_norm() {
285 let n: usize = 8;
286 let k: usize = 17;
287 let prec: usize = 84;
288 let mut a: Poly = Poly::new(n, k, prec);
289 let mut have: Vec<i64> = vec![i64::default(); n];
290 let mut source = Source::new([1; 32]);
291 have.iter_mut().for_each(|x| {
292 *x = source
293 .next_u64n(u64::MAX, u64::MAX)
294 .wrapping_sub(u64::MAX / 2 + 1) as i64;
295 });
296 a.set_i64(&have, 63);
297 //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
298 let mut want = vec![i64::default(); n];
299 //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
300 a.get_i64(&mut want);
301 izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
302 }
303 #[test]
304 fn test_normalize() {
305 let n: usize = 8;
306 let k: usize = 17;
307 let prec: usize = 84;
308 let mut a: Poly = Poly::new(n, k, prec);
309 let mut have: Vec<i64> = vec![i64::default(); n];
310 let mut source = Source::new([1; 32]);
311 have.iter_mut().for_each(|x| {
312 *x = source
313 .next_u64n(u64::MAX, u64::MAX)
314 .wrapping_sub(u64::MAX / 2 + 1) as i64;
315 });
316 a.set_i64(&have, 63);
317 let mut carry: Vec<u8> = vec![u8::default(); n * 8];
318 a.normalize(&mut carry);
319
320 let base_half = 1 << (k - 1);
321 a.data
322 .iter()
323 .for_each(|x| assert!(x.abs() <= base_half, "|x|={} > 2^(k-1)={}", x, base_half));
324 let mut want = vec![i64::default(); n];
325 a.get_i64(&mut want);
326 izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
327 }
328}