sumcheck.sage

# example taken from [thaler's](https://people.cs.georgetown.edu/jthaler/ProofsArgsAndZK.pdf) book chapter 3.7
# randomness in this file is completely twisted and not advisable to follow. transcript isn't implemented.

from sage.all import *
from sage.rings.polynomial.polydict import PolyDict
import unittest

p = 19
F = GF(p)
R.<x,y> = F['x, y']
G.<x> = F['x']


# G.<x> = F[]
# print(G)
# unipoly = G([0])
# print(unipoly)
# unipoly1 = G([1, 2])
# print(unipoly + unipoly1)

def fun(x, y):
	return ((1-x)*(1-y)) + (2*(1-x)*y) + (8*x*(1-y)) + (10*x*y)

P = R.interpolation(2, fun)
# print(R.interpolation(2, fun))
# print("poly:", P)
# print("monomials:", P.monomials())
# print("coefficients:", P.coefficients())
# print("variables:", P.variables())
# print(P(1, 1))

fw = [1, 2, 8, 10]
r = [int(F.random_element()), int(F.random_element())]

def to_bin(i, num):
	return list(f'{i:0{num}b}')

def to_int_list_bin(i, num):
	return [int(x) for x in to_bin(i, num)]

def delta_w(a, b):
	return (a*b)+((1-a)*(1-b))

def delta_step(i, r):
	res = 1
	len_r = len(r)
	# print(len_r, f'{i:0{len_r}b}')
	bin_i = to_bin(i, len_r)
	for j in range(len_r):
		res *= delta_w(r[j], int(bin_i[j]))

	# print("res: ", res)
	return res

def mle(fw, r, p):
	delta_w = 0
	for i in range(len(fw)):
		delta_w += fw[i] * delta_step(i, r)

	return delta_w % p

assert(mle(fw, r, p) == int(P(r)))

def memoize(r, n):
	if n == 1:
		return [delta_w(0, r[n-1]), delta_w(1, r[n-1])]
	else:
		vals = memoize(r, n-1)
		val = []
		for i in range(len(vals)):
			val.append(vals[i] * delta_w(0, r[n-1]))
			val.append(vals[i] * delta_w(1, r[n-1]))
		return val

def dynamic_mle(fw, r, p):
	delta_lookup = memoize(r, len(r))
	delta_w = 0
	for i in range(len(fw)):
		delta_w += fw[i] * delta_lookup[i]
	return delta_w % p

assert(dynamic_mle(fw, r, p) == int(P(r)))

def num_vars(P: PolyDict):
	"""
	returns number of variables in a multiliear polynomial
	"""
	return len(P.exponents()[0])

def max_degree(P: PolyDict):
	"""
	params:
	- `P`: multilinear polynomial

	returns:
	- `lookup_degree`: list containing max degree of each variable
	"""

	lookup_degree = [0] * num_vars(P)

	for term in P.dict():
		for i in range(len(term)):
			if term[i] > lookup_degree[i]:
				lookup_degree[i] = term[i]

	return lookup_degree

def evaluate(P: PolyDict, F, r):
	"""
	params:
	- `P`: multilinear polynomial
	- `F`: field
	- `r`: list of random variables

	returns:
	- `result`: evaluation of multilinear polynomial at a specific point
	"""
	terms = P.dict()
	result = F(0)

	for term, coeff in terms.items():
		prod = 1
		for i in range(len(term)):
			prod *= F(int(r[i])) ** term[i]
		result += coeff * prod

	return result

def evaluate_multipoly_hypercube(F, P):
	"""
	evaluate multilinear polynomial at all points in hypercube
	"""
	result = F(0)
	n = 2**num_vars(P)
	for i in range(n):
		result += evaluate(P, F, to_bin(i, num_vars(P)))
	return result

