support proof recorder in host merkle

DelphinusLab · Apr 11, 2024 · 01834b7 · 01834b7
1 parent 02ce400
commit 01834b7
Show file tree

Hide file tree

Showing 6 changed files with 55 additions and 18 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -11,7 +11,7 @@ halo2_proofs = { git = "https://github.com/DelphinusLab/halo2-gpu-specific.git",
 parity-wasm = { version = "0.42.0", features = ["sign_ext"] }
 wasmi = { path = "third-party/wasmi" }
 circuits-batcher = { git = "https://github.com/DelphinusLab/continuation-batcher.git" }
-zkwasm-host-circuits = { git = "https://github.com/DelphinusLab/zkWasm-host-circuits.git" }
+zkwasm-host-circuits = { git = "https://github.com/DelphinusLab/zkWasm-host-circuits.git", branch="synthesize-helper" }
 
 [profile.dev]
 opt-level = 3
diff --git a/crates/cli/src/app_builder.rs b/crates/cli/src/app_builder.rs
@@ -202,6 +202,7 @@ pub trait AppBuilder: CommandBuilder {
                                 context_outputs: context_output.clone(),
                                 indexed_witness: Rc::new(RefCell::new(HashMap::new())),
                                 tree_db: None,
+                                merkle_proof_recorder: None,
                             },
                             HostEnvConfig::default(),
                         )?;
@@ -255,6 +256,7 @@ pub trait AppBuilder: CommandBuilder {
                                 context_outputs: context_out.clone(),
                                 indexed_witness: Rc::new(RefCell::new(HashMap::new())),
                                 tree_db: None,
+                                merkle_proof_recorder: None,
                             },
                             HostEnvConfig::default(),
                         )?;

