fhe-study/arith/src/ntt.rs
arith/src/ntt.rsRaw
1//! Implementation of the NTT & iNTT, following the CT & GS algorighms, more details in
2//! https://eprint.iacr.org/2017/727.pdf, some notes at
3//! https://github.com/arnaucube/math/blob/master/notes_ntt.pdf .
4//!
5//! NOTE: initially I implemented it with fixed Q & N, given as constant
6//! generics; but once using real-world parameters, the stack could not handle
7//! it, so moved to use Vec instead of fixed-sized arrays, and adapted the NTT
8//! implementation to that too.
9
10use std::collections::HashMap;
11
12#[derive(Debug)]
13pub struct NTT {}
14
15use std::sync::{Mutex, OnceLock};
16
17static CACHE: OnceLock<Mutex<HashMap<(u64, usize), (Vec<u64>, Vec<u64>, u64)>>> = OnceLock::new();
18
19fn roots(q: u64, n: usize) -> (Vec<u64>, Vec<u64>, u64) {
20 let cache_lock = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
21 let mut cache = cache_lock.lock().unwrap();
22 if let Some(value) = cache.get(&(q, n)) {
23 return value.clone();
24 }
25
26 let n_inv: u64 = const_inv_mod(q, n as u64);
27 let root_of_unity: u64 = primitive_root_of_unity(q, 2 * n);
28 let roots_of_unity: Vec<u64> = roots_of_unity(q, n, root_of_unity);
29 let roots_of_unity_inv: Vec<u64> = roots_of_unity_inv(q, n, roots_of_unity.clone());
30 let value = (roots_of_unity, roots_of_unity_inv, n_inv);
31
32 cache.insert((q, n), value.clone());
33 value
34}
35
36impl NTT {
37 /// implements the Cooley-Tukey (CT) algorithm. Details at
38 /// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.1 of
39 /// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
40 pub fn ntt(q: u64, n: usize, a: &Vec<u64>) -> Vec<u64> {
41 debug_assert_eq!(n, a.len());
42
43 let (roots_of_unity, _, _) = roots(q, n);
44
45 let mut t = n / 2;
46 let mut m = 1;
47 let mut r: Vec<u64> = a.clone();
48 while m < n {
49 let mut k = 0;
50 for i in 0..m {
51 let S: u64 = roots_of_unity[m + i];
52 for j in k..k + t {
53 let U: u64 = r[j];
54 // let V: u64 = (r[j + t] * S) % q;
55 let V: u64 = ((r[j + t] as u128 * S as u128) % q as u128) as u64;
56 // compute r[j] = (U + V) % q:
57 r[j] = ((U as u128 + V as u128) % q as u128) as u64;
58 r[j + t] = (((q as u128 + U as u128) - V as u128) % q as u128) as u64;
59 // r[j] = U + V;
60 // if r[j] >= q {
61 // r[j] -= q;
62 // }
63 // // compute r[j + t] = (U - V) % q:
64 // if U >= V {
65 // r[j + t] = U - V;
66 // } else {
67 // r[j + t] = (q + U) - V;
68 // }
69 }
70 k = k + 2 * t;
71 }
72 t /= 2;
73 m *= 2;
74 }
75 r
76 }
77
78 /// implements the Cooley-Tukey (CT) algorithm. Details at
79 /// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.2 of
80 /// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
81 pub fn intt(q: u64, n: usize, a: &Vec<u64>) -> Vec<u64> {
82 debug_assert_eq!(n, a.len());
83
84 let (_, roots_of_unity_inv, n_inv) = roots(q, n);
85
86 let mut t = 1;
87 let mut m = n / 2;
88 let mut r: Vec<u64> = a.clone();
89 while m > 0 {
90 let mut k = 0;
91 for i in 0..m {
92 let S: u64 = roots_of_unity_inv[m + i];
93 for j in k..k + t {
94 let U: u64 = r[j];
95 let V: u64 = r[j + t];
96 r[j] = ((U as u128 + V as u128) % q as u128) as u64;
97 r[j + t] =
98 (((q as u128 + U as u128) - V as u128) * S as u128 % q as u128) as u64;
99 // compute r[j] = (U + V) % q:
100 // r[j] = U + V;
101 // if r[j] >= q {
102 // r[j] -= q;
103 // }
104 // // compute r[j + t] = ((U - V) * S) % q;
105 // if U >= V {
106 // r[j + t] = ((U - V) * S) % q;
107 // } else {
108 // r[j + t] = ((q + U - V) * S) % q;
109 // }
110 }
111 k += 2 * t;
112 }
113 t *= 2;
114 m /= 2;
115 }
116 for i in 0..n {
117 // r[i] = (r[i] * n_inv) % q;
118 r[i] = ((r[i] as u128 * n_inv as u128) % q as u128) as u64;
119 }
120 r
121 }
122}
123
124/// computes a primitive N-th root of unity using the method described by Thomas
125/// Pornin in https://crypto.stackexchange.com/a/63616
126const fn primitive_root_of_unity(q: u64, n: usize) -> u64 {
127 assert!(n.is_power_of_two());
128 assert!((q - 1) % n as u64 == 0);
129 let n_u64 = n as u64;
130
131 let mut k = 1;
132 while k < q {
133 // alternatively could get a random k at each iteration, if so, add the following if:
134 // `if k == 0 { continue; }`
135 let w = const_exp_mod(q, k, (q - 1) / n_u64);
136 if const_exp_mod(q, w, n_u64 / 2) != 1 {
137 return w; // w is a primitive N-th root of unity
138 }
139 k += 1;
140 }
141 panic!("No primitive root of unity");
142}
143
144fn roots_of_unity(q: u64, n: usize, w: u64) -> Vec<u64> {
145 let mut r: Vec<u64> = vec![0; n];
146 let mut i = 0;
147 let log_n = n.ilog2();
148 while i < n {
149 // (return the roots in bit-reverset order)
150 let j = ((i as u64).reverse_bits() >> (64 - log_n)) as usize;
151 r[i] = const_exp_mod(q, w, j as u64);
152 i += 1;
153 }
154 r
155}
156
157fn roots_of_unity_inv(q: u64, n: usize, v: Vec<u64>) -> Vec<u64> {
158 // assumes that the inputted roots are already in bit-reverset order
159 let mut r: Vec<u64> = vec![0; n];
160 let mut i = 0;
161 while i < n {
162 r[i] = const_inv_mod(q, v[i]);
163 i += 1;
164 }
165 r
166}
167
168/// returns x^k mod Q
169const fn const_exp_mod(q: u64, x: u64, k: u64) -> u64 {
170 // work on u128 to avoid overflow
171 let mut r = 1u128;
172 let mut x = x as u128;
173 let mut k = k as u128;
174 x = x % q as u128;
175 // exponentiation by square strategy
176 while k > 0 {
177 if k % 2 == 1 {
178 r = (r * x) % q as u128;
179 }
180 x = (x * x) % q as u128;
181 k /= 2;
182 }
183 r as u64
184}
185
186/// returns x^-1 mod Q
187const fn const_inv_mod(q: u64, x: u64) -> u64 {
188 // by Fermat's Little Theorem, x^-1 mod q \equiv x^{q-2} mod q
189 const_exp_mod(q, x, q - 2)
190}
191
192#[cfg(test)]
193mod tests {
194 use super::*;
195 use rand_distr::Distribution;
196
197 use anyhow::Result;
198
199 #[test]
200 fn test_ntt() -> Result<()> {
201 let q: u64 = 2u64.pow(16) + 1;
202 let n: usize = 4;
203
204 let a: Vec<u64> = vec![1u64, 2, 3, 4];
205
206 let a_ntt = NTT::ntt(q, n, &a);
207
208 let a_intt = NTT::intt(q, n, &a_ntt);
209
210 dbg!(&a);
211 dbg!(&a_ntt);
212 dbg!(&a_intt);
213 // dbg!(NTT::ROOT_OF_UNITY);
214 // dbg!(NTT::ROOTS_OF_UNITY);
215
216 assert_eq!(a, a_intt);
217 Ok(())
218 }
219
220 #[test]
221 fn test_ntt_loop() -> Result<()> {
222 let q: u64 = 2u64.pow(16) + 1;
223 let n: usize = 512;
224
225 use rand::distributions::Uniform;
226 let mut rng = rand::thread_rng();
227 let dist = Uniform::new(0_u64, q as u64);
228
229 for _ in 0..1000 {
230 let a: Vec<u64> = std::iter::repeat_with(|| dist.sample(&mut rng))
231 .take(n)
232 .collect();
233 let a_ntt = NTT::ntt(q, n, &a);
234 let a_intt = NTT::intt(q, n, &a_ntt);
235 assert_eq!(a, a_intt);
236 }
237 Ok(())
238 }
239}