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