From 18d5b6095a4ab54493f71c8583f559efe3d5d921 Mon Sep 17 00:00:00 2001
From: chriseth <chris@ethereum.org>
Date: Thu, 16 Nov 2023 14:48:28 +0100
Subject: [PATCH] Reuse existing inputs to query callback utility.

---
 halo2/Cargo.toml         |  1 +
 halo2/src/mock_prover.rs | 36 ++++++------------------------------
 2 files changed, 7 insertions(+), 30 deletions(-)

diff --git a/halo2/Cargo.toml b/halo2/Cargo.toml
index b36d75be2..03361a6d3 100644
--- a/halo2/Cargo.toml
+++ b/halo2/Cargo.toml
@@ -21,6 +21,7 @@ ast = { version = "0.1.0", path = "../ast" }
 [dev-dependencies]
 importer = { path = "../importer" }
 analysis = { path = "../analysis" }
+compiler = { path = "../compiler" }
 executor = { path = "../executor" }
 parser = { path = "../parser" }
 airgen = { path = "../airgen" }
diff --git a/halo2/src/mock_prover.rs b/halo2/src/mock_prover.rs
index 72886845e..1e379b2e2 100644
--- a/halo2/src/mock_prover.rs
+++ b/halo2/src/mock_prover.rs
@@ -58,39 +58,15 @@ mod test {
         let graph = airgen::compile(analysed);
         let pil = linker::link(graph).unwrap();
 
-        let query_callback = |query: &str| -> Option<Bn254Field> {
-            let items = query.split(',').map(|s| s.trim()).collect::<Vec<_>>();
-            match items[0] {
-                "\"input\"" => {
-                    assert_eq!(items.len(), 2);
-                    let index = items[1].parse::<usize>().unwrap();
-                    let value = inputs.get(index).cloned();
-                    if let Some(value) = value {
-                        log::trace!("Input query: Index {index} -> {value}");
-                    }
-                    value
-                }
-                "\"print\"" => {
-                    log::info!("Print: {}", items[1..].join(", "));
-                    Some(0.into())
-                }
-                "\"print_ch\"" => {
-                    print!("{}", items[1].parse::<u8>().unwrap() as char);
-                    Some(0.into())
-                }
-                "\"hint\"" => {
-                    assert_eq!(items.len(), 2);
-                    Some(Bn254Field::from_str(items[1]))
-                }
-                _ => None,
-            }
-        };
-
         let analyzed = pil_analyzer::analyze_string(&format!("{pil}"));
 
         let fixed = executor::constant_evaluator::generate(&analyzed);
-        let witness =
-            executor::witgen::WitnessGenerator::new(&analyzed, &fixed, query_callback).generate();
+        let witness = executor::witgen::WitnessGenerator::new(
+            &analyzed,
+            &fixed,
+            compiler::inputs_to_query_callback(inputs.to_vec()),
+        )
+        .generate();
 
         let fixed = to_owned_values(fixed);