fhe-study/arith/src/ntt_u64.rs
arith/src/ntt_u64.rsRaw
1//! This file implements the wrapper on top of the ntt.rs to be able to compute
2//! the NTT for non-prime modulus, specifically for modulus 2^64 (for u64).
3
4use crate::ntt::NTT as NTT_p;
5
6// const P0: u64 = 17293822569241362433;
7
8// const P0: u64 = 4611686018427387905;
9// const P1: u64 = 4611686018326724609;
10
11// const P0: u64 = 8070449433331580929; // max use 1<<62 -1
12// const P1: u64 = 8070450532384645121;
13
14// const P0: u64 = 0x80000000080001; // max use 1<<55 -1
15// const P1: u64 = 0x80000000130001;
16
17// const P0: u64 = ((1u128 << 64) - (1u128 << 28) + 1u128) as u64;
18// const P1: u64 = ((1u128 << 64) - (1u128 << 27) + 1u128) as u64;
19// const P2: u64 = ((1u128 << 64) - (1u128 << 26) + 1u128) as u64;
20
21const P0: u64 = ((1u128 << 32) - (1u128 << 18) + 1u128) as u64;
22const P1: u64 = ((1u128 << 32) - (1u128 << 17) + 1u128) as u64;
23const P2: u64 = ((1u128 << 32) - (1u128 << 16) + 1u128) as u64;
24// const P1: u64 = ((1u128 << 64) - (1u128 << 27) + 1u128) as u64;
25// const P2: u64 = (1 << 60) - (1 << 26) + 1;
26
27#[derive(Debug)]
28pub struct NTT {}
29
30impl NTT {
31 pub fn ntt(
32 n: usize,
33 a: &Vec<u64>,
34 ) -> (
35 Vec<u64>,
36 Vec<u64>,
37 Vec<u64>,
38 // Vec<u64>,
39 // Vec<u64>,
40 // Vec<u64>,
41 // Vec<u64>,
42 ) {
43 // TODO ensure that: a_i <P0
44
45 // apply modulus p_i
46 let a_0: Vec<u64> = a.iter().map(|a_i| a_i % P0).collect();
47 let a_1: Vec<u64> = a.iter().map(|a_i| a_i % P1).collect();
48 let a_2: Vec<u64> = a.iter().map(|a_i| a_i % P2).collect();
49 // let a_0: Vec<u64> = a.iter().map(|a_i| (a_i % P0 + P0) % P0).collect();
50 // let a_1: Vec<u64> = a.iter().map(|a_i| (a_i % P1 + P1) % P1).collect();
51 // let a_2: Vec<u64> = a.iter().map(|a_i| (a_i % P2 + P2) % P2).collect();
52 // let a_3: Vec<u64> = a.iter().map(|a_i| (a_i % P3 + P3) % P3).collect();
53 // let a_4: Vec<u64> = a.iter().map(|a_i| (a_i % P4 + P4) % P4).collect();
54 // let a_5: Vec<u64> = a.iter().map(|a_i| (a_i % P5 + P5) % P5).collect();
55 // let a_6: Vec<u64> = a.iter().map(|a_i| (a_i % P6 + P6) % P6).collect();
56
57 dbg!(&a_0, &a_1, &a_2);
58 let r_0 = NTT_p::ntt(P0, n, &a_0);
59 let r_1 = NTT_p::ntt(P1, n, &a_1);
60 let r_2 = NTT_p::ntt(P2, n, &a_2);
61 // let r_3 = NTT_p::ntt(P3, n, &a_3);
62 // let r_4 = NTT_p::ntt(P4, n, &a_4);
63 // let r_5 = NTT_p::ntt(P5, n, &a_5);
64 // let r_6 = NTT_p::ntt(P6, n, &a_6);
65
66 (r_0, r_1, r_2) //, r_3, r_4) //, r_5, r_6)
67 }
68
69 pub fn intt(
70 n: usize,
71 r: &(
72 Vec<u64>,
73 Vec<u64>,
74 Vec<u64>,
75 // Vec<u64>,
76 // Vec<u64>,
77 // Vec<u64>,
78 // Vec<u64>,
79 ),
80 ) -> Vec<u64> {
81 let a_0 = NTT_p::intt(P0, n, &r.0);
82 let a_1 = NTT_p::intt(P1, n, &r.1);
83 let a_2 = NTT_p::intt(P2, n, &r.2);
84 // let a_3 = NTT_p::intt(P3, n, &r.3);
85 // let a_4 = NTT_p::intt(P4, n, &r.4);
86 // let a_5 = NTT_p::intt(P5, n, &r.5);
87 // let a_6 = NTT_p::intt(P6, n, &r.6);
88
89 // Garner CRT for two moduli: combine (r1 mod p1, r2 mod p2) -> Z/(p1*p2)
90 // let inv_p1_mod_p2: u128 = inv_mod_u64(p1 % p2, p2) as u128;
91 // const INV_P1_MOD_P2: u128 = 4895217125691974194;
92
93 reconstruct(a_0, a_1, a_2) // , a_3, a_4) //, a_5, a_6)
94 }
95}
96
97fn reconstruct(
98 a0: Vec<u64>,
99 a1: Vec<u64>,
100 a2: Vec<u64>,
101 // a_3: Vec<u64>,
102 // a_4: Vec<u64>,
103 // a_5: Vec<u64>,
104 // a_6: Vec<u64>,
105) -> Vec<u64> {
106 // y_i = q/q_i
107 // let y0 = ((u64::MAX as u128 + 1) / P0 as u128);
108 // let y1 = ((u64::MAX as u128 + 1) / P1 as u128);
109 // let y2 = ((u64::MAX as u128 + 1) / P2 as u128) as u64;
110 // let y0: u128 = P1 as u128; // N_i =Q/P0 = P1*P2
111 // let y1: u128 = P0 as u128;
112 let y0: u128 = P1 as u128 * P2 as u128; // N_i =Q/P0 = P1*P2
113 let y1: u128 = P0 as u128 * P2 as u128;
114 let y2: u128 = P0 as u128 * P1 as u128;
115 // let y0 = (Q / P0 as u128) as u64;
116 // let y1 = (Q / P1 as u128) as u64;
117 // let y2 = ((u64::MAX as u128 + 1) / P2 as u128) as u64;
118 // let y3 = ((u64::MAX as u128 + 1) / P3 as u128) as u64;
119 // let y4 = ((u64::MAX as u128 + 1) / P4 as u128) as u64;
120
121 // y_i^-1 mod q_i = z_i
122 dbg!(P0, y0);
123 let z0: u128 = inv_mod(P0 as u128, y0); // M_i = N_i^-1 mod q_i
124 let z1: u128 = inv_mod(P1 as u128, y1);
125 let z2: u128 = inv_mod(P2 as u128, y2);
126 // let y2_inv = inv_mod(P2 as u128, y2);
127 // let y3_inv = inv_mod(P3 as u128, y3);
128 // let y4_inv = inv_mod(P4 as u128, y4);
129
130 // m1 = q1^-1 mod q2
131 // aux = (a2 - a1) * m1 mod q2
132 // a = a1 + (q1 * m1) * aux
133
134 // m1 == z1
135 // let m1 = inv_mod(P1 as u128, P0 as u128); // P0^-1 mod P1
136 // let aux: Vec<u64> = itertools::zip_eq(a0.clone(), a1.clone())
137 // .map(|(a0_i, a1_i)| ((a1_i - a0_i) * m1) % P1)
138 // .collect();
139 let p0: u128 = P0 as u128;
140 let p1: u128 = P1 as u128;
141 let p2: u128 = P2 as u128;
142 // let a: Vec<u64> = itertools::zip_eq(a0, a1)
143 // // .map(|(a1_i, aux_i)| a1_i + (P1 * m1) * aux_i)
144 // // .map(|(a0_i, aux_i)| a0_i + (P0 * m1) * aux_i)
145 // .map(|(a0_i, a1_i)| a0_i as u128 + ((p0 * z1) % p1) * (((a1_i - a0_i) as u128 * z1) % p1))
146 // .map(|v| v as u64)
147 // // let a: Vec<u64> = itertools::zip_eq(a0, a1)
148 // // .map(|(a0_i, a1_i)| {
149 // // ((((a0_i as u128 * z0) % P0 as u128) * y0) + (a1_i as u128 * z1 % P1 as u128) * y1)
150 // // as u64
151 // // })
152 // .collect();
153 // a
154
155 // dbg!(a0[0] as u128);
156 // dbg!(a0[0] as u128 * y0);
157 // dbg!(a0[0] as u128 * z0);
158 // dbg!(a0[0] as u128 * y0 * z0);
159 dbg!(&y0, &y1, &y2);
160 dbg!(&z0, &z1, &z2);
161 let Q = P0 as u128 * P1 as u128 * P2 as u128;
162 let a: Vec<u64> = itertools::multizip((a0, a1, a2))
163 .map(|(a0_i, a1_i, a2_i)| {
164 (a0_i as u128 * y0 * z0)// % Q
165 + (a1_i as u128 * y1 * z1)// % Q
166 + (a2_i as u128 * y2 * z2) // % Q
167 })
168 .map(|v| v % Q)
169 .map(|v| v as u32)
170 .map(|v| v as u64)
171 .collect();
172 a
173 // dbg!(&a);
174 // let Q = y2 * P2 as u128;
175 // let a: Vec<u128> = a.iter().map(|a_i| a_i % Q).collect();
176 // dbg!(&a);
177 // let q64 = 1_u128 << 64;
178 // let a: Vec<u64> = a.iter().map(|a_i| (a_i % q64) as u64).collect();
179 // a
180
181 /*
182 // x_i*z_i mod q_i
183 let r0: Vec<u64> = a_0.iter().map(|a_i| ((a_i * z0) % P0) * y0).collect();
184 let r1: Vec<u64> = a_1.iter().map(|a_i| ((a_i * z1) % P1) * y1).collect();
185 // let r0: Vec<u64> = a_0.iter().map(|a_i| ((a_i * z0) % P0) * y0).collect();
186 // let r1: Vec<u64> = a_1.iter().map(|a_i| ((a_i * z1) % P1) * y1).collect();
187 // let r2: Vec<u64> = a_2.iter().map(|a_i| ((a_i * y2_inv) % P2) * y2).collect();
188 // let r3: Vec<u64> = a_3.iter().map(|a_i| ((a_i * y3_inv) % P3) * y3).collect();
189 // let r4: Vec<u64> = a_4.iter().map(|a_i| ((a_i * y4_inv) % P4) * y4).collect();
190
191 let r: Vec<u64> = itertools::multizip((r0.iter(), r1.iter()))
192 .map(|(a, b)| a + b)
193 .collect();
194 // let r = r0;
195 //
196 dbg!(&r);
197
198 let p1p2: u128 = (P0 as u128) * (P1 as u128);
199 // let p1p2_inv: u128 = inv_mod((P0 % P1) as u128, P1) as u128;
200 let p1p2_inv: u128 = inv_mod((P0) as u128, P1) as u128;
201 dbg!(&p1p2);
202 dbg!(&p1p2_inv);
203 // let p1p2: u128 = P0 as u128 / 2; // PIHALF
204 let r = r
205 .iter()
206 .map(|c_i_u64| {
207 let c_i = *c_i_u64 as u128;
208 if c_i * 2 >= p1p2 {
209 // if c_i >= p1p2 {
210 c_i.wrapping_sub(p1p2) as u64
211 } else {
212 c_i as u64
213 }
214 })
215 .collect();
216 // let r: Vec<u64> = itertools::multizip((r0.iter(), r1.iter(), r2.iter(), r3.iter(), r4.iter()))
217 // .map(|(a, b, c, d, e)| a + b + c + d + e)
218 // .collect();
219 // let mut r = a_0 + y0_inv + a_1 * y1_inv + a_2 * y2_inv + a_3 * y3_inv + a_4 * y4_inv;
220
221 r
222 */
223}
224
225fn exp_mod(q: u128, x: u128, k: u128) -> u128 {
226 // work on u128 to avoid overflow
227 let mut r = 1u128;
228 let mut x = x.clone();
229 let mut k = k.clone();
230 x = x % q;
231 // exponentiation by square strategy
232 while k > 0 {
233 if k % 2 == 1 {
234 r = (r * x) % q;
235 }
236 x = (x * x) % q;
237 k /= 2;
238 }
239 r
240}
241/// returns x^-1 mod Q, assuming x and Q are coprime, generally Q is prime
242// fn inv_mod(q: u128, x: u128) -> u128 {
243// // by Fermat's Little Theorem, x^-1 mod q \equiv x^{q-2} mod q
244// // exp_mod(q, x, q - 2)
245// exp_mod(q, x, q - 2)
246// }
247
248fn inv_mod(m: u128, a: u128) -> u128 {
249 // if m == 1 {
250 // return Some(0);
251 // }
252
253 let mut m = m.clone();
254 let mut a = a.clone();
255 let m0 = m.clone();
256 let mut x0: i128 = 0;
257 let mut x1: i128 = 1;
258
259 while a > 1 {
260 let q = a / m;
261 let t = m.clone();
262
263 m = a % m;
264 a = t.clone();
265
266 let t = x0;
267 x0 = x1 - (q as i128) * x0;
268 x1 = t;
269 }
270
271 if x1 < 0 {
272 x1 += m0 as i128;
273 }
274
275 x1 as u128
276}
277// fn inv_mod(m: u128, a: u128) -> u128 {
278// let mut m = m.clone();
279// let mut a = a.clone();
280// let m0 = m.clone();
281// let mut x0 = 0;
282// let mut x1 = 1;
283//
284// if m == 1 {
285// return 0;
286// }
287//
288// while a > 1 {
289// let q = a / m;
290//
291// let mut t = m.clone();
292//
293// m = a % m;
294// a = t.clone();
295//
296// t = x0.clone();
297//
298// x0 = x1 - q * x0;
299//
300// x1 = t.clone();
301// }
302//
303// // if (x1 < 0) {
304// // x1 = x1 + m0
305// // }
306//
307// // return x1 % m0;
308// return x1 % m0;
309// }
310
311#[cfg(test)]
312mod tests {
313 use super::*;
314 use rand_distr::Distribution;
315
316 use anyhow::Result;
317
318 #[test]
319 fn test_inv_mod() -> Result<()> {
320 let x = 3;
321 let x_inv = inv_mod(P0 as u128, x);
322 dbg!(&P0);
323 dbg!(&x_inv);
324
325 // let r = x_inv * x;
326 // dbg!(&r);
327
328 Ok(())
329 }
330
331 #[test]
332 fn test_reconstruct() -> Result<()> {
333 let n: usize = 16;
334
335 use rand::Rng;
336 let mut rng = rand::thread_rng();
337 let a: Vec<u64> = (0..n)
338 // .map(|_| rng.gen_range(0..(1 << 64)))
339 .map(|_| rng.gen_range(0..(1 << 32)))
340 // .map(|_| rng.gen_range(0..16))
341 // .map(|_| rng.sample(rand::distributions::Standard))
342 .collect();
343
344 dbg!(a.len());
345
346 let a_0: Vec<u64> = a.iter().map(|a_i| a_i % P0).collect();
347 let a_1: Vec<u64> = a.iter().map(|a_i| a_i % P1).collect();
348 let a_2: Vec<u64> = a.iter().map(|a_i| a_i % P2).collect();
349 dbg!(&a_0, &a_1, &a_2);
350 let a_reconstructed = reconstruct(a_0, a_1, a_2);
351
352 dbg!(&a_reconstructed);
353 assert_eq!(a_reconstructed, a);
354
355 Ok(())
356 }
357
358 #[test]
359 fn test_dbg() -> Result<()> {
360 println!("{}", 1u128 << 64);
361 let n: usize = 2;
362
363 println!("{}", P0);
364 println!("{}", P1);
365 // let q = 1u128 << 64;
366 // assert!(P0 as u128 * P1 as u128 > (n as u128 * (q * q)) / 2);
367
368 // let a: Vec<u64> = vec![1u64, 2, 3, 4];
369 // let a: Vec<u64> = vec![1u64, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
370 // let a: Vec<u64> = vec![9u64, 8, 7, 6, 0, 9999, 0, 0, 0, 0, 0, 0, 6, 7, 8, 9];
371 use rand::Rng;
372 let mut rng = rand::thread_rng();
373 let a: Vec<u64> = (0..n)
374 // .map(|_| rng.gen_range(0..(1 << 64)))
375 // .map(|_| rng.gen_range(0..(1 << 32)))
376 .map(|_| rng.gen_range(0..16))
377 // .map(|_| rng.sample(rand::distributions::Standard))
378 .collect();
379
380 dbg!(a.len());
381
382 let a_0: Vec<u64> = a.iter().map(|a_i| a_i % P0).collect();
383 let a_1: Vec<u64> = a.iter().map(|a_i| a_i % P1).collect();
384 let a_2: Vec<u64> = a.iter().map(|a_i| a_i % P2).collect();
385 dbg!(&a_0, &a_1, &a_2);
386 // let a_0 = vec![3];
387 // let a_1 = vec![3];
388 // let a_2 = vec![3];
389 let a_intt = reconstruct(a_0, a_1, a_2);
390 // let a_ntt = NTT::ntt(n, &a);
391 // dbg!(&a_ntt);
392 //
393 // let a_intt = NTT::intt(n, &a_ntt);
394
395 dbg!(&a_intt);
396 assert_eq!(a_intt, a);
397
398 Ok(())
399 }
400}