circ_mat.py

import itertools as itt

from sage.all import *

# Fix k and generate random integers
k = 20
Zk = Zmod(2**k)
while True:
    # a^2 + b^2 = 1
    a = Zk.random_element()
    if (1 - a**2).is_square():
        b = (1 - a**2).sqrt()
        break

print(f'{a = } {b = }')

bin_a = bin(int(a))[2:].zfill(k)[::-1]
bin_b = bin(int(b))[2:].zfill(k)[::-1]

# Leak 50% of bits
bin_a, bin_b = list(bin_a), list(bin_b)
while bin_a.count('*') < k//2:
    bin_a[randint(0, k - 1)] = '*'
while bin_b.count('*') < k//2:
    bin_b[randint(0, k - 1)] = '*'

bin_a, bin_b = ''.join(bin_a), ''.join(bin_b)

# Create the polynomial ring
names = [f'x_{i}' for i in range(k)] + [f'y_{i}' for i in range(k)]
R, vs = PolynomialRing(Zk, names).objgens()
xs, ys = vs[:k], vs[k:]

alpha, beta = 0, 0
for i in range(k):
    if bin_a[i] in '01':
        alpha += 2**i * int(bin_a[i])
    else:
        alpha += 2**i * xs[i]

    if bin_b[i] in '01':
        beta += 2**i * int(bin_b[i])
    else:
        beta += 2**i * ys[i]

# Polyring caching
Rs = []
for i in range(k+1):
    Zi = Zmod(2**i)
    Ri = PolynomialRing(Zi, names)
    Rs.append(Ri)

# Find solutions
q = [(alpha, beta, 0)] # (a, b, c) where equation is correct up to 2**c

while True:
    nxt = q.pop(0)
    alpha, beta, i = nxt
    if i >= k:
        # Got to the end
        q.insert(0, nxt)
        break

    Ri = Rs[i+1]
    eq = Ri(alpha)**2 + Ri(beta)**2

    vs = eq.variables()
    if len(vs) == 0:
        # No new vars, just check the equation
        if eq == 1:
            q.append((alpha, beta, i+1))
        else:
            continue

    for s in itt.product([0, 1], repeat=len(vs)):
        sub_dic = {R(vs[i]):s[i] for i in range(len(vs))}
        alpha1, beta1 = alpha.subs(sub_dic), beta.subs(sub_dic)

        check = Ri(alpha1)**2 + Ri(beta1)**2
        if check == 1:
            q.append((alpha1, beta1, i+1))

# Check all solutions
q = set(q) # Lazy workaround for repeated solutions
print(f'{len(q) = }')
for sol in q:
    alpha, beta, _ = sol
    vs = (alpha + beta).variables()
    for s in itt.product([0, 1], repeat=len(vs)):
        sub_dic = {vs[i]:s[i] for i in range(len(vs))}
        alpha1, beta1 = int(alpha.subs(sub_dic)), int(beta.subs(sub_dic))

        if (alpha1, beta1) == (a, b):
            print(f'FOUND {alpha1 = } {beta1 = }')