Stone prover merkle trees

stone_prover_merkle_trees.py

from Crypto.Hash import keccak

p = 0x800000000000011000000000000000000000000000000000000000000000001
r = 2^256

# Example 1
# AIR: Fibonacci 2 columns
# Blow up factor = 2
# Queries = 1

# This has 2 columns and 4 rows.
# The rows are already in bit reverse order.
coset_0_trace = [
    [0x643e5520c60d06219b27b34da0856a2c23153efe9da75c6036f362c8f19615e, 
     0x165d7fb12913882268bb8cf470c81f42349fde7dec7b0a90526d142d6a61205],
    [0x1bc1aadf39f2faee64d84cb25f7a95d3dceac1016258a39fc90c9d370e69ea2, 
     0x69a2804ed6ec78ed9744730b8f37e0bdcb6021821384f56fad92ebd2959edf4],
    [0x4de0d56f9cf97dff326c26592fbd4ae9ee756080b12c51cfe4864e9b8734f43, 
     0x1bc1aadf39f2faee64d84cb25f7a95d3dceac1016258a39fc90c9d370e69e8e],
    [0x321f2a9063068310cd93d9a6d042b516118a9f7f4ed3ae301b79b16478cb0c6, 
     0x643e5520c60d06219b27b34da0856a2c23153efe9da75c6036f362c8f196186],
]

coset_1_trace = [
    [0x5b06ce63938c3b3432a7ff993923ec5a4e6529010f447f9edc11f2e48cb4b9a,
     0x4e7d646a0d30036c68cfd97fe07427247a580ad57bb4c10459244ce5bcf25b3],
    [0x6a909188b6b8a86c6786e3a3e30e1b763ec2ce9e6bda430538895226cde38f1,
    0x248bb5ad19ea33565bcd0377421f04d2d1fdc08f5a786e7e4ad823c52d67e5],
    [0x19df03447c7ab28a0e13b5c1285d41a1bbef0fa503e8bd8e8ab8d52ae0b4d9c,
    0x6d6f08cbe37ef1587cae9a8b577001d9fb857c41b9543a950bfa2e43483b725],
    [0x20899ccf39406bf557bd6701bb70b68db6e8f8bb80f87fcd60abe5c9c4b2de2,
     0x41cad76f3db26a25b4c4bbbd53f9e6b45d029cdfd54f7d7eb634029aa7fbb50]
]

def keccak_hash(data):
    k = keccak.new(digest_bits=int(256))
    k.update(data)
    return k


def make_leaf(f1, f2):
    b = bytes.fromhex(hex(f1)[2:].zfill(64)) + bytes.fromhex(hex(f2)[2:].zfill(64))
    return keccak_hash(b).digest()


def make_leaves(trace_mont):
    leaves = bytes()
    for row in range(len(trace_mont)):
        leaves += make_leaf(trace_mont[row][0], trace_mont[row][1])
    return leaves


def next_layer(last_layer):
    result = bytes()
    for i in range(0, len(last_layer), 64):
        result += keccak_hash(last_layer[i:i+64]).digest()
    return result

# The field elemnts are in standard form. But they should be on montgomery form.
trace_mont_coset_0 = [[(x * r) % p for x in row] for row in coset_0_trace]
trace_mont_coset_1 = [[(x * r) % p for x in row] for row in coset_1_trace]

left = next_layer(next_layer(make_leaves(trace_mont_coset_0)))
right = next_layer(next_layer(make_leaves(trace_mont_coset_1)))

# Commitment
print(keccak_hash(left + right).hexdigest())


# Example 2
# AIR: Fibonacci 2 columns
# Blow up factor = 4
# Queries = 1

# This cosets already have bit reverse order.
coset1 = [
    [0x643e5520c60d06219b27b34da0856a2c23153efe9da75c6036f362c8f19615e,0x165d7fb12913882268bb8cf470c81f42349fde7dec7b0a90526d142d6a61205],
    [0x1bc1aadf39f2faee64d84cb25f7a95d3dceac1016258a39fc90c9d370e69ea2,0x69a2804ed6ec78ed9744730b8f37e0bdcb6021821384f56fad92ebd2959edf4],
    [0x4de0d56f9cf97dff326c26592fbd4ae9ee756080b12c51cfe4864e9b8734f43,0x1bc1aadf39f2faee64d84cb25f7a95d3dceac1016258a39fc90c9d370e69e8e],
    [0x321f2a9063068310cd93d9a6d042b516118a9f7f4ed3ae301b79b16478cb0c6,0x643e5520c60d06219b27b34da0856a2c23153efe9da75c6036f362c8f196186]
]

