Panic when trying to generate proof with invalid witness

[Gorzorg]

When testing circuits written with Plonky2, one might want to see what happens when invalid witnesses are used. Unfortunately, it is very easy to trigger panics when this kind of tests is performed.

As far as I can tell, the culprit is an assert_eq! in the implementation of plonky2::iop::witness::WitnessWrite for plonky2::iop::witness::PartitionWitness.

The panic from that assert_eq! is triggered by the witness generation process in plonky2::iop::generator::generate_partial_witness.

It would be nice if the process didn't panic, and as far as I could see, an error propagation mechanism would be feasible, if one allowed changes in the signature of WitnessWrite::set_target so that the function returns a Result.

The obvious downside of this solution is that it would technically be a breaking change that affects all downstream crates that implement generators, given that calls <plonky2::iop::generator::GeneratedValues as WitnessWrite>::set_target are required for the implementation.

I say technically a breaking change because, since WitnessWrite::set_target currently returns (), there is no real reason to feed its output into the computation of other expressions, and therefore most downstream dependencies would only notice a bunch of new warnings on their code base. Still, that would be a breaking change and I am in no position to evaluate if it is beneficial.

---

In the meantime I had an idea to circumvent the problem in test environment. It's more of a hack than anything, but it seems to work.

The idea is to implement a super simple gate that only imposes the constraint vars.local_wires[0] - vars.local_wires[1] so that it is in a way logically equivalent to CircuitBuilder::connect. This gate does not implement any generators, so that the assert_eq! that causes the panic is not called. One can use that gate in tests instead of connect, so that the proof generator generates invalid proofs instead of raising panics.

Does it make sense as a workaround? If not, why?

Here follows my implementation of the gate in question. If the idea sketched above does not make sense, please ignore it. If it makes sense, is it implemented correctly?

use plonky2::{
    field::extension::Extendable,
    gates::gate::Gate,
    hash::hash_types::RichField,
    iop::ext_target::ExtensionTarget,
    plonk::{
        circuit_builder::CircuitBuilder,
        circuit_data::CommonCircuitData,
        vars::{EvaluationTargets, EvaluationVars},
    },
    util::serialization::{Buffer, IoResult},
};

#[cfg(test)]
use plonky2::{
    field::types::Field,
    iop::{
        target::Target,
        witness::{PartialWitness, WitnessWrite},
    },
    plonk::{
        circuit_data::CircuitConfig,
        config::{GenericConfig, PoseidonGoldilocksConfig},
    },
};

#[derive(Debug)]
struct TestEq;

impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for TestEq {
    fn id(&self) -> String {
        format!("{self:?}")
    }

    fn serialize(
        &self,
        _dst: &mut Vec<u8>,
        _common_data: &CommonCircuitData<F, D>,
    ) -> IoResult<()> {
        Ok(())
    }

    fn num_wires(&self) -> usize {
        2
    }

    fn num_constants(&self) -> usize {
        0
    }

    fn num_constraints(&self) -> usize {
        1
    }

    fn degree(&self) -> usize {
        1
    }

    fn deserialize(_src: &mut Buffer, _common_data: &CommonCircuitData<F, D>) -> IoResult<Self>
    where
        Self: Sized,
    {
        Ok(Self)
    }

    fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<<F as Extendable<D>>::Extension> {
        vec![vars.local_wires[0] - vars.local_wires[1]]
    }

    fn eval_unfiltered_circuit(
        &self,
        builder: &mut CircuitBuilder<F, D>,
        vars: EvaluationTargets<D>,
    ) -> Vec<ExtensionTarget<D>> {
        vec![builder.sub_extension(vars.local_wires[0], vars.local_wires[1])]
    }

    fn generators(
        &self,
        _row: usize,
        _local_constants: &[F],
    ) -> Vec<plonky2::iop::generator::WitnessGeneratorRef<F, D>> {
        vec![]
    }
}

#[test]
fn test_test_eq_gate() {
    const D: usize = 2;
    type BaseField = <PoseidonGoldilocksConfig as GenericConfig<D>>::F;

    let mut builder =
        CircuitBuilder::<BaseField, D>::new(CircuitConfig::standard_recursion_zk_config());

    let gate_row = builder.add_gate(TestEq, vec![]);

    let circuit = builder.build::<PoseidonGoldilocksConfig>();

    let mut correct_witness = PartialWitness::new();

    correct_witness.set_target(Target::wire(gate_row, 0), BaseField::ONE);

    let mut incorrect_witness = correct_witness.clone();

    incorrect_witness.set_target(Target::wire(gate_row, 1), BaseField::ZERO);
    correct_witness.set_target(Target::wire(gate_row, 1), BaseField::ONE);

    let incorrect_proof = circuit.prove(incorrect_witness);
    let correct_proof = circuit.prove(correct_witness);

    match incorrect_proof {
        Ok(proof) => match circuit.verify(proof) {
            Err(err) => {
                println!("The invalid proof is correctly identified as invalid because\n{err:?}");
            }
            _ => panic!("invalid proof is accepted!"),
        },
        _ => panic!("proof generation with invalid witness fails."),
    }

    match correct_proof {
        Ok(proof) => match circuit.verify(proof) {
            Ok(()) => {
                println!("The valid proof is correctly identified as valid")
            }
            _ => panic!("valid proof is rejected!"),
        },
        _ => panic!("proof generation with valid witness fails"),
    }
}