1# Primitive Root of Unity
 2def get_primitive_root_of_unity(F, n):
 3    # using the method described by Thomas Pornin in
 4    # https://crypto.stackexchange.com/a/63616
 5    q = F.order()
 6    for k in range(q):
 7        if k==0:
 8            continue
 9        g = F(k)
10        #  g = F.random_element()
11        if g==0:
12            continue
13        w = g ^ ((q-1)/n)
14        if w^(n/2) != 1:
15            return g, w
16
17# Roots of Unity
18def get_nth_roots_of_unity(n, primitive_w):
19    w = [0]*n
20    for i in range(n):
21        w[i] = primitive_w^i
22    return w
23
24
25# fft (Fast Fourier Transform) returns:
26# - nth roots of unity
27# - Vandermonde matrix for the nth roots of unity
28# - Inverse Vandermonde matrix
29def fft(F, n):
30    g, primitive_w = get_primitive_root_of_unity(F, n)
31
32    w = get_nth_roots_of_unity(n, primitive_w)
33
34    ft = matrix(F, n)
35    for j in range(n):
36        row = []
37        for k in range(n):
38            row.append(primitive_w^(j*k))
39        ft.set_row(j, row)
40    ft_inv = ft^-1
41    return w, ft, ft_inv
42
43
44# Fast polynomial multiplicaton using FFT
45def poly_mul(fa, fb, F, n):
46    w, ft, ft_inv = fft(F, n)
47
48    # compute evaluation points from polynomials fa & fb at the roots of unity
49    a_evals = []
50    b_evals = []
51    for i in range(n):
52        a_evals.append(fa(w[i]))
53        b_evals.append(fb(w[i]))
54
55    # multiply elements in a_evals by b_evals
56    c_evals = map(operator.mul, a_evals, b_evals)
57    c_evals = vector(c_evals)
58
59    # using FFT, convert the c_evals into fc(x)
60    fc_coef = c_evals*ft_inv
61    fc2=P(fc_coef.list())
62    return fc2, c_evals