coset2 = [
    [0x5b06ce63938c3b3432a7ff993923ec5a4e6529010f447f9edc11f2e48cb4b9a,0x4e7d646a0d30036c68cfd97fe07427247a580ad57bb4c10459244ce5bcf25b3],
    [0x6a909188b6b8a86c6786e3a3e30e1b763ec2ce9e6bda430538895226cde38f1,0x248bb5ad19ea33565bcd0377421f04d2d1fdc08f5a786e7e4ad823c52d67e5],
    [0x19df03447c7ab28a0e13b5c1285d41a1bbef0fa503e8bd8e8ab8d52ae0b4d9c,0x6d6f08cbe37ef1587cae9a8b577001d9fb857c41b9543a950bfa2e43483b725],
    [0x20899ccf39406bf557bd6701bb70b68db6e8f8bb80f87fcd60abe5c9c4b2de2,0x41cad76f3db26a25b4c4bbbd53f9e6b45d029cdfd54f7d7eb634029aa7fbb50]
]

coset3 = [
    [0x5d90b27ce9277b80ea6696ad713534cdcfd36b1676ca1194a01ade19c42f9e6,0x2794106fca0d21b76c1b356a4acda3885ae79310a825c3b6a82bc704b37f17a],
    [0x63fa9e0f2d77ac8e8a8bc8d52acdb417387334c156e0d862d27aa97b2c553d9,0x1d0de13479d05236f2bbe91d893b1726bdec4c76c0dafa2faf94cfba1e0f7ba],
    [0x53317c174e8937fff589d79d7c4ee7a726d007405bc09e0e7dadd3f59deb151,0x31783f51db0dc2e3470367f152a50f8c0f91c97a993ed364fdea3bc411ebb17],
    [0x6b43335c9ad7a3209583c8dfe7ae2f73d0e958e7d69477fa0fbca47571900fa,0x9e5cf09e114ca3e5a257986d95235c4d79a56fdfdc06eb4aa552d7d1c85bc1],
]

coset4 = [
    [0x74c69a32e0e5b4cb00ca0d40b567850588dcbb5f864511f8a920d0496f77a02,0x35b2d571312e23f298fc7f775107d8e06fb68f9011090f01676b8b2eb4dfb4f],
    [0x7f2c1eecd940a07029230b410ef20a1ab682141ffc62f70fd7ce0703e1493eb,0x262555edfd44d66ee4caca0dfc04d4804e65deee76ef0c181b60fab93d62e6e],
    [0x167e7520061fc6484ee9b147d884d7b4ea2fb4d61864971fb857d868c99e0e4,0x13ddc71606bab4079825a3e4f38a1bd9cee74c2a242e4f2f1e29e29dc8a4715],
    [0x758ed1c03fb9e7ac872936366321992ad6717baa64f35fd7c6b95049e5a1139,0x104a0d8acad252a6ea131295bf6936c572fc455753d995b75f09977a4518f3a],
]

mont_coset_1 = [[(x * r) % p for x in row] for row in coset1]
mont_coset_2 = [[(x * r) % p for x in row] for row in coset2]
mont_coset_3 = [[(x * r) % p for x in row] for row in coset3]
mont_coset_4 = [[(x * r) % p for x in row] for row in coset4]

nodes0 = next_layer(next_layer(make_leaves(mont_coset_1)))
nodes1 = next_layer(next_layer(make_leaves(mont_coset_2)))
nodes2 = next_layer(next_layer(make_leaves(mont_coset_3)))
nodes3 = next_layer(next_layer(make_leaves(mont_coset_4)))

# Commitment
print(next_layer(next_layer(nodes0 + nodes1 + nodes2 + nodes3)).hex())