User defined parameter k for halo2 mock prover (#155)

Ready for review. @leolara

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "chiquito"
-version = "0.1.2023101100"
+version = "0.1.2023101700"
 edition = "2021"
 license = "MIT OR Apache-2.0"
 authors = ["Leo Lara <[email protected]>"]

examples/fibonacci.py

@@ -79,6 +79,8 @@ def trace(self, n):
 
 fibo = Fibonacci()
 fibo_witness = fibo.gen_witness(7)
-fibo.halo2_mock_prover(fibo_witness)
+fibo.halo2_mock_prover(
+    fibo_witness, k=7
+)  # 2^k specifies the number of PLONKish table rows in Halo2
 another_fibo_witness = fibo.gen_witness(4)
-fibo.halo2_mock_prover(another_fibo_witness)
+fibo.halo2_mock_prover(another_fibo_witness, k=7)

examples/mimc7.py

@@ -167,4 +167,4 @@ def mapping(self, x_in_value, k_value):
 mimc7_super_witness = mimc7.gen_witness(F(1), F(2))
 # for key, value in mimc7_super_witness.items():
 #     print(f"{key}: {str(value)}")
-mimc7.halo2_mock_prover(mimc7_super_witness)
+mimc7.halo2_mock_prover(mimc7_super_witness, k = 10)

src/frontend/python/chiquito/dsl.py

@@ -70,7 +70,9 @@ def gen_witness(self: SuperCircuit, *args: Any) -> Dict[int, TraceWitness]:
         )  # so that we can generate different witness mapping in the next gen_witness() call
         return super_witness
 
-    def halo2_mock_prover(self: SuperCircuit, super_witness: Dict[int, TraceWitness]):
+    def halo2_mock_prover(
+        self: SuperCircuit, super_witness: Dict[int, TraceWitness], k: int = 16
+    ):
         for rust_id, witness in super_witness.items():
             if rust_id not in self.ast.sub_circuits:
                 raise ValueError(
@@ -79,7 +81,7 @@ def halo2_mock_prover(self: SuperCircuit, super_witness: Dict[int, TraceWitness]
             witness_json: str = witness.get_witness_json()
             super_witness[rust_id] = witness_json
         rust_chiquito.super_circuit_halo2_mock_prover(
-            list(self.ast.sub_circuits.keys()), super_witness
+            list(self.ast.sub_circuits.keys()), super_witness, k
         )
 
 
@@ -211,12 +213,12 @@ def gen_witness(self: Circuit, *args) -> TraceWitness:
     def get_ast_json(self: Circuit) -> str:
         return json.dumps(self.ast, cls=CustomEncoder, indent=4)
 
-    def halo2_mock_prover(self: Circuit, witness: TraceWitness):
+    def halo2_mock_prover(self: Circuit, witness: TraceWitness, k: int = 16):
         if self.rust_id == 0:
             ast_json: str = self.get_ast_json()
             self.rust_id: int = rust_chiquito.ast_to_halo2(ast_json)
         witness_json: str = witness.get_witness_json()
-        rust_chiquito.halo2_mock_prover(witness_json, self.rust_id)
+        rust_chiquito.halo2_mock_prover(witness_json, self.rust_id, k)
 
     def __str__(self: Circuit) -> str:
         return self.ast.__str__()

src/frontend/python/mod.rs

@@ -83,6 +83,7 @@ fn add_assignment_generator_to_rust_id(
 pub fn chiquito_super_circuit_halo2_mock_prover(
     rust_ids: Vec<UUID>,
     super_witness: HashMap<UUID, &str>,
+    k: usize,
 ) {
     let mut super_circuit_ctx = SuperCircuitContext::<Fr, ()>::default();
 
@@ -114,7 +115,7 @@ pub fn chiquito_super_circuit_halo2_mock_prover(
 
     let circuit = ChiquitoHalo2SuperCircuit::new(compiled, super_assignments);
 
-    let prover = MockProver::<Fr>::run(10, &circuit, circuit.instance()).unwrap();
+    let prover = MockProver::<Fr>::run(k as u32, &circuit, circuit.instance()).unwrap();
 
     let result = prover.verify_par();
 
@@ -138,7 +139,7 @@ fn rust_id_to_halo2(uuid: UUID) -> CircuitMapStore {
 
 /// Runs `MockProver` for a single circuit given JSON of `TraceWitness` and `rust_id` of the
 /// circuit.
-pub fn chiquito_halo2_mock_prover(witness_json: &str, rust_id: UUID) {
+pub fn chiquito_halo2_mock_prover(witness_json: &str, rust_id: UUID, k: usize) {
     let trace_witness: TraceWitness<Fr> =
         serde_json::from_str(witness_json).expect("Json deserialization to TraceWitness failed.");
     let (_, compiled, assignment_generator) = rust_id_to_halo2(rust_id);
@@ -147,7 +148,7 @@ pub fn chiquito_halo2_mock_prover(witness_json: &str, rust_id: UUID) {
         assignment_generator.map(|g| g.generate_with_witness(trace_witness)),
     );
 
-    let prover = MockProver::<Fr>::run(7, &circuit, circuit.instance()).unwrap();
+    let prover = MockProver::<Fr>::run(k as u32, &circuit, circuit.instance()).unwrap();
 
     let result = prover.verify_par();
 
@@ -1827,15 +1828,16 @@ fn ast_to_halo2(json: &PyString) -> u128 {
 }
 
 #[pyfunction]
-fn halo2_mock_prover(witness_json: &PyString, rust_id: &PyLong) {
+fn halo2_mock_prover(witness_json: &PyString, rust_id: &PyLong, k: &PyLong) {
     chiquito_halo2_mock_prover(
         witness_json.to_str().expect("PyString convertion failed."),
         rust_id.extract().expect("PyLong convertion failed."),
+        k.extract().expect("PyLong convertion failed."),
     );
 }
 
 #[pyfunction]
-fn super_circuit_halo2_mock_prover(rust_ids: &PyList, super_witness: &PyDict) {
+fn super_circuit_halo2_mock_prover(rust_ids: &PyList, super_witness: &PyDict, k: &PyLong) {
     let uuids = rust_ids
         .iter()
         .map(|rust_id| {
@@ -1864,7 +1866,11 @@ fn super_circuit_halo2_mock_prover(rust_ids: &PyList, super_witness: &PyDict) {
         })
         .collect::<HashMap<u128, &str>>();
 
-    chiquito_super_circuit_halo2_mock_prover(uuids, super_witness)
+    chiquito_super_circuit_halo2_mock_prover(
+        uuids,
+        super_witness,
+        k.extract().expect("PyLong convertion failed."),
+    )
 }
 
 #[pymodule]