zkrunner: Perform a rewrite with the cleaned up python bindings.

bin/zkrunner/README.md

@@ -1,35 +1,34 @@
-# What is this?
+zkrunner
+========
 
-`zkrunner` is a simple Python script that calls into the Darkfi Python SDK.
-The Python SDK provides APIs such as creating a circuit, assigning the witness to the circuit and more.
+`zkrunner` is a simple Python script using the DarkFi SDK Python
+bindings providing a CLI for prototyping zkas proofs.
 
-`zkrunner` uses the Python SDK to create a developement environment for zkas developer where:
-* the ZKAS developer provides the ZKAS binary code
-* the ZKAS developer provides the witness and assigns it accordingly
-* zkrunner:
-	* sets up the circuit from the binary
-	* generates both proving and verifying key
-	* creates the proof from the witness and proving key
-	* creates the public inputs from the witness
-	* verifies the proof using the public inputs and verifying key
-* zkrunner times each step as a basic performance benchmark
+## Usage
 
-This is so developers have an easier time to test their zkas circuit.
+Refer to the [README.md of the python bindings](../../src/sdk/python/README.md)
+to see how to install and use them. They're necessary for zkrunner to
+work properly.
 
-# Installation
+Help text:
 
-Follow the guide in src/sdk/python/README.md to install the Python bindings and virtual environment.
+```
+$ zkrunner.py -h
+```
 
-# Getting Started
+Running a demo:
 
-* Compile the ZKAS source to ZKAS binary
 ```
-cd <darkmap>
-zkas proof/set_v1.zk
+$ witness_gen.py > witness.json
+$ zkrunner.py -w witness.json opcodes.zk
 ```
-* Open up `zkrunner.py`, read over the TODOs and comments, provide the path to zkas binary code, witness and assign accordingly.
-* After installing Python bindings in your Python installation, simply run `python zkrunner.py [--verbose]`.
 
-# Notes
+The program expects a path to a `witness.json` file containing the
+information about witnesses and public inputs for the proof, and a
+path to a zkas circuit source code (does not have to be compiled).
+
+Once executed, zkrunner will attempt to create and verify the proof.
+
+## Creating witnesses
 
-* "witness" and "witnesses" are used interchangablely.
+Refer to the `witness_gen.py` file.

