extend ExecutionArg

DelphinusLab · Aug 1, 2024 · ff37c96 · ff37c96
1 parent 0860174
commit ff37c96
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 29 deletions.
diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs
@@ -1,8 +1,11 @@
 #![deny(warnings)]
 #![allow(clippy::too_many_arguments, clippy::while_let_on_iterator)]
 
+use std::cell::RefCell;
+use std::collections::HashMap;
 use std::fs;
 use std::path::PathBuf;
+use std::rc::Rc;
 
 use anyhow::Result;
 use app_builder::app;
@@ -82,6 +85,8 @@ fn main() -> Result<()> {
                     public_inputs,
                     private_inputs,
                     context_inputs,
+                    indexed_witness: Rc::new(RefCell::new(HashMap::default())),
+                    tree_db: None,
                 },
                 arg.running_arg.context_output,
                 arg.instruction_limit,
@@ -152,6 +157,8 @@ fn main() -> Result<()> {
                     public_inputs,
                     private_inputs,
                     context_inputs,
+                    indexed_witness: Rc::new(RefCell::new(HashMap::default())),
+                    tree_db: None,
                 },
                 arg.running_arg.context_output,
                 arg.mock_test,

diff --git a/crates/host/src/lib.rs b/crates/host/src/lib.rs
@@ -18,44 +18,31 @@ use std::collections::HashMap;
 use zkwasm_host_circuits::host::db::TreeDB;
 use zkwasm_host_circuits::proof::OpType;
 
-pub struct StandardExecutionArg {
-    /// Public inputs for `wasm_input(1)`
-    pub public_inputs: Vec<u64>,
-    /// Private inputs for `wasm_input(0)`
-    pub private_inputs: Vec<u64>,
-    /// Context inputs for `wasm_read_context()`
-    pub context_inputs: Vec<u64>,
-    /// indexed witness context
-    pub indexed_witness: Rc<RefCell<HashMap<u64, Vec<u64>>>>,
-    /// db src
-    pub tree_db: Option<Rc<RefCell<dyn TreeDB>>>,
-}
-
 #[derive(Serialize, Deserialize, Debug)]
 pub struct HostEnvConfig {
     pub ops: Vec<OpType>,
 }
 
 impl HostEnvConfig {
-    fn register_op(op: &OpType, env: &mut HostEnv) {
+    fn register_op(op: &OpType, env: &mut HostEnv, tree_db: Option<Rc<RefCell<dyn TreeDB>>>) {
         match op {
             OpType::BLS381PAIR => host::ecc_helper::bls381::pair::register_blspair_foreign(env),
             OpType::BLS381SUM => host::ecc_helper::bls381::sum::register_blssum_foreign(env),
             OpType::BN256PAIR => host::ecc_helper::bn254::pair::register_bn254pair_foreign(env),
             OpType::BN256SUM => host::ecc_helper::bn254::sum::register_bn254sum_foreign(env),
             OpType::POSEIDONHASH => host::hash_helper::poseidon::register_poseidon_foreign(env),
             OpType::MERKLE => {
-                host::merkle_helper::merkle::register_merkle_foreign(env, None);
-                host::merkle_helper::datacache::register_datacache_foreign(env, None);
+                host::merkle_helper::merkle::register_merkle_foreign(env, tree_db.clone());
+                host::merkle_helper::datacache::register_datacache_foreign(env, tree_db);
             }
             OpType::JUBJUBSUM => host::ecc_helper::jubjub::sum::register_babyjubjubsum_foreign(env),
             OpType::KECCAKHASH => host::hash_helper::keccak256::register_keccak_foreign(env),
         }
     }
 
-    fn register_ops(&self, env: &mut HostEnv) {
+    fn register_ops(&self, env: &mut HostEnv, tree_db: Option<Rc<RefCell<dyn TreeDB>>>) {
         for op in &self.ops {
-            Self::register_op(op, env);
+            Self::register_op(op, env, tree_db.clone());
         }
     }
 }
@@ -92,7 +79,7 @@ impl HostEnvBuilder for StandardHostEnvBuilder {
             &mut env,
             Rc::new(RefCell::new(HashMap::new())),
         );
-        host_env_config.register_ops(&mut env);
+        host_env_config.register_ops(&mut env, None);
 
         env.finalize();
 
@@ -104,20 +91,13 @@ impl HostEnvBuilder for StandardHostEnvBuilder {
         let host_env_config = HostEnvConfig {
             ops: self.ops.clone(),
         };
-        let arg = StandardExecutionArg {
-            public_inputs: arg.public_inputs,
-            private_inputs: arg.private_inputs,
-            context_inputs: arg.context_inputs,
-            indexed_witness: Rc::new(RefCell::new(HashMap::new())),
-            tree_db: None,
-        };
 
         register_wasm_input_foreign(&mut env, arg.public_inputs, arg.private_inputs);
         register_require_foreign(&mut env);
         register_log_foreign(&mut env);
         register_context_foreign(&mut env, arg.context_inputs);
         host::witness_helper::register_witness_foreign(&mut env, arg.indexed_witness);
-        host_env_config.register_ops(&mut env);
+        host_env_config.register_ops(&mut env, arg.tree_db);
 
         env.finalize();
 

diff --git a/crates/playground/examples/binary_search.rs b/crates/playground/examples/binary_search.rs
@@ -8,6 +8,9 @@ use delphinus_zkwasm::runtime::host::default_env::ExecutionArg;
 use delphinus_zkwasm::runtime::host::HostEnvBuilder;
 use delphinus_zkwasm::runtime::monitor::table_monitor::TableMonitor;
 use pairing_bn256::bn256::Fr;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::rc::Rc;
 
 const K: u32 = MIN_K;
 
@@ -21,6 +24,8 @@ fn main() -> Result<()> {
             public_inputs: vec![0],
             private_inputs: vec![],
             context_inputs: vec![],
+            indexed_witness: Rc::new(RefCell::new(HashMap::default())),
+            tree_db: None,
         },
     );
     let mut monitor = TableMonitor::new(K, &vec![], TraceBackend::Memory, &env);

diff --git a/crates/playground/examples/context.rs b/crates/playground/examples/context.rs
@@ -8,6 +8,9 @@ use delphinus_zkwasm::runtime::host::default_env::ExecutionArg;
 use delphinus_zkwasm::runtime::host::HostEnvBuilder;
 use delphinus_zkwasm::runtime::monitor::table_monitor::TableMonitor;
 use pairing_bn256::bn256::Fr;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::rc::Rc;
 
 const K: u32 = MIN_K;
 
@@ -22,6 +25,8 @@ fn main() -> Result<()> {
                 public_inputs: vec![],
                 private_inputs: vec![],
                 context_inputs: vec![2, 1],
+                indexed_witness: Rc::new(RefCell::new(HashMap::default())),
+                tree_db: None,
             },
         );
 
@@ -44,6 +49,8 @@ fn main() -> Result<()> {
                 public_inputs: vec![],
                 private_inputs: vec![],
                 context_inputs: context_output.0,
+                indexed_witness: Rc::new(RefCell::new(HashMap::default())),
+                tree_db: None,
             },
         );
 

