add prover options to export the witness (+ gates + selectors) for third party visualization

Cargo.toml

@@ -11,6 +11,7 @@ log = { version = "0.4.14", default-features = false }
 num = { version = "0.4", default-features = false, features = ["rand"] }
 rand = { version = "0.8.4", default-features = false }
 serde = { version = "1.0", default-features = false, features = ["derive"] }
+serde_json = { version = "1.0" }
 static_assertions = { version = "1.1.0", default-features = false }
 unroll = { version = "0.1.5", default-features = false }
 

plonky2/Cargo.toml

@@ -29,6 +29,7 @@ num = { workspace = true }
 rand = { workspace = true }
 rand_chacha = { version = "0.3.1", optional = true, default-features = false }
 serde = { workspace = true, features = ["rc"] }
+serde_json = { workspace = true }
 static_assertions = { workspace = true }
 unroll = { workspace = true }
 web-time = { version = "1.0.0", optional = true }

plonky2/examples/fibonacci.rs

@@ -4,6 +4,7 @@ use plonky2::iop::witness::{PartialWitness, WitnessWrite};
 use plonky2::plonk::circuit_builder::CircuitBuilder;
 use plonky2::plonk::circuit_data::CircuitConfig;
 use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+use plonky2::plonk::prover::ProverOptions;
 
 /// An example of using Plonky2 to prove a statement of the form
 /// "I know the 100th element of the Fibonacci sequence, starting with constants a and b."
@@ -38,7 +39,12 @@ fn main() -> Result<()> {
     pw.set_target(initial_b, F::ONE)?;
 
     let data = builder.build::<C>();
-    let proof = data.prove(pw)?;
+
+    let prover_opts = ProverOptions {
+        export_witness: Some(String::from("witness.json")),
+    };
+  
+    let proof = data.prove_with_options(pw, &prover_opts)?;
 
     println!(
         "100th Fibonacci number mod |F| (starting with {}, {}) is: {}",

plonky2/src/plonk/circuit_data.rs

@@ -46,7 +46,7 @@ use crate::plonk::circuit_builder::CircuitBuilder;
 use crate::plonk::config::{GenericConfig, Hasher};
 use crate::plonk::plonk_common::PlonkOracle;
 use crate::plonk::proof::{CompressedProofWithPublicInputs, ProofWithPublicInputs};
-use crate::plonk::prover::prove;
+use crate::plonk::prover::{prove,ProverOptions,default_prover_options};
 use crate::plonk::verifier::verify;
 use crate::util::serialization::{
     Buffer, GateSerializer, IoResult, Read, WitnessGeneratorSerializer, Write,
@@ -196,6 +196,17 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
             &self.common,
             inputs,
             &mut TimingTree::default(),
+            &default_prover_options,
+        )
+    }
+
+    pub fn prove_with_options(&self, inputs: PartialWitness<F>, opts: &ProverOptions) -> Result<ProofWithPublicInputs<F, C, D>> {
+        prove::<F, C, D>(
+            &self.prover_only,
+            &self.common,
+            inputs,
+            &mut TimingTree::default(),
+            opts,
         )
     }
 
@@ -294,6 +305,17 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
             &self.common,
             inputs,
             &mut TimingTree::default(),
+            &default_prover_options,
+        )
+    }
+
+    pub fn prove_with_options(&self, inputs: PartialWitness<F>, opts: &ProverOptions) -> Result<ProofWithPublicInputs<F, C, D>> {
+        prove::<F, C, D>(
+            &self.prover_only,
+            &self.common,
+            inputs,
+            &mut TimingTree::default(),
+            opts,
         )
     }
 }

plonky2/src/plonk/prover.rs

@@ -9,6 +9,9 @@ use anyhow::{ensure, Result};
 use hashbrown::HashMap;
 use plonky2_maybe_rayon::*;
 
+use serde::Serialize;
+use crate::util::serialization::json::write_json_file;
+
 use super::circuit_builder::{LookupChallenges, LookupWire};
 use crate::field::extension::Extendable;
 use crate::field::polynomial::{PolynomialCoeffs, PolynomialValues};