bin/zkrunner/witness_gen.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+# This file is part of DarkFi (https://dark.fi)
+#
+# Copyright (C) 2020-2023 Dyne.org foundation
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+Example witness generation for some circuit.
+Here we generate them for opcodes.zk
+
+This reflects /darkfi/tests/zkvm_opcodes.rs
+"""
+import json
+from darkfi_sdk.pasta import Ep, Fp, Fq, nullifier_k, EpAffine, mod_r_p
+from darkfi_sdk.crypto import poseidon_hash, pedersen_commitment_u64
+from darkfi_sdk.merkle import MerkleTree
+
+# Creating base elements and scalars
+value = 666
+value_blind = Fq.random()
+blind = Fp.random()
+secret = Fp.random()
+a = Fp.from_u64(42)
+b = Fp.from_u64(69)
+
+# Creating a Merkle tree (internally using bridgetree)
+tree = MerkleTree()
+c0 = Fp.random()
+c1 = Fp.random()
+c2 = poseidon_hash([Fp.one(), Fp.from_u64(2), blind])
+c3 = Fp.random()
+
+# Appending and marking leaves in the Merkle tree
+tree.append(c0)
+tree.mark()
+tree.append(c1)
+tree.append(c2)
+leaf_pos = tree.mark()
+tree.append(c3)
+tree.mark()
+
+# Calculating the tree root and authentication path
+root = tree.root(0)
+path = tree.witness(leaf_pos, 0)
+
+# Elliptic curve multiplication
+ephem_secret = Fp.random()
+pubkey = Ep.from_affine(nullifier_k()) * mod_r_p(ephem_secret)
+
+ephem_public = pubkey * mod_r_p(ephem_secret)
+ephem_x, ephem_y = EpAffine.from_projective(ephem_public).coordinates()
+
+value_commit = pedersen_commitment_u64(value, value_blind)
+value_coords = EpAffine.from_projective(value_commit).coordinates()
+d = poseidon_hash([Fp.one(), blind, value_coords[0], value_coords[1]])
+
+public = Ep.from_affine(nullifier_k()) * mod_r_p(secret)
+pub_x, pub_y = EpAffine.from_projective(public).coordinates()
+
+# Create the object representing the JSON witnesses file.
+w = {}
+
+# Private witnesses for the proof
+# yapf: disable
+w["witnesses"] = [
+    {"Base": str(Fp.from_u64(value))},
+    {"Scalar": str(value_blind)},
+    {"Base": str(blind)},
+    {"Base": str(a)},
+    {"Base": str(b)},
+    {"Base": str(secret)},
+    {"EcNiPoint": EpAffine.from_projective(pubkey).coordinates_str()},
+    {"Base": str(ephem_secret)},
+    {"Uint32": leaf_pos},
+    {"MerklePath": [str(i) for i in path]},
+    {"Base": str(Fp.one())},
+]
+
+# Public inputs for the proof
+# yapf: disable
+w["instances"] = [
+    str(value_coords[0]),
+    str(value_coords[1]),
+    str(c2),
+    str(d),
+    str(root),
+    str(pub_x),
+    str(pub_y),
+    str(ephem_x),
+    str(ephem_y),
+    str(a),
+    str(Fp.zero()),
+]
+
+# Printing the expected JSON file used by zkrunner.
+print(json.dumps(w, indent=2))

bin/zkrunner/zkrunner.py

@@ -1,197 +1,152 @@
-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'
-}
+#!/usr/bin/env python3
+# This file is part of DarkFi (https://dark.fi)
+#
+# Copyright (C) 2020-2023 Dyne.org foundation
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+Python tool to prototype zkVM proofs given zkas source code and necessary
+witness values in JSON format.
+"""
+import json
+from darkfi_sdk.pasta import Fp, Fq, Ep
+from darkfi_sdk.zkas import (MockProver, ZkBinary, ZkCircuit, ProvingKey,
+                             Proof, VerifyingKey)
+
+
+def main(witness_file, source_file, mock=False):
+    """main zkrunner logic"""
+    # We will first attempt to decode the witnesses from the JSON file.
+    # Refer to the `witness_gen.py` file to see what the format of this
+    # file should be.
+    print("Decoding witnesses...")
+    with open(witness_file, "r", encoding="utf-8") as json_file:
+        witness_data = json.load(json_file)
+
+    # Then we attempt to compile the given zkas code and create a
+    # zkVM circuit. This compiling logic happens in the Python bindings'
+    # `ZkBinary::new` function, and should be equivalent to the actual
+    # `zkas` binary provided in the DarkFi codebase.
+    print("Compiling zkas code...")
+    with open(source_file, "r", encoding="utf-8") as zkas_file:
+        zkas_source = zkas_file.read()
+
+    # This line will compile the source code
+    zkbin = ZkBinary(source_file, zkas_source)
+
+    # Construct the initial circuit object.
+    circuit = ZkCircuit(zkbin)
+
+    # If we want to build an actual proof, we'll need a proving key
+    # and a verifying key.
+    # circuit.verifier_build() is called so that the inital circuit
+    # (which contains no witnesses) actually calls empty_witnesses()
+    # in order to have the correct code path when the circuit gets
+    # synthesized.
+    if not mock:
+        print("Building proving key...")
+        proving_key = ProvingKey.build(zkbin.k(), circuit.verifier_build())
+
+        print("Building verifying key...")
+        verifying_key = VerifyingKey.build(zkbin.k(), circuit.verifier_build())
+
+    # Now we scan through the parsed JSON witness file and
+    # build our "heap". These will be appended to the initial
+    # circuit and decide the code path for the prover.
+    for witness in witness_data["witnesses"]:
+        assert len(witness) == 1
+        if value := witness.get("EcPoint"):
+            circuit.witness_ecpoint(Ep(value))
+
+        elif value := witness.get("EcNiPoint"):
+            assert len(value) == 2
+            xcoord, ycoord = Fp(value[0]), Fp(value[1])
+            circuit.witness_ecnipoint(Ep(xcoord, ycoord))
+
+        elif value := witness.get("Base"):
+            circuit.witness_base(Fp(value))
+
+        elif value := witness.get("Scalar"):
+            circuit.witness_scalar(Fq(value))
+
+        elif value := witness.get("MerklePath"):
+            path = [Fp(i) for i in value]
+            assert len(path) == 32
+            circuit.witness_merklepath(path)
+
+        elif value := witness.get("Uint32"):
+            circuit.witness_uint32(value)
+
+        elif value := witness.get("Uint64"):
+            circuit.witness_uint64(value)
 
