darkfi/bin/zkrunner/zkrunner.py

from argparse import ArgumentParser
from darkfi_sdk_py.affine import Affine
from darkfi_sdk_py.base import Base
from darkfi_sdk_py.point import Point
from darkfi_sdk_py.proof import Proof
from darkfi_sdk_py.proving_key import ProvingKey
from darkfi_sdk_py.scalar import Scalar
from darkfi_sdk_py.verifying_key import VerifyingKey
from darkfi_sdk_py.zk_binary import ZkBinary
from darkfi_sdk_py.zk_circuit import ZkCircuit
from pprint import pprint
from sys import getsizeof
from time import time
import argparse


def heap_add(heap, element):
    heap.append(element)
    vprint(f"HEAP: {heap}")


def pubins_add(pubins, element):
    pubins.append(element)
    vprint(f"PUBLIC INPUTS: {pubins}")


def get_pubins(statements, witnesses, constant_count, literals):
    # Python heap for executing zk statements
    heap = [None] * constant_count + witnesses
    pubins = []
    for stmt in statements:
        vprint(f"STATEMENT: {stmt}")
        opcode, args = stmt[0], stmt[1]
        if opcode == 'BaseAdd':
            a = heap[args[0][1]]
            b = heap[args[1][1]]
            heap_add(heap, a + b)
        elif opcode == 'BaseMul':
            a = heap[args[0][1]]
            b = heap[args[1][1]]
            heap_add(heap, a * b)
        elif opcode == 'BaseSub':
            a = heap[args[0][1]]
            b = heap[args[1][1]]
            heap_add(heap, a - b)
        elif opcode == 'EcAdd':
            a = heap[args[0][1]]
            b = heap[args[1][1]]
            heap_add(heap, a + b)
        elif opcode == 'EcMul':
            a = heap[args[0][1]]
            heap_add(heap, Point.mul_r_generator(a))
        elif opcode in {'EcMulBase', 'EcMulVarBase'}:
            i = args[0][1]
            base = heap[i]
            product = Point.mul_base(base)
            heap_add(heap, product)
        elif opcode == 'EcMulShort':
            value = heap[args[0][1]]
            heap_add(heap, Point.mul_short(value))
        elif opcode == 'EcGetX':
            i = args[0][1]
            point = heap[i]
            x, _ = point.to_affine().coordinates()
            heap_add(heap, x)
        elif opcode == 'EcGetY':
            i = args[0][1]
            point = heap[i]
            _, y = point.to_affine().coordinates()
            heap_add(heap, y)
        elif opcode == 'PoseidonHash':
            messages = [heap[m[1]] for m in args]
            heap_add(heap, Base.poseidon_hash(messages))
        elif opcode == 'MerkleRoot':
            i = heap[args[0][1]]
            p = heap[args[1][1]]
            a = heap[args[2][1]]
            heap_add(heap, Base.merkle_root(i, p, a))
        elif opcode == 'ConstrainInstance':
            i = args[0][1]
            element = heap[i]
            pubins_add(pubins, element)
        elif opcode == 'WitnessBase':
            type = args[0][0]
            assert type == 'Lit', f"type should LitType instead of {type}"
            i = args[0][1]
            element = int(literals[i][1])  # (LitType, Lit)
            base = Base.from_u64(element)
            heap_add(heap, base)
        elif opcode == 'CondSelect':
            cnd = heap[args[0][1]]
            thn = heap[args[1][1]]
            els = heap[args[2][1]]
            assert cnd == Base.from_u64(0) or cnd == Base.from_u64(
                1), "Failed bool check"
            res = thn if cnd == Base.from_u64(1) else els
            heap_add(heap, res)
        elif opcode in IGNORED_OPCODES:
            vprint(f"IGNORE: {opcode}")
        else:
            vprint(f"NO IMPLEMENTATION: {opcode}")
    return pubins


def bincode_data(bincode):
    with open(bincode, "rb") as f:
        bincode = f.read()
        zkbin = ZkBinary.decode(bincode)
        return {
            "zkbin": zkbin,
            "namespace": zkbin.namespace(),
            "witnesses": zkbin.witnesses(),
            "constant_count": zkbin.constant_count(),
            "statements": zkbin.opcodes(),
            "literals": zkbin.literals(),
            "k": zkbin.k()
        }


IGNORED_OPCODES = {
    'Noop', 'RangeCheck', 'LessThanStrict', 'LessThanLoose', 'BoolCheck',
    'ConstrainEqualBase', 'ConstrainEqualPoint', 'DebugPrint'
}

if __name__ == "__main__":

    ##### Your Inputs #####

    bincode_path = "opcodes.no-nipoint.zk.bin"
    witnesses = [
        Base.from_u64(3),
        Scalar.from_u64(4),
        Base.from_u64(5),
        Base.from_u64(6),
        Base.from_u64(7),
        Base.from_u64(8),
        10,
        [Base.from_u64(42)] * 32,
        Base.from_u64(1),
    ]

    ##### Setup #####

    bincode_data_ = bincode_data(bincode_path)
    zkbin = bincode_data_['zkbin']
    statements = bincode_data_['statements']
    constant_count = bincode_data_['constant_count']
    literals = bincode_data_['literals']
    K = bincode_data_['k']

    # Verbosity
    parser = argparse.ArgumentParser()
    verbose = parser.add_argument(
        '--verbose', action='store_true', help='verbose switch')
    args = parser.parse_args()
    vprint = print if args.verbose else lambda *a, **k: None

    ##### Proving #####

    pubins = get_pubins(statements, witnesses, constant_count, literals)

    zkcircuit = ZkCircuit(zkbin)
    zkcircuit.witness_base(witnesses[0])
    zkcircuit.witness_scalar(witnesses[1])
    zkcircuit.witness_base(witnesses[2])
    zkcircuit.witness_base(witnesses[3])
    zkcircuit.witness_base(witnesses[4])
    zkcircuit.witness_base(witnesses[5])
    zkcircuit.witness_u32(witnesses[6])
    zkcircuit.witness_merkle_path(witnesses[7])
    zkcircuit.witness_base(witnesses[8])
    zkcircuit = zkcircuit.build(zkbin)

    print("Making proving key.....")
    start = time()
    proving_key = ProvingKey.build(K, zkcircuit)
    print(f"Time for making proving key: {time() - start}")

    print("Proving.....")
    start = time()
    proof = Proof.create(proving_key, [zkcircuit], pubins)
    print(f"Time for proving: {time() - start}")

    ##### Verifiying #####

    zkcircuit_v = zkcircuit.verifier_build(zkbin)

    print(f"Making verifying key.....")
    start = time()
    verifying_key = VerifyingKey.build(K, zkcircuit_v)
    print(f"Time for making verifying key: {time() - start}")

    print("Verifying.....")
    start = time()
    proof.verify(verifying_key, pubins)
    print(f"Time for verifying {time() - start}")