diff --git a/crates/playground/examples/fibonacci.rs b/crates/playground/examples/fibonacci.rs
@@ -8,6 +8,9 @@ use delphinus_zkwasm::runtime::host::default_env::ExecutionArg;
 use delphinus_zkwasm::runtime::host::HostEnvBuilder;
 use delphinus_zkwasm::runtime::monitor::table_monitor::TableMonitor;
 use pairing_bn256::bn256::Fr;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::rc::Rc;
 
 const K: u32 = MIN_K;
 
@@ -21,6 +24,8 @@ fn main() -> Result<()> {
             public_inputs: vec![5],
             private_inputs: vec![],
             context_inputs: vec![],
+            indexed_witness: Rc::new(RefCell::new(HashMap::default())),
+            tree_db: None,
         },
     );
     let mut monitor = TableMonitor::new(K, &vec![], TraceBackend::Memory, &env);

diff --git a/crates/playground/examples/phantom.rs b/crates/playground/examples/phantom.rs
@@ -8,6 +8,9 @@ use delphinus_zkwasm::runtime::host::default_env::ExecutionArg;
 use delphinus_zkwasm::runtime::host::HostEnvBuilder;
 use delphinus_zkwasm::runtime::monitor::table_monitor::TableMonitor;
 use pairing_bn256::bn256::Fr;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::rc::Rc;
 
 const K: u32 = MIN_K;
 
@@ -21,6 +24,8 @@ fn main() -> Result<()> {
             public_inputs: vec![2],
             private_inputs: vec![],
             context_inputs: vec![],
+            indexed_witness: Rc::new(RefCell::new(HashMap::default())),
+            tree_db: None,
         },
     );
     let mut monitor = TableMonitor::new(K, &vec!["search".to_string()], TraceBackend::Memory, &env);

diff --git a/crates/zkwasm/src/runtime/host/default_env.rs b/crates/zkwasm/src/runtime/host/default_env.rs
@@ -1,3 +1,9 @@
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::rc::Rc;
+
+use zkwasm_host_circuits::host::db::TreeDB;
+
 use crate::foreign::context::runtime::register_context_foreign;
 use crate::foreign::log_helper::register_log_foreign;
 use crate::foreign::require_helper::register_require_foreign;
@@ -6,15 +12,17 @@ use crate::foreign::wasm_input_helper::runtime::register_wasm_input_foreign;
 use super::host_env::HostEnv;
 use super::HostEnvBuilder;
 
-// TODO: remove me after refine tracer
-#[derive(Clone)]
 pub struct ExecutionArg {
     /// Public inputs for `wasm_input(1)`
     pub public_inputs: Vec<u64>,
     /// Private inputs for `wasm_input(0)`
     pub private_inputs: Vec<u64>,
     /// Context inputs for `wasm_read_context()`
     pub context_inputs: Vec<u64>,
+    /// indexed witness context
+    pub indexed_witness: Rc<RefCell<HashMap<u64, Vec<u64>>>>,
+    /// db src
+    pub tree_db: Option<Rc<RefCell<dyn TreeDB>>>,
 }
 
 pub struct DefaultHostEnvBuilder;