A0-4073: Add a basic circuit runner

circuit_runner/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "circuit_runner"
+version = "0.0.0"
+edition = "2021"
+authors = ["Cardinal"]
+documentation = "https://docs.rs/?"
+homepage = "https://alephzero.org"
+license = "Apache-2.0"
+categories = ["cryptography"]
+keywords = ["cryptography", "snark", "zero-knowledge", "liminal", "shielder"]
+repository = "https://github.com/Cardinal-Cryptography/zk-apps"
+description = "A collection of halo2-based relations for use in liminal."
+
+[dependencies]
+rand = "=0.8"
+serde = { version = "=1.0", default-features = false, features = ["derive"] }
+serde_json = "=1.0"
+
+halo2-base =  { git = "https://github.com/axiom-crypto/halo2-lib", branch = "community-edition" }
+snark-verifier-sdk = { git = "https://github.com/axiom-crypto/snark-verifier.git", branch = "community-edition" }
+
+# For measuring performance
+ark-std = { version = "=0.4.0", features = ["print-trace"] }

circuit_runner/examples/square.rs

@@ -0,0 +1,82 @@
+use circuit_runner::{mock_run, run, CircuitGenerator};
+use halo2_base::{
+    gates::{circuit::builder::BaseCircuitBuilder, GateChip, GateInstructions},
+    halo2_proofs::halo2curves::{bn256::Fr, ff::PrimeField},
+    utils::BigPrimeField,
+    AssignedValue,
+};
+
+#[derive(Clone, Debug)]
+struct WitnessInput<F: BigPrimeField> {
+    x: F,
+}
+
+impl<F: BigPrimeField> Default for WitnessInput<F> {
+    fn default() -> Self {
+        WitnessInput { x: F::from_u128(0) }
+    }
+}
+
+struct SquareCircuit<F: BigPrimeField> {
+    _marker: std::marker::PhantomData<F>,
+}
+
+impl<F: BigPrimeField> Default for SquareCircuit<F> {
+    fn default() -> Self {
+        SquareCircuit {
+            _marker: std::marker::PhantomData,
+        }
+    }
+}
+
+// Checks that x*x = y, for x being the witness, and y the instance
+impl<F: BigPrimeField> CircuitGenerator<F> for SquareCircuit<F> {
+    type WitnessInput = WitnessInput<F>;
+
+    fn run(
+        &self,
+        builder: &mut BaseCircuitBuilder<F>,
+        input: Self::WitnessInput,
+        make_public: &mut Vec<AssignedValue<F>>,
+    ) {
+        let ctx = builder.main(0);
+        let cx = ctx.load_witness(input.x);
+        let gate = GateChip::<F>::default();
+        let cy = gate.mul(ctx, cx, cx);
+        make_public.push(cy);
+    }
+}
+
+fn mock_run_circuit() {
+    let k_srs = 6;
+    let k_circuit = 5;
+    let cg = SquareCircuit::<Fr>::default();
+    let input = WitnessInput {
+        x: Fr::from_u128(2),
+    };
+
+    let instance = vec![Fr::from_u128(4)];
+    println!("Running on correct data, this should succeed:");
+    mock_run(&cg, input.clone(), instance, k_circuit, k_srs);
+
+    println!("Running on incorrect data, this should panic:");
+    let instance = vec![Fr::from_u128(5)];
+    mock_run(&cg, input, instance, k_circuit, k_srs);
+}
+
+fn run_circuit() {
+    let k_srs = 6;
+    let k_circuit = 5;
+    let cg = SquareCircuit::<Fr>::default();
+    let input = WitnessInput {
+        x: Fr::from_u128(2),
+    };
+
+    let instance = vec![Fr::from_u128(4)];
+    run(&cg, input.clone(), instance, k_circuit, k_srs);
+}
+
+fn main() {
+    run_circuit();
+    mock_run_circuit();
+}

circuit_runner/rust-toolchain.toml

@@ -0,0 +1,2 @@
+[toolchain]
+channel = "nightly-2023-12-21"

circuit_runner/src/lib.rs

