zkrunner/pydrk: add --trace argument to zkrunner, and appropriate changes to the py bindings.

bin/zkrunner/zkrunner.py

@@ -25,8 +25,22 @@ from darkfi_sdk.pasta import Fp, Fq, Ep
 from darkfi_sdk.zkas import (MockProver, ZkBinary, ZkCircuit, ProvingKey,
                              Proof, VerifyingKey)
 
+def eprint(fstr, *args):
+    print("error: " + fstr, *args, file=sys.stderr)
+
+def show_trace(opcodes, trace):
+    print(f"{'Line':<4} {'Opcode':<22} {'Type':<10} {'Values'}")
+    for i, (opcode, (optype, args)) in enumerate(zip(opcodes, trace)):
+        if args:
+            args = ", ".join([str(arg) for arg in args])
+            args = f"[{args}]"
+        else:
+            args = ""
+        opcode = str(opcode)
+        optype = str(optype)
+        print(f"{i:<4} {opcode:<22} {optype:<10} {args}")
 
-def main(witness_file, source_file, mock=False):
+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.
     # Refer to the `witness_gen.py` file to see what the format of this
@@ -96,11 +110,17 @@ def main(witness_file, source_file, mock=False):
             circuit.witness_uint64(value)
 
         else:
-            raise ValueError(f"Invalid Witness type for witness {witness}")
+            eprint(f"Invalid Witness type for witness {witness}")
+            return -1
 
     # circuit.prover_build() will actually construct the circuit
     # with the values witnessed above.
     circuit = circuit.prover_build()
+    if trace:
+        if mock:
+            eprint(f"Debug trace can only be enabled with --prove")
+            return -2
+        circuit.enable_trace()
 
     # Instances are our public inputs for the proof and they're also
     # part of the JSON file.
@@ -114,6 +134,9 @@ def main(witness_file, source_file, mock=False):
         print("Proving knowledge of witnesses...")
         proof = Proof.create(proving_key, [circuit], instances)
 
+        if trace:
+            show_trace(zkbin.opcodes(), circuit.opvalues())
+
         print("Verifying ZK proof...")
         proof.verify(verifying_key, instances)
 
@@ -121,11 +144,12 @@ def main(witness_file, source_file, mock=False):
     else:
         print("Running MockProver...")
         proof = MockProver.run(zkbin.k(), circuit, instances)
+
         print("Verifying MockProver...")
         proof.verify()
 
     print("Proof verified successfully!")
-
+    return 0
 
 if __name__ == "__main__":
     from argparse import ArgumentParser
@@ -151,6 +175,13 @@ if __name__ == "__main__":
         action="store_true",
         help="Actually create a real proof instead of using MockProver",
     )
+    parser.add_argument(
+        "--trace",
+        action="store_true",
+        help="Enable debug trace (only works with --prove enabled)",
+    )
 
     args = parser.parse_args()
-    main(args.witness, args.SOURCE, mock=not args.prove)
+    sys.exit(main(args.witness, args.SOURCE, mock=not args.prove,
+                  trace=args.trace))
+

src/sdk/python/src/zkas.rs

@@ -28,6 +28,16 @@ use rand::rngs::OsRng;
 
 use super::pasta::{Ep, Fp, Fq};
 
+#[pyclass]
+pub struct ZkOpcode(zkas::Opcode);
+
+#[pymethods]
+impl ZkOpcode {
+    fn __str__(&self) -> PyResult<String> {
+        Ok(self.0.name().to_string())
+    }
+}
+
 #[pyclass]
 /// Decoded zkas bincode
 pub struct ZkBinary(decoder::ZkBinary);
@@ -71,6 +81,29 @@ impl ZkBinary {
     fn k(&self) -> u32 {
         self.0.k
     }
+
+    fn opcodes(&self) -> Vec<ZkOpcode> {
+        return self.0.opcodes.iter().map(|op| ZkOpcode(op.0)).collect()
+    }
+}
+
+#[pyclass]
+enum DebugOpValue {
+    EcPoint,
+    Base,
+    Void,
+}
+
+#[pymethods]
+impl DebugOpValue {
+    fn __str__(&self) -> PyResult<String> {
+        let name = match self {
+            DebugOpValue::EcPoint => "EcPoint",
+            DebugOpValue::Base => "Base",
+            DebugOpValue::Void => "Void",
+        };
+        Ok(name.to_string())
+    }
 }
 
 #[pyclass]
@@ -136,6 +169,26 @@ impl ZkCircuit {
     fn witness_uint64(&mut self, w: u64) {
         self.1.push(zk::vm::Witness::Uint64(Value::known(w)));
     }
+
+    fn enable_trace(&mut self) {
+        self.0.enable_trace();
+    }
+
+    fn opvalues(&self) -> Vec<(DebugOpValue, Vec<Fp>)> {
+        let opvalue_binding = self.0.tracer.opvalues.borrow();
+        let opvalues = opvalue_binding.as_ref().unwrap();
+        let mut result = Vec::new();
+        for opvalue in opvalues {
+            match opvalue {
+                zk::DebugOpValue::EcPoint(x, y) => {
+                    result.push((DebugOpValue::EcPoint, vec![Fp(*x), Fp(*y)]))
+                }
+                zk::DebugOpValue::Base(v) => result.push((DebugOpValue::Base, vec![Fp(*v)])),
+                zk::DebugOpValue::Void => result.push((DebugOpValue::Void, vec![])),
+            }
+        }
+        result
+    }
 }
 
 #[pyclass]
