From 7d7f01da46c994a76185ab5d705ebd9545fc91b1 Mon Sep 17 00:00:00 2001
From: Tamir Hemo <tamir@succinct.xyz>
Date: Tue, 19 Sep 2023 11:32:13 -0700
Subject: [PATCH] refactor prove method

---
 plonky2/src/plonk/prover.rs | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/plonky2/src/plonk/prover.rs b/plonky2/src/plonk/prover.rs
index b77f7aa5..c53ea78e 100644
--- a/plonky2/src/plonk/prover.rs
+++ b/plonky2/src/plonk/prover.rs
@@ -112,6 +112,21 @@ pub fn prove<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D:
     common_data: &CommonCircuitData<F, D>,
     inputs: PartialWitness<F>,
     timing: &mut TimingTree,
+) -> Result<ProofWithPublicInputs<F, C, D>> {
+    let partition_witness = timed!(
+        timing,
+        &format!("run {} generators", prover_data.generators.len()),
+        generate_partial_witness(inputs, prover_data, common_data)
+    );
+
+    prove_with_partition_witness(prover_data, common_data, partition_witness, timing)
+}
+
+pub fn prove_with_partition_witness<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>(
+    prover_data: &ProverOnlyCircuitData<F, C, D>,
+    common_data: &CommonCircuitData<F, D>,
+    mut partition_witness: PartitionWitness<F>,
+    timing: &mut TimingTree,
 ) -> Result<ProofWithPublicInputs<F, C, D>>
 where
     C::Hasher: Hasher<F>,
@@ -123,12 +138,6 @@ where
     let quotient_degree = common_data.quotient_degree();
     let degree = common_data.degree();
 
-    let mut partition_witness = timed!(
-        timing,
-        &format!("run {} generators", prover_data.generators.len()),
-        generate_partial_witness(inputs, prover_data, common_data)
-    );
-
     set_lookup_wires(prover_data, common_data, &mut partition_witness);
 
     let public_inputs = partition_witness.get_targets(&prover_data.public_inputs);