# use this: https://doc.sagemath.org/html/en/reference/polynomial_rings/sage/rings/polynomial/polydict.html#sage.rings.polynomial.polydict.PolyDict
class Sumcheck:
	def __init__(self, F, P: PolyDict):
		"""
		params:
		- `F`: field
		- `P`: multilinear poly
		"""
		self.F = F
		self.P = P
		self.r = [] # random vars
		G.<x> = F['x'] # unipoly
		self.G = G
		self.x = x # poly invariate

	def evaluate_term(self, term, points):
		"""
		evaluate a monomial in multilinear polynomial at a specific point and substitue
		initial variables with random variables generated by verifier.
		"""
		len_r = len(self.r)
		coeff = F(1)
		poly = self.G([0])
		for i in range(len(term)):
			if i < len_r:
				coeff = coeff * self.r[i] ** (term[i])
			elif i == len_r:
				poly = self.x ** (term[i])
			else:
				coeff = coeff * (int(points[i - len_r]) ** (term[i]))

		poly = poly * coeff
		return poly

	def evaluate_gj(self, points):
		"""
		evaluate multilinear polynomial at a specific point and substitue initial
		variables with random vars provided by verifier
		"""
		terms = self.P.dict()

		uniPoly = self.G([0])

		for term, coeff in terms.items():
			unipoly_term = self.evaluate_term(term, points)
			uniPoly = uniPoly + coeff*unipoly_term

		return uniPoly

	def prove_interactive(self, r):
		"""
		takes a random variable r and creates univariate polynomial by evaluating
		poly `G(r0, r1, ..., rk, xk+1, ..., xn)`.
		"""
		if r != None:
			self.r.append(r)

		v = num_vars(self.P) - len(self.r)
		unipoly = self.G([0])

		for i in range(0, 2**(v-1)):
			poly = self.evaluate_gj(to_bin(i, v))
			# print(poly)
			unipoly = unipoly + poly

		return unipoly

	def prove_non_interactive(self):
		randomness = []
		uni_polys = []
		r = None
		for i in range(num_vars(self.P)):
			unipoly = self.prove_interactive(r)
			uni_polys.append(unipoly)
			r = F.random_element()
			randomness.append(r)

		return (uni_polys, randomness)

	def verify_interactive(self, c):
		"""
		runs sumcheck protocol and verifies the result
		"""
		lookup_degree_table = max_degree(self.P)

		gi = self.prove_interactive(None)
		print("round 0: ", gi)

		# print(gi, gi(0), gi(1))
		expected = gi(0) + gi(1)
		# print(expected)
		assert(expected == c)
		assert(gi.degree() <= lookup_degree_table[0])

		for i in range(1, num_vars(self.P)):
			r = F.random_element()
			actual = gi(r)
			gi = self.prove_interactive(r)
			print(f"round {i}: {gi}")
			expected = gi(0) + gi(1)
			assert(actual == expected)
			assert(gi.degree() <= lookup_degree_table[i])

		r = F.random_element()
		expected = gi(r)
		self.r.append(r)
		actual = evaluate(self.P, self.F, self.r)
		print("last round: ", actual, expected)
		return actual == expected

	def verify_non_interactive(self, proof, r, sum):
		"""
		runs sumcheck protocol non interactively and verifies the result
		TODO: implement transcript
		"""

		lookup_degree_table = max_degree(self.P)

		for i in range(num_vars(self.P)):
			unipoly = proof[i]
			expected = unipoly(0) + unipoly(1)
			assert(expected == sum)
			assert(unipoly.degree() <= lookup_degree_table[i])

			sum = unipoly(r[i])

		expected = sum
		actual = evaluate(self.P, self.F, r)
		assert(expected == actual)


class TestSumcheck(unittest.TestCase):
	def test_sumcheck(self):
		multipoly1 = PolyDict({(3, 0, 0): F(2), (1, 0, 1): F(1), (0, 1, 1): F(1)})
		print("starting sumcheck with poly: ", multipoly1.latex(['x1', 'x2', 'x3']))
		sumcheck = Sumcheck(F, multipoly1)
		evaluation = evaluate_multipoly_hypercube(sumcheck.F, sumcheck.P)
		assert(sumcheck.verify_interactive(evaluation))

		multipoly = PolyDict({(2, 3, 2): F(2), (1, 0, 1): F(3), (2, 1, 0): F(4)})
		print("starting sumcheck with poly: ", multipoly1.latex(['x1', 'x2', 'x3']))
		sumcheck = Sumcheck(F, multipoly)
		evaluation = evaluate_multipoly_hypercube(sumcheck.F, sumcheck.P)
		assert(sumcheck.verify_interactive(evaluation))

	def test_poly_funcs(self):
		multipoly1 = PolyDict({(3, 0, 0): F(2), (1, 0, 1): F(1), (0, 1, 1): F(1)})

		num = num_vars(multipoly1)
		assert(num == 3)

		degree = max_degree(multipoly1)
		assert(degree == [3, 1, 1])

		evaluation = evaluate(multipoly1, F, [1, 0, 1])
		assert(evaluation == F(3))

	def test_non_interactive(self):
		multipoly1 = PolyDict({(3, 0, 0): F(2), (1, 0, 1): F(1), (0, 1, 1): F(1)})
		print("starting sumcheck with poly: ", multipoly1.latex(['x1', 'x2', 'x3']))
		sumcheck = Sumcheck(F, multipoly1)
		proof, r = sumcheck.prove_non_interactive()
		sum = evaluate_multipoly_hypercube(sumcheck.F, sumcheck.P)
		sumcheck.verify_non_interactive(proof, r, sum)

# if __name__ == '__main__':
#     unittest.main()