@@ -181,13 +234,50 @@ impl Proof {
         instances: Vec<&PyCell<Fp>>,
     ) -> Self {
         let pk = pk.borrow().deref().0.clone();
-        let circuits: Vec<zk::vm::ZkCircuit> =
-            circuits.iter().map(|c| c.borrow().deref().0.clone()).collect();
+
+        // Ugh this is so annoying. The halo2 API expects &[] of values.
+        // We carefully unpack the current Vec, then replace its contents back again.
+        // I've left the old code below to see what we did before.
+        //
+        //   let circuits: Vec<zk::vm::ZkCircuit> =
+        //       circuits.iter().map(|c| c.borrow().deref().0.clone()).collect();
+        //
+        // The alternative is to make your own container as documented here:
+        // https://pyo3.rs/v0.19.2/class/protocols.html?highlight=__getitem__#mapping--sequence-types
+        let zkbin = decoder::ZkBinary {
+            namespace: "".to_string(),
+            k: 0,
+            constants: Vec::new(),
+            literals: Vec::new(),
+            witnesses: Vec::new(),
+            opcodes: Vec::new(),
+        };
+        let empty_circuit = zk::vm::ZkCircuit::new(Vec::new(), &zkbin);
+        let curr_circuits: Vec<ZkCircuit> = circuits
+            .iter()
+            .map(|c| c.replace(ZkCircuit(empty_circuit.clone(), Vec::new(), zkbin.clone())))
+            .collect();
+
+        let mut ucircuits = Vec::new();
+        let mut other_stuff = Vec::new();
+        for circ in curr_circuits.into_iter() {
+            ucircuits.push(circ.0);
+            other_stuff.push((circ.1, circ.2));
+        }
+        //////////////
+
         let instances: Vec<pallas::Base> = instances.iter().map(|i| i.borrow().deref().0).collect();
 
         let proof =
-            zk::proof::Proof::create(&pk, circuits.as_slice(), instances.as_slice(), &mut OsRng)
+            zk::proof::Proof::create(&pk, ucircuits.as_slice(), instances.as_slice(), &mut OsRng)
                 .unwrap();
+
+        // Now replace the "stuff" back again
+        for (old_circ, (circ, stuff)) in
+            circuits.iter().zip(ucircuits.into_iter().zip(other_stuff.into_iter()))
+        {
+            old_circ.replace(ZkCircuit(circ, stuff.0, stuff.1));
+        }
         Self(proof)
     }
 

src/zk/mod.rs

@@ -33,6 +33,7 @@ pub use proof::{Proof, ProvingKey, VerifyingKey};
 
 /// Trace computation of intermediate values in circuit
 mod tracer;
+pub use tracer::DebugOpValue;
 
 pub mod halo2 {
     pub use halo2_proofs::{

src/zk/tracer.rs

@@ -1,7 +1,4 @@
-use std::{
-    cell::RefCell,
-    ops::{Deref, DerefMut},
-};
+use std::cell::RefCell;
 
 use darkfi_sdk::{crypto::constants::OrchardFixedBases, pasta::pallas};
 use halo2_gadgets::ecc as ecc_gadget;

src/zkas/opcode.rs

@@ -156,6 +156,35 @@ impl Opcode {
         }
     }
 
+    pub fn name(&self) -> &str {
+        match self {
+            Self::Noop => "noop",
+            Self::EcAdd => "ec_add",
+            Self::EcMul => "ec_mul",
+            Self::EcMulBase => "ec_mul_base",
+            Self::EcMulShort => "ec_mul_short",
+            Self::EcMulVarBase => "ec_mul_var_base",
+            Self::EcGetX => "ec_get_x",
+            Self::EcGetY => "ec_get_y",
+            Self::PoseidonHash => "poseidon_hash",
+            Self::MerkleRoot => "merkle_root",
+            Self::BaseAdd => "base_add",
+            Self::BaseMul => "base_mul",
+            Self::BaseSub => "base_sub",
+            Self::WitnessBase => "witness_base",
+            Self::RangeCheck => "range_check",
+            Self::LessThanStrict => "less_than_strict",
+            Self::LessThanLoose => "less_than_loose",
+            Self::BoolCheck => "bool_check",
+            Self::CondSelect => "cond_select",
+            Self::ZeroCondSelect => "zero_cond",
+            Self::ConstrainEqualBase => "constrain_equal_base",
+            Self::ConstrainEqualPoint => "constrain_equal_point",
+            Self::ConstrainInstance => "constrain_instance",
+            Self::DebugPrint => "debug",
+        }
+    }
+
     /// Return a tuple of vectors of types that are accepted by a specific opcode.
     /// `r.0` is the return type(s), and `r.1` is the argument type(s).
     pub fn arg_types(&self) -> (Vec<VarType>, Vec<VarType>) {