diff --git a/crates/cli/test_cli.sh b/crates/cli/test_cli.sh
@@ -7,8 +7,8 @@ rm -rf params/*.data
 rm -rf output/*.data
 
 # Single test
-RUST_LOG=info cargo run --release --features cuda -- --host default -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm setup
-RUST_LOG=info cargo run --release --features cuda -- --host default -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm checksum
+RUST_LOG=info cargo run --release --features cuda -- --host default -k 22 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm setup
+RUST_LOG=info cargo run --release --features cuda -- --host default -k 22 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm checksum
 
-RUST_LOG=info cargo run --release --features cuda -- --host default -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm single-prove --public 133:i64 --public 2:i64
-RUST_LOG=info cargo run --release --features cuda -- --host default -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm single-verify
+RUST_LOG=info cargo run --release --features cuda -- --host default -k 22 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm single-prove --public 133:i64 --public 2:i64
+RUST_LOG=info cargo run --release --features cuda -- --host default -k 22 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm single-verify
diff --git a/crates/host/src/host/merkle_helper/merkle.rs b/crates/host/src/host/merkle_helper/merkle.rs
@@ -4,6 +4,7 @@ use delphinus_zkwasm::runtime::host::ForeignContext;
 use delphinus_zkwasm::runtime::host::ForeignStatics;
 use halo2_proofs::pairing::bn256::Fr;
 use std::cell::RefCell;
+use std::io::Write;
 use std::rc::Rc;
 use wasmi::tracer::Observer;
 use zkwasm_host_circuits::circuits::host::HostOpSelector;
@@ -19,6 +20,7 @@ use zkwasm_host_circuits::host::ForeignInst::MerkleSet;
 use zkwasm_host_circuits::host::ForeignInst::MerkleSetRoot;
 use zkwasm_host_circuits::host::Reduce;
 use zkwasm_host_circuits::host::ReduceRule;
+use zkwasm_host_circuits::proof::write_merkle_proof;
 
 const MERKLE_TREE_HEIGHT: usize = 32;
 
@@ -33,6 +35,7 @@ pub struct MerkleContext {
     pub mongo_merkle: Option<merklehelper::MongoMerkle<MERKLE_TREE_HEIGHT>>,
     pub mongo_datahash: datahelper::MongoDataHash,
     pub tree_db: Option<Rc<RefCell<dyn TreeDB>>>,
+    pub proof_recorder: Option<Rc<RefCell<dyn Write>>>,
     pub used_round: usize,
 }
 
@@ -41,7 +44,10 @@ fn new_reduce(rules: Vec<ReduceRule<Fr>>) -> Reduce<Fr> {
 }
 
 impl MerkleContext {
-    pub fn new(tree_db: Option<Rc<RefCell<dyn TreeDB>>>) -> Self {
+    pub fn new(
+        tree_db: Option<Rc<RefCell<dyn TreeDB>>>,
+        proof_recorder: Option<Rc<RefCell<dyn Write>>>,
+    ) -> Self {
         MerkleContext {
             set_root: new_reduce(vec![ReduceRule::Bytes(vec![], 4)]),
             get_root: new_reduce(vec![ReduceRule::Bytes(vec![], 4)]),
@@ -53,6 +59,7 @@ impl MerkleContext {
             mongo_merkle: None,
             mongo_datahash: datahelper::MongoDataHash::construct([0; 32], tree_db.clone()),
             tree_db,
+            proof_recorder,
             used_round: 0,
         }
     }
@@ -112,8 +119,13 @@ impl MerkleContext {
                 .as_mut()
                 .expect("merkle db not initialized");
             let hash = self.set.rules[0].bytes_value().unwrap();
-            mt.update_leaf_data_with_proof(index, &hash)
+            let proof = mt
+                .update_leaf_data_with_proof(index, &hash)
                 .expect("Unexpected failure: update leaf with proof fail");
+            self.proof_recorder.as_mut().map(|v| {
+                let mut writer = v.borrow_mut();
+                write_merkle_proof(&mut *writer, proof);
+            });
         }
     }
 
@@ -124,9 +136,15 @@ impl MerkleContext {
             .mongo_merkle
             .as_ref()
             .expect("merkle db not initialized");
-        let (leaf, _) = mt
+        let (leaf, proof) = mt
             .get_leaf_with_proof(index)
             .expect("Unexpected failure: get leaf fail");
+
+        self.proof_recorder.as_mut().map(|v| {
+            let mut writer = v.borrow_mut();
+            write_merkle_proof(&mut *writer, proof);
+        });
+
         let values = leaf.data_as_u64();
         if self.data_cursor == 0 {
             self.data = values;
@@ -149,10 +167,15 @@ impl ForeignContext for MerkleContext {
 }
 
 use specs::external_host_call_table::ExternalHostCallSignature;
-pub fn register_merkle_foreign(env: &mut HostEnv, tree_db: Option<Rc<RefCell<dyn TreeDB>>>) {
-    let foreign_merkle_plugin = env
-        .external_env
-        .register_plugin("foreign_merkle", Box::new(MerkleContext::new(tree_db)));
+pub fn register_merkle_foreign(
+    env: &mut HostEnv,
+    tree_db: Option<Rc<RefCell<dyn TreeDB>>>,
+    proof_recorder: Option<Rc<RefCell<dyn Write>>>,
+) {
+    let foreign_merkle_plugin = env.external_env.register_plugin(
+        "foreign_merkle",
+        Box::new(MerkleContext::new(tree_db, proof_recorder)),
+    );
 
     env.external_env.register_function(
         "merkle_setroot",

diff --git a/crates/host/src/lib.rs b/crates/host/src/lib.rs
@@ -14,6 +14,7 @@ use delphinus_zkwasm::runtime::host::HostEnvBuilder;
 use serde::Deserialize;
 use serde::Serialize;
 use std::collections::HashMap;
+use std::io::Write;
 use std::sync::Arc;
 use std::sync::Mutex;
 use zkwasm_host_circuits::host::db::TreeDB;
@@ -32,6 +33,8 @@ pub struct ExecutionArg {
     pub indexed_witness: Rc<RefCell<HashMap<u64, Vec<u64>>>>,
     /// db src
     pub tree_db: Option<Rc<RefCell<dyn TreeDB>>>,
+    /// merkle_proof_tracker:
+    pub merkle_proof_recorder: Option<Rc<RefCell<dyn Write>>>,
 }
 
 impl ContextOutput for ExecutionArg {
@@ -60,7 +63,12 @@ impl Default for HostEnvConfig {
 }
 
 impl HostEnvConfig {
-    fn register_ops(&self, env: &mut HostEnv, tree_db: Option<Rc<RefCell<dyn TreeDB>>>) {
+    fn register_ops(
+        &self,
+        env: &mut HostEnv,
+        tree_db: Option<Rc<RefCell<dyn TreeDB>>>,
+        merkle_proof_recorder: Option<Rc<RefCell<dyn Write>>>,
+    ) {
         for op in &self.ops {
             match op {
                 OpType::BLS381PAIR => host::ecc_helper::bls381::pair::register_blspair_foreign(env),
@@ -69,7 +77,11 @@ impl HostEnvConfig {
                 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, tree_db.clone());
+                    host::merkle_helper::merkle::register_merkle_foreign(
+                        env,
+                        tree_db.clone(),
+                        merkle_proof_recorder.clone(),
+                    );
                     host::merkle_helper::datacache::register_datacache_foreign(
                         env,
                         tree_db.clone(),
@@ -96,7 +108,7 @@ impl HostEnvBuilder for StandardHostEnvBuilder {
         register_require_foreign(&mut env);
         register_log_foreign(&mut env);
         register_context_foreign(&mut env, vec![], Arc::new(Mutex::new(vec![])));
-        envconfig.register_ops(&mut env, None);
+        envconfig.register_ops(&mut env, None, None);
         host::witness_helper::register_witness_foreign(
             &mut env,
             Rc::new(RefCell::new(HashMap::new())),
@@ -114,7 +126,7 @@ impl HostEnvBuilder for StandardHostEnvBuilder {
         register_log_foreign(&mut env);
         register_context_foreign(&mut env, arg.context_inputs, arg.context_outputs);
         host::witness_helper::register_witness_foreign(&mut env, arg.indexed_witness);
-        envconfig.register_ops(&mut env, arg.tree_db);
+        envconfig.register_ops(&mut env, arg.tree_db, arg.merkle_proof_recorder);
         env.finalize();
 
         (env, wasm_runtime_io)