ntt.mx raw
1 package ring
2
3 // Number-Theoretic Transform for R_q = Z_q[x]/(x^n + 1).
4 // Cooley-Tukey forward, Gentleman-Sande inverse, precomputed twiddle factors.
5
6 type nttTables struct {
7 n int32
8 q uint32
9 logN int32
10 psiPow []uint32
11 psiInvPow []uint32
12 omegaPow []uint32
13 omegaInvPow []uint32
14 bitrevPerm []int32
15 invN uint32
16 }
17
18 var tableCache map[[2]uint32]*nttTables
19
20 func main() {
21 tableCache = map[[2]uint32]*nttTables{}
22 }
23
24 func getTables(p Params) (t *nttTables) {
25 key := [2]uint32{uint32(p.N), p.Q}
26 t, ok := tableCache[key]
27 if ok {
28 return t
29 }
30 t = newNTTTables(p)
31 tableCache[key] = t
32 return t
33 }
34
35 func newNTTTables(p Params) (t *nttTables) {
36 n := p.N
37 q := p.Q
38 psi := p.RootOfUnity
39 logN := log2(n)
40
41 psiPow := []uint32{:2 * n}
42 psiPow[0] = 1
43 for i := int32(1); i < 2*n; i++ {
44 psiPow[i] = mulMod(psiPow[i-1], psi, q)
45 }
46
47 psiInv := powMod(psi, q-2, q)
48 psiInvPow := []uint32{:2 * n}
49 psiInvPow[0] = 1
50 for i := int32(1); i < 2*n; i++ {
51 psiInvPow[i] = mulMod(psiInvPow[i-1], psiInv, q)
52 }
53
54 omega := mulMod(psi, psi, q)
55 omegaPow := []uint32{:n}
56 omegaPow[0] = 1
57 for i := int32(1); i < n; i++ {
58 omegaPow[i] = mulMod(omegaPow[i-1], omega, q)
59 }
60
61 omegaInv := powMod(omega, q-2, q)
62 omegaInvPow := []uint32{:n}
63 omegaInvPow[0] = 1
64 for i := int32(1); i < n; i++ {
65 omegaInvPow[i] = mulMod(omegaInvPow[i-1], omegaInv, q)
66 }
67
68 bitrevPerm := []int32{:n}
69 for i := int32(0); i < n; i++ {
70 bitrevPerm[i] = bitrev(i, logN)
71 }
72
73 return &nttTables{
74 n: n,
75 q: q,
76 logN: logN,
77 psiPow: psiPow,
78 psiInvPow: psiInvPow,
79 omegaPow: omegaPow,
80 omegaInvPow: omegaInvPow,
81 bitrevPerm: bitrevPerm,
82 invN: powMod(uint32(n), q-2, q),
83 }
84 }
85
86 func NTT(a *Poly) {
87 if a.isNTT {
88 return
89 }
90 t := getTables(a.params)
91 n := t.n
92 q := t.q
93 c := a.Coeffs
94
95 for i := int32(0); i < n; i++ {
96 c[i] = mulMod(c[i], t.psiPow[i], q)
97 }
98
99 for i := int32(0); i < n; i++ {
100 j := t.bitrevPerm[i]
101 if i < j {
102 c[i], c[j] = c[j], c[i]
103 }
104 }
105
106 for length := int32(1); length < n; length <<= 1 {
107 step := n / (2 * length)
108 for start := int32(0); start < n; start += 2 * length {
109 for j := int32(0); j < length; j++ {
110 tw := t.omegaPow[(j*step)%n]
111 idx0 := start + j
112 idx1 := idx0 + length
113 u := c[idx0]
114 v := mulMod(c[idx1], tw, q)
115 c[idx0] = addMod(u, v, q)
116 c[idx1] = subMod(u, v, q)
117 }
118 }
119 }
120
121 a.isNTT = true
122 }
123
124 func INTT(a *Poly) {
125 if !a.isNTT {
126 return
127 }
128 t := getTables(a.params)
129 n := t.n
130 q := t.q
131 c := a.Coeffs
132
133 for i := int32(0); i < n; i++ {
134 j := t.bitrevPerm[i]
135 if i < j {
136 c[i], c[j] = c[j], c[i]
137 }
138 }
139
140 for length := int32(1); length < n; length <<= 1 {
141 step := n / (2 * length)
142 for start := int32(0); start < n; start += 2 * length {
143 for j := int32(0); j < length; j++ {
144 tw := t.omegaInvPow[(j*step)%n]
145 idx0 := start + j
146 idx1 := idx0 + length
147 u := c[idx0]
148 v := mulMod(c[idx1], tw, q)
149 c[idx0] = addMod(u, v, q)
150 c[idx1] = subMod(u, v, q)
151 }
152 }
153 }
154
155 for i := int32(0); i < n; i++ {
156 c[i] = mulMod(c[i], t.invN, q)
157 c[i] = mulMod(c[i], t.psiInvPow[i], q)
158 }
159
160 a.isNTT = false
161 }
162
163 func Mul(a, b *Poly) (c *Poly) {
164 if a.isNTT && b.isNTT {
165 return MulPointwise(a, b)
166 }
167
168 aNTT := a.Clone()
169 bNTT := b.Clone()
170 NTT(aNTT)
171 NTT(bNTT)
172 c = MulPointwise(aNTT, bNTT)
173 INTT(c)
174 return c
175 }
176
177 func powMod(base, exp, q uint32) (result uint32) {
178 result = 1
179 b := base % q
180 for e := exp; e > 0; e >>= 1 {
181 if e&1 == 1 {
182 result = mulMod(result, b, q)
183 }
184 b = mulMod(b, b, q)
185 }
186 return result
187 }
188
189 func log2(n int32) (r int32) {
190 n >>= 1
191 for n > 0 {
192 r++
193 n >>= 1
194 }
195 return r
196 }
197
198 func bitrev(x, bits int32) (r int32) {
199 for i := int32(0); i < bits; i++ {
200 r = (r << 1) | (x & 1)
201 x >>= 1
202 }
203 return r
204 }
205