@@ -17,7 +20,7 @@ use crate::field::zero_poly_coset::ZeroPolyOnCoset;
 use crate::fri::oracle::PolynomialBatch;
 use crate::gates::lookup::LookupGate;
 use crate::gates::lookup_table::LookupTableGate;
-use crate::gates::selectors::LookupSelectors;
+use crate::gates::selectors::{LookupSelectors};
 use crate::hash::hash_types::RichField;
 use crate::iop::challenger::Challenger;
 use crate::iop::generator::generate_partial_witness;
@@ -111,11 +114,47 @@ pub fn set_lookup_wires<
     Ok(())
 }
 
+//------------------------------------------------------------------------------
+
+// debugging features in the prover
+#[derive(Debug,Clone)]
+pub struct ProverOptions {
+    pub export_witness: Option<String>,      // export the full witness into the given file
+}
+
+pub const default_prover_options: ProverOptions = ProverOptions {
+    export_witness: None,  
+};
+
+// things we want to export to be used by third party tooling
+#[derive(Debug,Clone,Serialize)]
+struct ThingsToExport<F> {
+    gates:           Vec<String>,
+    selector_vector: Vec<usize>,
+    matrix:          Vec<Vec<F>>,
+}
+
+// the idea is to export the witness (plus some more information),
+// so that third-party tools can for example visualize it
+fn collect_things_to_export<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>(
+    common_data:   &CommonCircuitData<F, D>,
+    wires_matrix:  &MatrixWitness<F>,
+) -> ThingsToExport<F> {
+    ThingsToExport {
+        gates:           common_data.gates.iter().map(|g| g.0.short_id()).collect(),
+        selector_vector: common_data.selectors_info.selector_vector.clone(),
+        matrix:          wires_matrix.wire_values.clone(),
+    }
+}
+
+//------------------------------------------------------------------------------
+
 pub fn prove<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>(
     prover_data: &ProverOnlyCircuitData<F, C, D>,
     common_data: &CommonCircuitData<F, D>,
     inputs: PartialWitness<F>,
     timing: &mut TimingTree,
+    prover_options: &ProverOptions,
 ) -> Result<ProofWithPublicInputs<F, C, D>>
 where
     C::Hasher: Hasher<F>,
@@ -127,7 +166,7 @@ where
         generate_partial_witness(inputs, prover_data, common_data)?
     );
 
-    prove_with_partition_witness(prover_data, common_data, partition_witness, timing)
+    prove_with_partition_witness(prover_data, common_data, partition_witness, timing, prover_options)
 }
 
 pub fn prove_with_partition_witness<
@@ -139,6 +178,7 @@ pub fn prove_with_partition_witness<
     common_data: &CommonCircuitData<F, D>,
     mut partition_witness: PartitionWitness<F>,
     timing: &mut TimingTree,
+    prover_options: &ProverOptions,
 ) -> Result<ProofWithPublicInputs<F, C, D>>
 where
     C::Hasher: Hasher<F>,
@@ -184,6 +224,16 @@ where
         )
     );
 
+    // export witness etc for third party tooling
+    match &prover_options.export_witness {
+        None        => (),
+        Some(fname) => {
+            let things_to_export = collect_things_to_export::<F, C, D>( &common_data, &witness );
+            write_json_file(&fname, &things_to_export);
+            println!("exported witness to `{}`",fname);
+        },
+    }
+
     let mut challenger = Challenger::<F, C::Hasher>::new();
 
     // Observe the instance.

plonky2/src/util/serialization/json.rs

@@ -0,0 +1,14 @@
+
+use std::fs::File;
+use std::io::{BufWriter, Write};
+
+use serde::Serialize;
+use serde_json;
+
+pub fn write_json_file<T: Serialize>(fname: &str, what: &T) -> std::io::Result<()> {
+    let file = File::create(fname)?;
+    let mut writer = BufWriter::new(file);
+    serde_json::to_writer(&mut writer, &what)?;
+    writer.flush()?;
+    Ok(())
+}

plonky2/src/util/serialization/mod.rs

@@ -4,6 +4,8 @@ pub mod generator_serialization;
 #[macro_use]
 pub mod gate_serialization;
 
+pub mod json;
+
 #[cfg(not(feature = "std"))]
 use alloc::{collections::BTreeMap, sync::Arc, vec, vec::Vec};
 use core::convert::Infallible;