Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
112 lines (101 sloc) 4.39 KB
from fft import fft, mul_polys
# Calculates modular inverses [1/values[0], 1/values[1] ...]
def multi_inv(values, modulus):
partials = [1]
for i in range(len(values)):
partials.append(partials[-1] * values[i] % modulus)
inv = pow(partials[-1], modulus - 2, modulus)
outputs = [0] * len(values)
for i in range(len(values), 0, -1):
outputs[i-1] = partials[i-1] * inv % modulus
inv = inv * values[i-1] % modulus
return outputs
# Generates q(x) = poly(k * x)
def p_of_kx(poly, modulus, k):
o = []
power_of_k = 1
for x in poly:
o.append(x * power_of_k % modulus)
power_of_k = (power_of_k * k) % modulus
return o
# Return (x - root**positions[0]) * (x - root**positions[1]) * ...
# possibly with a constant factor offset
def _zpoly(positions, modulus, roots_of_unity):
# If there are not more than 4 positions, use the naive
# O(n^2) algorithm as it is faster
if len(positions) <= 4:
root = [1]
for pos in positions:
x = roots_of_unity[pos]
root.insert(0, 0)
for j in range(len(root)-1):
root[j] -= root[j+1] * x
return [x % modulus for x in root]
else:
# Recursively find the zpoly for even indices and odd
# indices, operating over a half-size subgroup in each
# case
left = _zpoly([x//2 for x in positions if x%2 == 0],
modulus, roots_of_unity[::2])
right = _zpoly([x//2 for x in positions if x%2 == 1],
modulus, roots_of_unity[::2])
invroot = roots_of_unity[-1]
# Offset the result for the odd indices, and combine
# the two
o = mul_polys(left, p_of_kx(right, modulus, invroot),
modulus, roots_of_unity[1])
# Deal with the special case where mul_polys returns zero
# when it should return x ^ (2 ** k) - 1
if o == [0] * len(o):
return [1] + [0] * (len(o) - 1) + [modulus - 1]
else:
return o
def zpoly(positions, modulus, root_of_unity):
# Precompute roots of unity
rootz = [1, root_of_unity]
while rootz[-1] != 1:
rootz.append((rootz[-1] * root_of_unity) % modulus)
return _zpoly(positions, modulus, rootz[:-1])
def erasure_code_recover(vals, modulus, root_of_unity):
# Generate the polynomial that is zero at the roots of unity
# corresponding to the indices where vals[i] is None
import poly_utils
z = zpoly([i for i in range(len(vals)) if vals[i] is None],
modulus, root_of_unity)
zvals = fft(z, modulus, root_of_unity)
# Pointwise-multiply (vals filling in zero at missing spots) * z
# By construction, this equals vals * z
vals_with_zeroes = [x or 0 for x in vals]
p_times_z_vals = [x*y % modulus for x,y in zip(vals_with_zeroes, zvals)]
p_times_z = fft(p_times_z_vals, modulus, root_of_unity, inv=True)
# Keep choosing k values until the algorithm does not fail
# Check only with primitive roots of unity
for k in range(2, modulus):
if pow(k, (modulus - 1) // 2, modulus) == 1:
continue
invk = pow(k, modulus - 2, modulus)
# Convert p_times_z(x) and z(x) into new polynomials
# q1(x) = p_times_z(k*x) and q2(x) = z(k*x)
# These are likely to not be 0 at any of the evaluation points.
p_times_z_of_kx = [x * pow(k, i, modulus) % modulus
for i, x in enumerate(p_times_z)]
p_times_z_of_kx_vals = fft(p_times_z_of_kx, modulus, root_of_unity)
z_of_kx = [x * pow(k, i, modulus) for i, x in enumerate(z)]
z_of_kx_vals = fft(z_of_kx, modulus, root_of_unity)
# Compute q1(x) / q2(x) = p(k*x)
inv_z_of_kv_vals = multi_inv(z_of_kx_vals, modulus)
p_of_kx_vals = [x*y % modulus for x,y in
zip(p_times_z_of_kx_vals, inv_z_of_kv_vals)]
p_of_kx = fft(p_of_kx_vals, modulus, root_of_unity, inv=True)
# Given q3(x) = p(k*x), recover p(x)
p_of_x = [x * pow(invk, i, modulus) % modulus
for i, x in enumerate(p_of_kx)]
output = fft(p_of_x, modulus, root_of_unity)
# Check that the output matches the input
success = True
for inpd, outd in zip(vals, output):
success *= (inpd is None or inpd == outd)
if not success:
continue
# Output the evaluations if all good
return output
You can’t perform that action at this time.