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.
9use crate::{ring::RingParam, ring_nq::Rq, zq::Zq};
10
11use std::collections::HashMap;
12
13#[derive(Debug)]
14pub struct NTT {}
15
16use std::sync::{Mutex, OnceLock};
17
18static CACHE: OnceLock<Mutex<HashMap<(u64, usize), (Vec<Zq>, Vec<Zq>, Zq)>>> = OnceLock::new();
19
20fn roots(q: u64, n: usize) -> (Vec<Zq>, Vec<Zq>, Zq) {
21 // Initialize CACHE with an empty HashMap on first use
22 let cache_lock = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
23
24 // Lock the HashMap for this thread
25 let mut cache = cache_lock.lock().unwrap();
26
27 if let Some(value) = cache.get(&(q, n)) {
28 // Found an existing value — return a clone
29 return value.clone();
30 }
31
32 // Not found — compute the new triple
33 let n_inv: Zq = Zq {
34 q,
35 v: const_inv_mod(q, n as u64),
36 };
37 let root_of_unity: u64 = primitive_root_of_unity(q, 2 * n);
38 let roots_of_unity: Vec<Zq> = roots_of_unity(q, n, root_of_unity);
39 let roots_of_unity_inv: Vec<Zq> = roots_of_unity_inv(q, n, roots_of_unity.clone());
40
41 let value = (roots_of_unity, roots_of_unity_inv, n_inv);
42
43 // Store and return
44 cache.insert((q, n), value.clone());
45 value
46}
47
48impl NTT {
49 /// implements the Cooley-Tukey (CT) algorithm. Details at
50 /// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.1 of
51 /// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
52 pub fn ntt(a: &Rq) -> Rq {
53 let (q, n) = (a.param.q, a.param.n);
54 let (roots_of_unity, _, _) = roots(q, n);
55
56 let mut t = n / 2;
57 let mut m = 1;
58 let mut r: Vec<Zq> = a.coeffs.clone();
59 while m < n {
60 let mut k = 0;
61 for i in 0..m {
62 let S: Zq = roots_of_unity[m + i];
63 for j in k..k + t {
64 let U: Zq = r[j];
65 let V: Zq = r[j + t] * S;
66 r[j] = U + V;
67 r[j + t] = U - V;
68 }
69 k = k + 2 * t;
70 }
71 t /= 2;
72 m *= 2;
73 }
74 // Rq::from_vec((a.q, n), r)
75 Rq {
76 param: RingParam { q, n },
77 coeffs: r,
78 evals: None,
79 }
80 }
81
82 /// implements the Cooley-Tukey (CT) algorithm. Details at
83 /// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.2 of
84 /// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
85 pub fn intt(a: &Rq) -> Rq {
86 let (q, n) = (a.param.q, a.param.n);
87 let (_, roots_of_unity_inv, n_inv) = roots(q, n);
88
89 let mut t = 1;
90 let mut m = n / 2;
91 let mut r: Vec<Zq> = a.coeffs.clone();
92 while m > 0 {
93 let mut k = 0;
94 for i in 0..m {
95 let S: Zq = roots_of_unity_inv[m + i];
96 for j in k..k + t {
97 let U: Zq = r[j];
98 let V: Zq = r[j + t];
99 r[j] = U + V;
100 r[j + t] = (U - V) * S;
101 }
102 k += 2 * t;
103 }
104 t *= 2;
105 m /= 2;
106 }
107 for i in 0..n {
108 r[i] = r[i] * n_inv;
109 }
110 // Rq::from_vec((a.q, n), r)
111 Rq {
112 param: RingParam { q, n },
113 coeffs: r,
114 evals: None,
115 }
116 }
117}
118
119/// computes a primitive N-th root of unity using the method described by Thomas
120/// Pornin in https://crypto.stackexchange.com/a/63616
121const fn primitive_root_of_unity(q: u64, n: usize) -> u64 {
122 assert!(n.is_power_of_two());
123 assert!((q - 1) % n as u64 == 0);
124 let n_u64 = n as u64;
125
126 let mut k = 1;
127 while k < q {
128 // alternatively could get a random k at each iteration, if so, add the following if:
129 // `if k == 0 { continue; }`
130 let w = const_exp_mod(q, k, (q - 1) / n_u64);
131 if const_exp_mod(q, w, n_u64 / 2) != 1 {
132 return w; // w is a primitive N-th root of unity
133 }
134 k += 1;
135 }
136 panic!("No primitive root of unity");
137}
138
139fn roots_of_unity(q: u64, n: usize, w: u64) -> Vec<Zq> {
140 let mut r: Vec<Zq> = vec![Zq { q, v: 0 }; n];
141 let mut i = 0;
142 let log_n = n.ilog2();
143 while i < n {
144 // (return the roots in bit-reverset order)
145 let j = ((i as u64).reverse_bits() >> (64 - log_n)) as usize;
146 r[i] = Zq {
147 q,
148 v: const_exp_mod(q, w, j as u64),
149 };
150 i += 1;
151 }
152 r
153}
154
155fn roots_of_unity_inv(q: u64, n: usize, v: Vec<Zq>) -> Vec<Zq> {
156 // assumes that the inputted roots are already in bit-reverset order
157 let mut r: Vec<Zq> = vec![Zq { q, v: 0 }; n];
158 let mut i = 0;
159 while i < n {
160 r[i] = Zq {
161 q,
162 v: const_inv_mod(q, v[i].v),
163 };
164 i += 1;
165 }
166 r
167}
168
169/// returns x^k mod Q
170const fn const_exp_mod(q: u64, x: u64, k: u64) -> u64 {
171 // work on u128 to avoid overflow
172 let mut r = 1u128;
173 let mut x = x as u128;
174 let mut k = k as u128;
175 x = x % q as u128;
176 // exponentiation by square strategy
177 while k > 0 {
178 if k % 2 == 1 {
179 r = (r * x) % q as u128;
180 }
181 x = (x * x) % q as u128;
182 k /= 2;
183 }
184 r as u64
185}
186
187/// returns x^-1 mod Q
188const fn const_inv_mod(q: u64, x: u64) -> u64 {
189 // by Fermat's Little Theorem, x^-1 mod q \equiv x^{q-2} mod q
190 const_exp_mod(q, x, q - 2)
191}
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196 use crate::Ring;
197
198 use anyhow::Result;
199
200 #[test]
201 fn test_ntt() -> Result<()> {
202 let q: u64 = 2u64.pow(16) + 1;
203 let n: usize = 4;
204 let param = RingParam { q, n };
205
206 let a: Vec<u64> = vec![1u64, 2, 3, 4];
207 let a: Rq = Rq::from_vec_u64(¶m, a);
208
209 let a_ntt = NTT::ntt(&a);
210
211 let a_intt = NTT::intt(&a_ntt);
212
213 dbg!(&a);
214 dbg!(&a_ntt);
215 dbg!(&a_intt);
216 // dbg!(NTT::ROOT_OF_UNITY);
217 // dbg!(NTT::ROOTS_OF_UNITY);
218
219 assert_eq!(a, a_intt);
220 Ok(())
221 }
222
223 #[test]
224 fn test_ntt_loop() -> Result<()> {
225 let q: u64 = 2u64.pow(16) + 1;
226 let n: usize = 512;
227 let param = RingParam { q, n };
228
229 use rand::distributions::Uniform;
230 let mut rng = rand::thread_rng();
231 let dist = Uniform::new(0_f64, q as f64);
232
233 for _ in 0..1000 {
234 let a: Rq = Rq::rand(&mut rng, dist, ¶m);
235 let a_ntt = NTT::ntt(&a);
236 let a_intt = NTT::intt(&a_ntt);
237 assert_eq!(a, a_intt);
238 }
239 Ok(())
240 }
241
242 // #[test]
243 // fn test_ntt_loop_2() -> Result<()> {
244 // // let q: u64 = 2u64.pow(16) + 1;
245 // // let n: usize = 512;
246 // let q: u64 = 35184371138561;
247 // let n: usize = 1 << 14;
248 // let param = RingParam { q, n };
249 //
250 // use rand::distributions::Uniform;
251 // let mut rng = rand::thread_rng();
252 // let dist = Uniform::new(0_f64, q as f64);
253 //
254 // let a: Rq = Rq::rand(&mut rng, dist, ¶m);
255 // let start = std::time::Instant::now();
256 // for _ in 0..10_000 {
257 // let a_ntt = NTT::ntt(&a);
258 // let a_intt = NTT::intt(&a_ntt);
259 // assert_eq!(a, a_intt);
260 // }
261 // dbg!(start.elapsed());
262 // Ok(())
263 // }
264}