-if __name__ == "__main__":
+        else:
+            raise ValueError(f"Invalid Witness type for witness {witness}")
 
-    # TODO: relative path to your zkas binary
-    bincode_path = "set_v1.zk.bin"
-
-    ##### 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']
-
-    ##### TODO: Your Inputs #####
-
-    # TODO: list of witnesses, in the same order as in the zkas circuit witness section
-    witnesses = [
-        Base.from_u64(42),
-        Base.from_u64(1),
-        Base.from_u64(1),
-        Base.from_u64(1),
-        Base.from_u64(1),
-    ]
-
-    zkcircuit = ZkCircuit(zkbin)
-
-    # TODO: call the corresponding witness_* prefixed methods to assign the witness
-    # to the circuit. For a complete list, `rgrep witness_ <darkfi>/src/sdk/python`
-    zkcircuit.witness_base(witnesses[0])
-    zkcircuit.witness_base(witnesses[1])
-    zkcircuit.witness_base(witnesses[2])
-    zkcircuit.witness_base(witnesses[3])
-    zkcircuit.witness_base(witnesses[4])
-
-    zkcircuit = zkcircuit.build(zkbin)
-
-    # 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
+    # circuit.prover_build() will actually construct the circuit
+    # with the values witnessed above.
+    circuit = circuit.prover_build()
 
-    ##### Proving #####
+    # Instances are our public inputs for the proof and they're also
+    # part of the JSON file.
+    instances = []
+    for instance in witness_data["instances"]:
+        instances.append(Fp(instance))
 
-    pubins = get_pubins(statements, witnesses, constant_count, literals)
+    # If we're building an actual proof, we'll use the ProvingKey to
+    # prove and our VerifyingKey to verify the proof.
+    if not mock:
+        print("Proving knowledge of witnesses...")
+        proof = Proof.create(proving_key, [circuit], instances)
 
-    vprint("Making proving key.....")
-    start = time()
-    proving_key = ProvingKey.build(K, zkcircuit)
-    print(f"Time for making proving key: {time() - start}")
+        print("Verifying ZK proof...")
+        proof.verify(verifying_key, instances)
 
-    vprint("Proving.....")
-    start = time()
-    proof = Proof.create(proving_key, [zkcircuit], pubins)
-    # TODO: consider persisting the proof for making a transaction
-    print(f"Time for proving: {time() - start}")
+    # Otherwise, we'll simply run the MockProver:
+    else:
+        print("Running MockProver...")
+        proof = MockProver.run(zkbin.k(), circuit, instances)
+        print("Verifying MockProver...")
+        proof.verify()
 
-    ##### Verifiying #####
+    print("Proof verified successfully!")
 
-    zkcircuit_v = zkcircuit.verifier_build(zkbin)
 
-    vprint(f"Making verifying key.....")
-    start = time()
-    verifying_key = VerifyingKey.build(K, zkcircuit_v)
-    print(f"Time for making verifying key: {time() - start}")
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser(
+        prog="zkrunner",
+        description="Python util for running zk proofs",
+        epilog="This tool is only for prototyping purposes",
+    )
+
+    parser.add_argument(
+        "SOURCE",
+        help="Path to zkas source code",
+    )
+    parser.add_argument(
+        "-w",
+        "--witness",
+        required=True,
+        help="Path to JSON file holding witnesses",
+    )
+    parser.add_argument(
+        "--prove",
+        action="store_true",
+        help="Actually create a real proof instead of using MockProver",
+    )
 
-    vprint("Verifying.....")
-    start = time()
-    vprint(f"PUBLIC INPUTS: {pubins}")
-    proof.verify(verifying_key, pubins)
-    print(f"Time for verifying {time() - start}")
+    args = parser.parse_args()
+    main(args.witness, args.SOURCE, mock=not args.prove)