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 // compute r[j] = (U + V) % q:
56 r[j] = U + V;
57 if r[j] >= q {
58 r[j] -= q;
59 }
60 // compute r[j + t] = (U - V) % q:
61 if U >= V {
62 r[j + t] = U - V;
63 } else {
64 r[j + t] = (q + U) - V;
65 }
66 }
67 k = k + 2 * t;
68 }
69 t /= 2;
70 m *= 2;
71 }
72 r
73 }
74
75 /// implements the Cooley-Tukey (CT) algorithm. Details at
76 /// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.2 of
77 /// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
78 pub fn intt(q: u64, n: usize, a: &Vec<u64>) -> Vec<u64> {
79 debug_assert_eq!(n, a.len());
80
81 let (_, roots_of_unity_inv, n_inv) = roots(q, n);
82
83 let mut t = 1;
84 let mut m = n / 2;
85 let mut r: Vec<u64> = a.clone();
86 while m > 0 {
87 let mut k = 0;
88 for i in 0..m {
89 let S: u64 = roots_of_unity_inv[m + i];
90 for j in k..k + t {
91 let U: u64 = r[j];
92 let V: u64 = r[j + t];
93 // compute r[j] = (U + V) % q:
94 r[j] = U + V;
95 if r[j] >= q {
96 r[j] -= q;
97 }
98 // compute r[j + t] = ((U - V) * S) % q;
99 if U >= V {
100 r[j + t] = ((U - V) * S) % q;
101 } else {
102 r[j + t] = ((q + U - V) * S) % q;
103 }
104 }
105 k += 2 * t;
106 }
107 t *= 2;
108 m /= 2;
109 }
110 for i in 0..n {
111 r[i] = (r[i] * n_inv) % q;
112 }
113 r
114 }
115}
116
117/// computes a primitive N-th root of unity using the method described by Thomas
118/// Pornin in https://crypto.stackexchange.com/a/63616
119const fn primitive_root_of_unity(q: u64, n: usize) -> u64 {
120 assert!(n.is_power_of_two());
121 assert!((q - 1) % n as u64 == 0);
122 let n_u64 = n as u64;
123
124 let mut k = 1;
125 while k < q {
126 // alternatively could get a random k at each iteration, if so, add the following if:
127 // `if k == 0 { continue; }`
128 let w = const_exp_mod(q, k, (q - 1) / n_u64);
129 if const_exp_mod(q, w, n_u64 / 2) != 1 {
130 return w; // w is a primitive N-th root of unity
131 }
132 k += 1;
133 }
134 panic!("No primitive root of unity");
135}
136
137fn roots_of_unity(q: u64, n: usize, w: u64) -> Vec<u64> {
138 let mut r: Vec<u64> = vec![0; n];
139 let mut i = 0;
140 let log_n = n.ilog2();
141 while i < n {
142 // (return the roots in bit-reverset order)
143 let j = ((i as u64).reverse_bits() >> (64 - log_n)) as usize;
144 r[i] = const_exp_mod(q, w, j as u64);
145 i += 1;
146 }
147 r
148}
149
150fn roots_of_unity_inv(q: u64, n: usize, v: Vec<u64>) -> Vec<u64> {
151 // assumes that the inputted roots are already in bit-reverset order
152 let mut r: Vec<u64> = vec![0; n];
153 let mut i = 0;
154 while i < n {
155 r[i] = const_inv_mod(q, v[i]);
156 i += 1;
157 }
158 r
159}
160
161/// returns x^k mod Q
162const fn const_exp_mod(q: u64, x: u64, k: u64) -> u64 {
163 // work on u128 to avoid overflow
164 let mut r = 1u128;
165 let mut x = x as u128;
166 let mut k = k as u128;
167 x = x % q as u128;
168 // exponentiation by square strategy
169 while k > 0 {
170 if k % 2 == 1 {
171 r = (r * x) % q as u128;
172 }
173 x = (x * x) % q as u128;
174 k /= 2;
175 }
176 r as u64
177}
178
179/// returns x^-1 mod Q
180const fn const_inv_mod(q: u64, x: u64) -> u64 {
181 // by Fermat's Little Theorem, x^-1 mod q \equiv x^{q-2} mod q
182 const_exp_mod(q, x, q - 2)
183}
184
185#[cfg(test)]
186mod tests {
187 use super::*;
188 use rand_distr::Distribution;
189
190 use anyhow::Result;
191
192 #[test]
193 fn test_ntt() -> Result<()> {
194 let q: u64 = 2u64.pow(16) + 1;
195 let n: usize = 4;
196
197 let a: Vec<u64> = vec![1u64, 2, 3, 4];
198
199 let a_ntt = NTT::ntt(q, n, &a);
200
201 let a_intt = NTT::intt(q, n, &a_ntt);
202
203 dbg!(&a);
204 dbg!(&a_ntt);
205 dbg!(&a_intt);
206 // dbg!(NTT::ROOT_OF_UNITY);
207 // dbg!(NTT::ROOTS_OF_UNITY);
208
209 assert_eq!(a, a_intt);
210 Ok(())
211 }
212
213 #[test]
214 fn test_ntt_loop() -> Result<()> {
215 let q: u64 = 2u64.pow(16) + 1;
216 let n: usize = 512;
217
218 use rand::distributions::Uniform;
219 let mut rng = rand::thread_rng();
220 let dist = Uniform::new(0_u64, q as u64);
221
222 for _ in 0..1000 {
223 let a: Vec<u64> = std::iter::repeat_with(|| dist.sample(&mut rng))
224 .take(n)
225 .collect();
226 let a_ntt = NTT::ntt(q, n, &a);
227 let a_intt = NTT::intt(q, n, &a_ntt);
228 assert_eq!(a, a_intt);
229 }
230 Ok(())
231 }
232}