@@ -0,0 +1,245 @@
+use ark_std::{end_timer, start_timer};
+use halo2_base::{
+    gates::circuit::{builder::BaseCircuitBuilder, BaseCircuitParams, CircuitBuilderStage},
+    halo2_proofs::{
+        dev::MockProver,
+        halo2curves::bn256::{Bn256, Fr, G1Affine},
+        plonk::{keygen_pk, keygen_vk, verify_proof, Circuit, ProvingKey, VerifyingKey},
+        poly::{
+            commitment::ParamsProver,
+            kzg::{
+                commitment::{KZGCommitmentScheme, ParamsKZG},
+                multiopen::VerifierSHPLONK,
+                strategy::SingleStrategy,
+            },
+        },
+    },
+    utils::BigPrimeField,
+    AssignedValue,
+};
+use rand::{rngs::StdRng, SeedableRng};
+use snark_verifier_sdk::{
+    halo2::{gen_snark_shplonk, PoseidonTranscript},
+    NativeLoader, Snark,
+};
+
+pub fn gen_srs(k: u32) -> ParamsKZG<Bn256> {
+    ParamsKZG::<Bn256>::setup(k, StdRng::from_seed(Default::default()))
+}
+
+/// An interface for circuit generator using halo2-lib
+/// See the examples in the `circuit_runner/examples` directory for usage
+pub trait CircuitGenerator<F: BigPrimeField> {
+    type WitnessInput: Default + Clone;
+    fn run(
+        &self,
+        builder: &mut BaseCircuitBuilder<F>,
+        input: Self::WitnessInput,
+        make_public: &mut Vec<AssignedValue<F>>,
+    );
+}
+
+fn create_circuit<F: BigPrimeField, CG: CircuitGenerator<F>>(
+    cg: &CG,
+    stage: CircuitBuilderStage,
+    lookup_bits: Option<usize>,
+    // It seems to me this `k` does not need to be the same as `k` from srs, but it should not be larger
+    // Supposedly this is max number of rows the halo2-lib circuit generator will use in a single column
+    // If the circuit is larger than this, it will be split into multiple columns
+    k_circuit: usize,
+    input: CG::WitnessInput,
+) -> BaseCircuitBuilder<F> {
+    assert!(!matches!(stage, CircuitBuilderStage::Prover));
+    let mut builder = BaseCircuitBuilder::from_stage(stage);
+    builder.set_k(k_circuit);
+    if let Some(lookup_bits) = lookup_bits {
+        builder.set_lookup_bits(lookup_bits);
+    }
+    builder.set_instance_columns(1);
+
+    let mut assigned_instances = vec![];
+    cg.run(&mut builder, input, &mut assigned_instances);
+    // Not sure what's the point of the below code
+    if !assigned_instances.is_empty() {
+        assert_eq!(
+            builder.assigned_instances.len(),
+            1,
+            "num_instance_columns != 1"
+        );
+        builder.assigned_instances[0] = assigned_instances;
+    }
+    if !stage.witness_gen_only() {
+        // now `builder` contains the execution trace, and we are ready to actually create the circuit
+        // minimum rows is the number of rows used for blinding factors. This depends on the circuit itself, but we can guess the number and change it if something breaks (default 9 usually works)
+        let minimum_rows = 20;
+        builder.calculate_params(Some(minimum_rows));
+    }
+    builder
+}
+
+pub fn create_circuit_prover<F: BigPrimeField, CG: CircuitGenerator<F>>(
+    cg: &CG,
+    params: BaseCircuitParams,
+    break_points: Vec<Vec<usize>>,
+    input: CG::WitnessInput,
+) -> BaseCircuitBuilder<F> {
+    let mut builder = BaseCircuitBuilder::from_stage(CircuitBuilderStage::Prover);
+    builder.set_params(params);
+    builder.set_break_points(break_points);
+    builder.set_instance_columns(1);
+    let mut assigned_instances = vec![];
+    cg.run(&mut builder, input, &mut assigned_instances);
+    // Not sure what's the point of the below code
+    if !assigned_instances.is_empty() {
+        assert_eq!(
+            builder.assigned_instances.len(),
+            1,
+            "num_instance_columns != 1"
+        );
+        builder.assigned_instances[0] = assigned_instances;
+    }
+    builder
+}
+
+pub fn create_circuit_keygen<F: BigPrimeField, CG: CircuitGenerator<F>>(
+    cg: &CG,
+    lookup_bits: Option<usize>,
+    k_circuit: usize,
+) -> BaseCircuitBuilder<F> {
+    create_circuit(
+        cg,
+        CircuitBuilderStage::Keygen,
+        lookup_bits,
+        k_circuit,
+        CG::WitnessInput::default(),
+    )
+}
+
+pub fn create_circuit_mock<F: BigPrimeField, CG: CircuitGenerator<F>>(
+    cg: &CG,
+    lookup_bits: Option<usize>,
+    k_circuit: usize,
+    input: CG::WitnessInput,
+) -> BaseCircuitBuilder<F> {
+    create_circuit(cg, CircuitBuilderStage::Mock, lookup_bits, k_circuit, input)
+}
+
+#[derive(Clone)]
+struct KeygenResult {
+    params: ParamsKZG<Bn256>,
+    pk: ProvingKey<G1Affine>,
+    c_params: BaseCircuitParams,
+    break_points: Vec<Vec<usize>>,
+}
+
+fn keygen<CG: CircuitGenerator<Fr>>(cg: &CG, k_circuit: usize, k_srs: u32) -> KeygenResult {
+    assert!(k_circuit <= k_srs as usize, "k_circuit > k_srs");
+
+    let srs_time = start_timer!(|| "Generating srs");
+    let params = gen_srs(k_srs);
+    end_timer!(srs_time);
+
+    let generate_time = start_timer!(|| "Generating circuit");
+    let circuit = create_circuit_keygen(cg, None, k_circuit);
+    end_timer!(generate_time);
+
+    let vk_time = start_timer!(|| "Generating vk");
+    let vk = keygen_vk(&params, &circuit).unwrap();
+    end_timer!(vk_time);
+
+    let pk_time = start_timer!(|| "Generating pk");
+    let pk = keygen_pk(&params, vk.clone(), &circuit).unwrap();
+    end_timer!(pk_time);
+
+    let c_params = circuit.params();
+    let break_points = circuit.break_points();
+
+    KeygenResult {
+        params,
+        pk,
+        c_params,
+        break_points,
+    }
+}
+
+fn generate_snark<CG: CircuitGenerator<Fr>>(
+    cg: &CG,
+    kg: KeygenResult,
+    input: CG::WitnessInput,
+) -> Snark {
+    let KeygenResult {
+        params,
+        pk,
+        c_params,
+        break_points,
+        ..
+    } = kg;
+    let circuit = create_circuit_prover(cg, c_params, break_points, input);
+    gen_snark_shplonk(&params, &pk, circuit, None::<String>)
+}
+
+fn verify_snark(
+    vk: VerifyingKey<G1Affine>,
+    snark: Snark,
+    params: ParamsKZG<Bn256>,
+    instance: Vec<Fr>,
+) -> bool {
+    let verifier_params = params.verifier_params();
+    let strategy = SingleStrategy::new(&params);
+    let mut transcript = PoseidonTranscript::<NativeLoader, &[u8]>::new::<0>(&snark.proof[..]);
+    let instance = &instance;
+    let verify_time = start_timer!(|| "Verifying proof");
+    let res = verify_proof::<
+        KZGCommitmentScheme<Bn256>,
+        VerifierSHPLONK<'_, Bn256>,
+        _,
+        _,
+        SingleStrategy<'_, Bn256>,
+    >(
+        verifier_params,
+        &vk,
+        strategy,
+        &[&[instance]],
+        &mut transcript,
+    );
+    end_timer!(verify_time);
+    res.is_ok()
+}
+
+/// k_srs -- 2^{k_srs} is the size of the srs
+/// k_circuit -- 2^{k_circuit} is the max number of rows the halo2-lib circuit generator will use in a single column
+pub fn mock_run<F: BigPrimeField, CG: CircuitGenerator<F>>(
+    cg: &CG,
+    input: CG::WitnessInput,
+    instance: Vec<F>,
+    k_circuit: usize,
+    k_srs: u32,
+) {
+    let circuit = create_circuit_mock(cg, None, k_circuit, input);
+    let expected_instance_len = circuit.assigned_instances[0].len();
+    assert_eq!(instance.len(), expected_instance_len);
+    MockProver::run(k_srs, &circuit, vec![instance])
+        .unwrap()
+        .assert_satisfied();
+    println!("Mock run successful");
+}
+
+/// k_srs -- 2^{k_srs} is the size of the srs
+/// k_circuit -- 2^{k_circuit} is the max number of rows the halo2-lib circuit generator will use in a single column
+pub fn run<CG: CircuitGenerator<Fr>>(
+    cg: &CG,
+    input: CG::WitnessInput,
+    instance: Vec<Fr>,
+    k_circuit: usize,
+    k_srs: u32,
+) {
+    println!("Generating key...");
+    let kg = keygen(cg, k_circuit, k_srs);
+    let params = kg.params.clone();
+    let vk = kg.pk.get_vk().clone();
+    println!("Generating snark...");
+    let snark = generate_snark(cg, kg, input.clone());
+    println!("Verifying snark...");
+    let res = verify_snark(vk, snark, params, instance);
+    println!("Snark verification result: {}", res);
+}