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