create zkrender tool to plot circuit layouts

Cargo.lock

@@ -1993,6 +1993,8 @@ dependencies = [
  "darkfi",
  "darkfi-sdk",
  "halo2_gadgets",
+ "halo2_proofs",
+ "plotters",
  "pyo3",
  "rand 0.8.5",
 ]

bin/zkrunner/zkrender.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+# This file is part of DarkFi (https://dark.fi)
+#
+# Copyright (C) 2020-2024 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 render zkVM circuit layouts given zkas source code.
+"""
+import sys
+from darkfi_sdk.zkas import ZkBinary, ZkCircuit
+from zkrunner import eprint, load_circuit_witness
+
+def main(witness_file, source_file, output, width, height, font_size):
+    print("Compiling zkas code...")
+    with open(source_file, "r", encoding="utf-8") as zkas_file:
+        zkas_source = zkas_file.read()
+
+    zkbin = ZkBinary(source_file, zkas_source)
+    circuit = ZkCircuit(zkbin)
+    print("Decoding witnesses...")
+    load_circuit_witness(circuit, witness_file)
+    circuit = circuit.prover_build()
+
+    if not circuit.render(zkbin.k(), output, width, height, font_size):
+        eprint("Rendering failed")
+    print(f"Written output to '{output}'")
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser(
+        prog="zkrender",
+        description="Python util for rendering zk circuits",
+        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(
+        "OUTPUT",
+        help="Path to output image",
+    )
+
+    parser.add_argument(
+        "-W", "--width", type=int,
+        default=800,
+        help="Image width",
+    )
+    parser.add_argument(
+        "-H", "--height", type=int,
+        default=600,
+        help="Image width",
+    )
+    parser.add_argument(
+        "-f", "--font-size", type=int,
+        default=20,
+        help="Image width",
+    )
+
+    args = parser.parse_args()
+    sys.exit(main(args.witness, args.SOURCE, args.OUTPUT,
+                  args.width, args.height, args.font_size))
+

bin/zkrunner/zkrunner.py

@@ -40,45 +40,16 @@ def show_trace(opcodes, trace):
         optype = str(optype)
         print(f"{i:<4} {opcode:<22} {optype:<10} {args}")
 
-def main(witness_file, source_file, mock=False, trace=False):
-    """main zkrunner logic"""
-    # We will first attempt to decode the witnesses from the JSON file.
+def load_circuit_witness(circuit, witness_file):
+    # We 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...")
     if witness_file == "-":
         witness_data = json.load(sys.stdin)
     else:
         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.
@@ -113,6 +84,44 @@ def main(witness_file, source_file, mock=False, trace=False):
             eprint(f"Invalid Witness type for witness {witness}")
             return -1
 
+    # 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))
+    return instances
+
+def main(witness_file, source_file, mock=False, trace=False):
+    """main zkrunner logic"""
+    # 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)
+    print("Decoding witnesses...")
+    instances = load_circuit_witness(circuit, witness_file)
+
+    # 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())
+
     # circuit.prover_build() will actually construct the circuit
     # with the values witnessed above.
     circuit = circuit.prover_build()
@@ -122,12 +131,6 @@ def main(witness_file, source_file, mock=False, trace=False):
             return -2
         circuit.enable_trace()
 
-    # 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))
-
     # If we're building an actual proof, we'll use the ProvingKey to
     # prove and our VerifyingKey to verify the proof.
     if not mock:

src/sdk/python/Cargo.toml

@@ -16,6 +16,8 @@ doc = false
 [dependencies]
 darkfi = {path = "../../../", features = ["zk", "zkas"]}
 darkfi-sdk = {path = "../"}
+halo2_proofs = {version = "0.3.0", features = ["dev-graph", "sanity-checks"]}
 halo2_gadgets = "0.3.0"
+plotters = "0.3.5"
 pyo3 = "0.20.3"
 rand = "0.8.5"

src/sdk/python/src/zkas.rs

@@ -189,6 +189,24 @@ impl ZkCircuit {
         }
         result
     }
+
+    fn render(&self, k: u32, filename: &str, width: u32, height: u32, font_size: u32) -> bool {
+        use plotters::prelude::*;
+        let root = BitMapBackend::new(filename, (width, height)).into_drawing_area();
+        if let Err(_) = root.fill(&WHITE) {
+            return false
+        }
+        let root = if let Ok(root) = root.titled("circuit", ("sans-serif", font_size)) {
+            root
+        } else {
+            return false
+        };
+
+        match halo2_proofs::dev::CircuitLayout::default().render(k, &self.0, &root) {
+            Ok(()) => true,
+            Err(_) => false,
+        }
+    }
 }
 
 #[pyclass]