add module zk::debug, and move export_witness_json() and zkas_type_checks() inside it.

src/error.rs

@@ -197,6 +197,12 @@ pub enum Error {
     #[error("halo2 plonk error: {0}")]
     PlonkError(String),
 
+    #[error("Wrong witness type at index: {0}")]
+    WrongWitnessType(usize),
+
+    #[error("Incorrect public inputs count")]
+    IncorrectPublicInputsCount,
+
     #[error("Unable to decrypt mint note: {0}")]
     NoteDecryptionFailed(String),
 

src/zk/debug.rs

@@ -0,0 +1,105 @@
+use darkfi_sdk::pasta::pallas;
+use log::error;
+
+#[cfg(feature = "tinyjson")]
+use {
+    std::{collections::HashMap, fs::File, io::Write, path::Path},
+    tinyjson::JsonValue::{Array as JsonArray, Object as JsonObj, String as JsonStr},
+};
+
+use super::{Witness, ZkCircuit};
+use crate::{zkas, Error, Result};
+
+#[cfg(feature = "tinyjson")]
+/// Export witness.json which can be used by zkrunner for debugging circuits
+pub fn export_witness_json<P: AsRef<Path>>(
+    output_path: P,
+    prover_witnesses: &Vec<Witness>,
+    public_inputs: &Vec<pallas::Base>,
+) {
+    let mut witnesses = Vec::new();
+    for witness in prover_witnesses {
+        let mut value_json = HashMap::new();
+        match witness {
+            Witness::Base(value) => {
+                value.map(|w1| {
+                    value_json.insert("Base".to_string(), JsonStr(format!("{:?}", w1)));
+                    w1
+                });
+            }
+            Witness::Scalar(value) => {
+                value.map(|w1| {
+                    value_json.insert("Scalar".to_string(), JsonStr(format!("{:?}", w1)));
+                    w1
+                });
+            }
+            _ => unimplemented!(),
+        }
+        witnesses.push(JsonObj(value_json));
+    }
+
+    let mut instances = Vec::new();
+    for instance in public_inputs {
+        instances.push(JsonStr(format!("{:?}", instance)));
+    }
+
+    let witnesses_json = JsonArray(witnesses);
+    let instances_json = JsonArray(instances);
+    let witness_json = JsonObj(HashMap::from([
+        ("witnesses".to_string(), witnesses_json),
+        ("instances".to_string(), instances_json),
+    ]));
+    // This is a debugging method. We don't care about .expect() crashing.
+    let json = witness_json.format().expect("cannot create debug json");
+    let mut output = File::create(output_path).expect("cannot write file");
+    output.write_all(json.as_bytes()).expect("write failed");
+}
+
+/// Call this before `Proof::create()` to perform type checks on the witnesses and check
+/// the amount of provided instances are correct.
+pub fn zkas_type_checks(
+    circuit: &ZkCircuit,
+    binary: &zkas::ZkBinary,
+    instances: &Vec<pallas::Base>,
+) -> Result<()> {
+    for (i, (circuit_witness, binary_witness)) in
+        circuit.witnesses.iter().zip(binary.witnesses.iter()).enumerate()
+    {
+        let is_pass = match circuit_witness {
+            Witness::EcPoint(_) => *binary_witness == zkas::VarType::EcPoint,
+            Witness::EcNiPoint(_) => *binary_witness == zkas::VarType::EcNiPoint,
+            Witness::EcFixedPoint(_) => *binary_witness == zkas::VarType::EcFixedPoint,
+            Witness::Base(_) => *binary_witness == zkas::VarType::Base,
+            Witness::Scalar(_) => *binary_witness == zkas::VarType::Scalar,
+            Witness::MerklePath(_) => *binary_witness == zkas::VarType::MerklePath,
+            Witness::Uint32(_) => *binary_witness == zkas::VarType::Uint32,
+            Witness::Uint64(_) => *binary_witness == zkas::VarType::Uint64,
+        };
+        if !is_pass {
+            error!(
+                "Incorrect witness type at index {}. Expected '{}', instead got '{}'.",
+                i,
+                binary_witness.name(),
+                circuit_witness.name()
+            );
+            return Err(Error::WrongWitnessType(i))
+        }
+    }
+
+    // Count number of public instances
+    let mut instances_count = 0;
+    for opcode in &circuit.opcodes {
+        if let (zkas::Opcode::ConstrainInstance, _) = opcode {
+            instances_count += 1;
+        }
+    }
+    if instances.len() != instances_count {
+        error!(
+            "Wrong number of public inputs. Should be {}, but instead got {}.",
+            instances_count,
+            instances.len()
+        );
+        return Err(Error::IncorrectPublicInputsCount)
+    }
+    Ok(())
+}

src/zk/mod.rs

@@ -35,6 +35,11 @@ pub use proof::{Proof, ProvingKey, VerifyingKey};
 mod tracer;
 pub use tracer::DebugOpValue;
 
+mod debug;
+#[cfg(feature = "tinyjson")]
+pub use debug::export_witness_json;
+pub use debug::zkas_type_checks;
+
 pub mod halo2 {
     pub use halo2_proofs::{
         arithmetic::Field,
@@ -58,55 +63,3 @@ where
         |mut region| region.assign_advice(|| "load private", column, 0, || value),
     )
 }
-
-#[cfg(feature = "tinyjson")]
-use darkfi_sdk::pasta::pallas;
-#[cfg(feature = "tinyjson")]
-use std::{collections::HashMap, fs::File, io::Write, path::Path};
-#[cfg(feature = "tinyjson")]
-use tinyjson::JsonValue::{Array as JsonArray, Object as JsonObj, String as JsonStr};
-
-#[cfg(feature = "tinyjson")]
-/// Export witness.json which can be used by zkrunner for debugging circuits
-pub fn export_witness_json<P: AsRef<Path>>(
-    output_path: P,
-    prover_witnesses: &Vec<Witness>,
-    public_inputs: &Vec<pallas::Base>,
-) {
-    let mut witnesses = Vec::new();
-    for witness in prover_witnesses {
-        let mut value_json = HashMap::new();
-        match witness {
-            Witness::Base(value) => {
-                value.map(|w1| {
-                    value_json.insert("Base".to_string(), JsonStr(format!("{:?}", w1)));
-                    w1
-                });
-            }
-            Witness::Scalar(value) => {
-                value.map(|w1| {
-                    value_json.insert("Scalar".to_string(), JsonStr(format!("{:?}", w1)));
-                    w1
-                });
-            }
-            _ => unimplemented!(),
-        }
-        witnesses.push(JsonObj(value_json));
-    }
-
-    let mut instances = Vec::new();
-    for instance in public_inputs {
-        instances.push(JsonStr(format!("{:?}", instance)));
-    }
-
-    let witnesses_json = JsonArray(witnesses);
-    let instances_json = JsonArray(instances);
-    let witness_json = JsonObj(HashMap::from([
-        ("witnesses".to_string(), witnesses_json),
-        ("instances".to_string(), instances_json),
-    ]));
-    // This is a debugging method. We don't care about .expect() crashing.
-    let json = witness_json.format().expect("cannot create debug json");
-    let mut output = File::create(output_path).expect("cannot write file");
-    output.write_all(json.as_bytes()).expect("write failed");
-}

src/zk/vm.rs

@@ -65,7 +65,6 @@ use super::{
     tracer::ZkTracer,
 };
 use crate::zkas::{
-    self,
     types::{HeapType, LitType},
     Opcode, ZkBinary,
 };
@@ -264,9 +263,9 @@ pub struct ZkParams {
 #[derive(Clone)]
 pub struct ZkCircuit {
     constants: Vec<String>,
-    witnesses: Vec<Witness>,
+    pub(super) witnesses: Vec<Witness>,
     literals: Vec<(LitType, String)>,
-    opcodes: Vec<(Opcode, Vec<(HeapType, usize)>)>,
+    pub(super) opcodes: Vec<(Opcode, Vec<(HeapType, usize)>)>,
     pub tracer: ZkTracer,
 }
 
@@ -274,51 +273,18 @@ impl ZkCircuit {
     pub fn new(witnesses: Vec<Witness>, circuit_code: &ZkBinary) -> Self {
         let constants = circuit_code.constants.iter().map(|x| x.1.clone()).collect();
         let literals = circuit_code.literals.clone();
-        let self_ = Self {
+        Self {
             constants,
             witnesses,
             literals,
             opcodes: circuit_code.opcodes.clone(),
             tracer: ZkTracer::new(true),
-        };
-        self_.check_witness_types(&circuit_code.witnesses);
-        self_
+        }
     }
 
     pub fn enable_trace(&mut self) {
         self.tracer.init();
     }
-
-    // Temporary, should be moved into Proof::create() but then we need to store binary_witnesses
-    // inside ZkCircuit.
-    fn check_witness_types(&self, binary_witnesses: &Vec<zkas::VarType>) {
-        for (circuit_witness, binary_witness) in self.witnesses.iter().zip(binary_witnesses.iter())
-        {
-            let is_pass = match circuit_witness {
-                Witness::EcPoint(_) => *binary_witness == zkas::VarType::EcPoint,
-                Witness::EcNiPoint(_) => *binary_witness == zkas::VarType::EcNiPoint,
-                Witness::EcFixedPoint(_) => *binary_witness == zkas::VarType::EcFixedPoint,
-                Witness::Base(_) => *binary_witness == zkas::VarType::Base,
-                Witness::Scalar(_) => *binary_witness == zkas::VarType::Scalar,
-                Witness::MerklePath(_) => *binary_witness == zkas::VarType::MerklePath,
-                Witness::Uint32(_) => *binary_witness == zkas::VarType::Uint32,
-                Witness::Uint64(_) => *binary_witness == zkas::VarType::Uint64,
-            };
-            if is_pass {
-                // return Err(Error::IncorrectWitnessType)
-                panic!("incorrect type passed in");
-            }
-        }
-
-        // Count number of public instances
-        let mut instances_count = 0;
-        for opcode in &self.opcodes {
-            if let (Opcode::ConstrainInstance, _) = opcode {
-                instances_count += 1;
-            }
-        }
-        // if instances.len() != instances_count { ... }
-    }
 }
 
 impl Circuit<pallas::Base> for ZkCircuit {

src/zk/vm_heap.rs

@@ -48,6 +48,21 @@ pub enum Witness {
     Uint64(Value<u64>),
 }
 
+impl Witness {
+    pub fn name(&self) -> &str {
+        match self {
+            Self::EcPoint(_) => "EcPoint",
+            Self::EcNiPoint(_) => "EcNiPoint",
+            Self::EcFixedPoint(_) => "EcFixedPoint",
+            Self::Base(_) => "Base",
+            Self::Scalar(_) => "Scalar",
+            Self::MerklePath(_) => "MerklePath",
+            Self::Uint32(_) => "Uint32",
+            Self::Uint64(_) => "Uint64",
+        }
+    }
+}
+
 pub enum Literal {
     Uint64(Value<u64>),
 }

src/zkas/types.rs

@@ -100,6 +100,25 @@ impl VarType {
             _ => None,
         }
     }
+
+    pub fn name(&self) -> &str {
+        match self {
+            Self::Dummy => "Dummy",
+            Self::EcPoint => "EcPoint",
+            Self::EcFixedPoint => "EcFixedPoint",
+            Self::EcFixedPointShort => "EcFixedPointShort",
+            Self::EcFixedPointBase => "EcFixedPointBase",
+            Self::EcNiPoint => "EcNiPoint",
+            Self::Base => "Base",
+            Self::BaseArray => "BaseArray",
+            Self::Scalar => "Scalar",
+            Self::ScalarArray => "ScalarArray",
+            Self::MerklePath => "MerklePath",
+            Self::Uint32 => "Uint32",
+            Self::Uint64 => "Uint64",
+            Self::Any => "Any",
+        }
+    }
 }
 
 /// Literal types supported by the zkas VM