babysnark_opt.py

#| # Baby SNARK with FFT optimizations
#| This is the same algorithm as BabySNARK, except optimized so that
#| the overhead is quasilinear rather than quadratic (in the number of
#| constraints).
from babysnark import *

from finitefield.finitefield import FiniteField
from finitefield.polynomial import polynomialsOver
from ssbls12 import Fp, Poly, Group
G = Group.G
GT = Group.GT

#| # Choosing roots of unity
#| The BLS12-381 is chosen in part because it's FFT friendly. To use radix-2
#| FFT, we need to find m^th roots of unity, where m is a power of two, and
#| m is the degree bound of the polynomial we want to represent.
#|
#| In the BLS12-381, we can find primitive n^th roots of unity, for any
#| power of two n up to n <= 2^**32.
#| This follows because for the ssbls12-381 exponent field Fp, we have
#|    2^32 divides (p - 1).
from polynomial_evalrep import get_omega, polynomialsEvalRep, RowDictSparseMatrix

omega_base = get_omega(Fp, 2**32, seed=0)

# These are globals. They're set to be large enough for small examples, and in
# babysnark_setup we'll change them for other sizes anyway.
mpow2 = 128  # nearest power of 2 greather than m, the number of constraints
omega = omega_base ** (2**32 // mpow2)

# This overwrite a global parameter imported from babysnark.
ROOTS.clear()
ROOTS += [omega**i for i in range(mpow2)]

def vanishing_poly(omega, n):
    # For the special case of evaluating at all n powers of omega,
    # the vanishing poly has a special form.
    #  t(X) = (X-1)(X-omega)....(X-omega^(n-1)) = X^n - 1
    return Poly([Fp(-1)] + [Fp(0)]*(n-1) + [Fp(1)])

#| # Evaluation Representation of polynomials
#|
#| This representation is sparse - it only stores the non-zero values,
#| so if it has lots of roots among the powers of omega,
def _sparse_poly_demo():
    PolyEvalRep = polynomialsEvalRep(Fp, omega, mpow2)

    # Example polynomial that has roots at most of the powers of omega
    xs = [omega**1, omega**4, omega**5]
    ys = [Fp(3), Fp(5), Fp(1)]
    f_rep = PolyEvalRep(xs, ys)

    for i in [0,2,3,6,7]:
        assert f_rep(omega**i) == Fp(0)
    for i, x in enumerate(xs):
        assert f_rep(x) == ys[i]

    # Convert to coeffs and back
    f = f_rep.to_coeffs()
    assert f_rep.to_coeffs() == PolyEvalRep.from_coeffs(f).to_coeffs()

    # Check f and f_rep are consistent
    tau = random_fp()
    assert f(tau) == f_rep(tau)

    # print('f_rep:', f_rep)
    # print('f:', f)
_sparse_poly_demo()

#| # Sparse representation of Square Constraint Systems
#|
#| In order to reach our goal of quasilinear overhead (in the number of gates),
#| we can't just use the dense numpy matrix U to represent our constraints.
#|
#| Suppose we use the construction from boolean circuits, and that the circuit
#| has an average fan-in of some constant (e.g., average fan-in of ~2). Then the
#| matrix U, will be sparse, with only O(m + n) non-zero values despite being
#| an (m * n) matrix.
#|

# Matrix is m-by-n, but contains only avgPerN*n non-zero values in expectation.
# The first column is all non-zero.
def random_sparse_matrix(m, n, avgPerN=2):
    U = RowDictSparseMatrix(m, n, Fp(0))

    # First fill the first column
    for row in range(m):
        U[row,0] = random_fp()

    # Then pick randomly for the rest
    for _ in range(avgPerN * n - 1):
        row = random.randrange(m)
        col = random.randrange(n)
        U[row,col] = random_fp()

    return U

def generate_solved_instance(m, n):
    """
    Generates a random circuit and satisfying witness
    U, (stmt, wit)
    """
    # Generate a, U
    a = np.array([random_fp() for _ in range(n)])
    U = random_sparse_matrix(m, n)

    # Normalize U to satisfy constraints
    Ua2 = U.dot(a) * U.dot(a)
    for (i,j), val in U.items():
        U[i,j] /= Ua2[i].sqrt()

    assert((U.dot(a) * U.dot(a) == 1).all())
    Ud = U.to_dense()
    assert((Ud.dot(a) * Ud.dot(a) == 1).all())
    return U, a

#-
# Example
m, n = 10, 12
U = random_sparse_matrix(m, n)
U, a = generate_solved_instance(m, n)
print(U)
print(U.to_dense())

#| # Baby SNARK with quasilinear overhead

# Setup
def babysnarkopt_setup(U, n_stmt):
    (m, n) = U.shape
    assert n_stmt < n

    # Generate roots for each gate
    # TODO: Handle padding?
    global mpow2, omega
    mpow2 = m
    assert mpow2 & mpow2 - 1 == 0, "mpow2 must be a power of 2"
    omega = omega_base ** (2**32 // mpow2)
    PolyEvalRep = polynomialsEvalRep(Fp, omega, mpow2)

    global ROOTS
    if len(ROOTS) != m:
        ROOTS.clear()
        ROOTS += [omega**i for i in range(m)]

    # Generate polynomials u from columns of U
    Us = [PolyEvalRep((), ()) for _ in range(n)]
    for (i,k), y in U.items():
        x = ROOTS[i]
        Us[k] += PolyEvalRep([x],[y])

    # Trapdoors
    global tau, beta, gamma
    tau   = random_fp()
    beta  = random_fp()
    gamma = random_fp()

    # CRS elements
    CRS = [G * (tau ** i) for i in range(m+1)] + \
          [G * gamma, G * (beta * gamma)] + \
          [G * (beta * Ui(tau)) for Ui in Us[n_stmt:]]

    # Precomputation
    # Note: This is not considered part of the trusted setup, since it
    # could be computed direcftly from the G * (tau **i) terms.

    # Compute the target poly term
    t = vanishing_poly(omega, m)
    T = G * t(tau)

    # Evaluate the Ui's corresponding to statement values
    Uis = [G * Ui(tau) for Ui in Us]
    precomp = Uis, T

    return CRS, precomp


# Prover
def babysnarkopt_prover(U, n_stmt, CRS, precomp, a):
    (m, n) = U.shape
    assert n == len(a)
    assert len(CRS) == (m + 1) + 2 + (n - n_stmt)

    taus = CRS[:m+1]
    bUis = CRS[-(n-n_stmt):]

    Uis, T = precomp

    # Target is the vanishing polynomial
    mpow2 = m
    assert mpow2 & mpow2 - 1 == 0, "mpow2 must be a power of 2"
    omega  = omega_base ** (2**32 // mpow2)
    omega2 = omega_base ** (2**32 // (2*mpow2))
    PolyEvalRep = polynomialsEvalRep(Fp, omega, mpow2)
    t = vanishing_poly(omega, mpow2)
    
    # 1. Find the polynomial p(X)

    # First compute v
    v = PolyEvalRep((),())
    for (i,k), y in U.items():
        x = ROOTS[i]
        v += PolyEvalRep([x], [y]) * a[k]

    # Now we need to convert between representations to multiply and divide
    PolyEvalRep2 = polynomialsEvalRep(Fp, omega2, 2*mpow2)
    roots2 = [omega2**i for i in range(2*mpow2)]
    ONE = PolyEvalRep2(roots2, [Fp(1) for _ in roots2])

    vv = v.to_coeffs()
    v2 = PolyEvalRep2.from_coeffs(v.to_coeffs())
    p2 = v2 * v2 - ONE
    p = p2.to_coeffs()

    # Find the polynomial h by dividing p / t
    h = PolyEvalRep2.divideWithCoset(p, t)
    # assert p == h * t

    # 2. Compute the H term
    H = evaluate_in_exponent(taus, h)

    # 3. Compute the Vw terms, using precomputed Uis
    Vw = sum([Uis[k] * a[k] for k in range(n_stmt, n)], G*0)
    # assert G * vw(tau) == Vw

    # 4. Compute the Bw terms
    Bw = sum([bUis[k-n_stmt] * a[k] for k in range(n_stmt, n)], G*0)
    # assert G * (beta * vw(tau)) == Bw

    # V = G * vv(tau)
    # assert H.pair(T) * GT == V.pair(V)

    # print('H:', H)
    # print('Bw:', Bw)
    # print('Vw:', Vw)
    return H, Bw, Vw

#-
if __name__ == '__main__':
    # Sample a problem instance
    print("Generating a Square Span Program instance")
    n_stmt = 4
    m,n = (16, 6)
    U, a = generate_solved_instance(m, n)
    a_stmt = a[:n_stmt]
    print('U:', repr(U))
    print('a_stmt:', a_stmt)
    print('m x n:', m * n)

    # Setup
    print("Computing Setup (optimized)...")
    CRS, precomp = babysnarkopt_setup(U, n_stmt)
    print("CRS length:", len(CRS))

    # Prover
    print("Proving (optimized)...")
    H, Bw, Vw = babysnarkopt_prover(U, n_stmt, CRS, precomp, a)

    if 0:  # Uncomment this to cross-check the optimized with the reference
        # Alternate prover
        print("Proving (reference)...")
        H_, Bw_, Vw_ = babysnark_prover(U.to_dense(), n_stmt, CRS, precomp, a)
        assert (H_, Bw_, Vw_) == (H, Bw, Vw)

    # Verifier
    print("[opt] Verifying (optimized)...")
    babysnark_verifier(U, CRS, precomp, a[:n_stmt], (H, Bw, Vw))