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