fri_folds.sage

p = 2^8 + 1
F = GF(p)
n = 32
k = 16

C = codes.ReedSolomonCode(F, n, k)
C_folded = codes.ReedSolomonCode(F, n//2, k//2)
D = C.decoder("BerlekampWelch")
D_folded = C_folded.decoder("BerlekampWelch")

g_full = F.multiplicative_generator()

# A generator of an order n subgroup (undefined if one doesn't exist).
def g_n(n):
    return g_full^((p - 1) / n)

# An order n subgroup, if one exists (undefined if one doesn't exist).
def subgroup(n):
    g = g_n(n)
    return [g^i for i in range(n)]

# Evaluates the given function over an order n subgroup (undefined if one doesn't exist).
def subgroup_evals(f, n):
    return vector(f(x) for x in subgroup(n))

# Hamming weight.
def weight(v):
    return sum([bool(x) for x in v])

# Hamming distance.
def dist(v, w):
    return weight(v - w)

# Evalutes
#     f_e(x^2) = (f(x) + f(-x)) / 2
# where the inputs and outputs are evaluations over the corresponding subgroup.
def f_e(f):
    n = len(f)
    return vector((f[i] + f[i + n//2]) / F(2) for i in range(n//2))

# Evalutes
#     f_o(x^2) = (f(x) - f(-x)) / 2x
# where the inputs and outputs are evaluations over the corresponding subgroup.
def f_o(f):
    n = len(f)
    g = g_n(n)
    return vector((f[i] - f[i + n//2]) / (F(2) * g^i) for i in range(n//2))

# Performs an alpha-fold on f, i.e.
#     fold(f, alpha)(x) = f_e(x) + alpha * f_o(x)
# where the inputs and outputs are evaluations over the corresponding subgroup.
def fold(f, alpha):
    fe = f_e(f)
    fo = f_o(f)
    return vector(e + alpha * o for (e, o) in zip(f_e(f), f_o(f)))

# A pair of polynomials which intersect at 14 points:
# g^0, ..., g^6, g^(n/2), ..., g^(n/2 + 6).
def f_1(x):
    return 0
def f_2(x):
    g = g_n(n)
    prod = F(1)
    for i_left in range(7):
        i_right = n//2 + i_left
        prod *= (x - g^i_left)
        prod *= (x - g^i_right)
    return prod

# A particular vector v of length n, which has
# - 7 pairs that agrees with both f_1 and f_2
# - 4 pairs that agree with f_1 only
# - 4 pairs that agree with f_2 only
# - 1 pair that agrees with neither
def candidate_v():
    v = vector([0]*n)
    for i_left in range(n//2):
        i_right = i_left + n//2
        x_left = g_n(n)^i_left
        x_right = g_n(n)^i_right

        if i_left < 7:
            # f_1 and f_2 agree on this pair, so we can use either's value.
            assert f_1(x_left) == f_2(x_left)
            assert f_1(x_right) == f_2(x_right)
            v[i_left] = f_1(x_left)
            v[i_right] = f_1(x_right)
        elif i_left == 7:
            # This is the "random" pair, where we introduce an error which a certain alpha can cancel.
            v[i_left] = F(79) # "random"
            v[i_right] = F(156) # "random"
        elif i_left % 2 == 0:
            v[i_left] = f_1(x_left)
            v[i_right] = f_1(x_right)
        else:
            v[i_left] = f_2(x_left)
            v[i_right] = f_2(x_right)
    return v

def try_decode(D, v):
    try:
        return D_folded.decode_to_code(v)
    except sage.coding.decoder.DecodingError:
        return None

# Given some v outside the UDR, search for successful folds, where fold(v, alpha) is within the UDR.
def search_for_successful_folds(v):
    for alpha in F:
        folded_v_decoded = try_decode(D_folded, fold(v, alpha))

        if folded_v_decoded != None:
            f_1_match = folded_v_decoded == fold(subgroup_evals(f_1, n), alpha)
            f_2_match = folded_v_decoded == fold(subgroup_evals(f_2, n), alpha)

            print('Found decodable fold using alpha={}; '.format(alpha), end='')

            if f_1_match and f_2_match:
                print('fold(v) matches both fold(f_1) and fold(f_2)')
            elif f_1_match:
                print('fold(v) matches fold(f_1) only')
            elif f_2_match:
                print('fold(v) matches fold(f_2) only')

if __name__ == '__main__':
    v = candidate_v()
    search_for_successful_folds(v)