feat: benchmark jellyfish merkle tree

Cargo.toml

@@ -21,3 +21,5 @@ chrono = "0.4.38"
 hex = "0.4.3"
 eth_trie = "0.4.0"
 ethereum-trie = "=0.5.0"
+jmt = { git = "https://github.com/jingjunLi/jmt", branch = "feat-bench-mock" }
+sha2 = { version = "0.10"} 

src/jmt/mod.rs

@@ -0,0 +1,89 @@
+use async_trait::async_trait;
+use jmt::{JellyfishMerkleTree, KeyHash, MockTreeStore};
+use sha2::Sha256;
+
+use crate::database::Database;
+use std::sync::Arc;
+use tokio::sync::Mutex;
+
+pub struct JMTDatabase {
+    db: Arc<Mutex<MockTreeStore>>,
+    version: Arc<Mutex<u64>>,
+}
+
+#[async_trait]
+impl Database for JMTDatabase {
+    async fn open(datadir: String) -> Result<Self, String> {
+        let mockdb = MockTreeStore::default();
+        Ok(JMTDatabase {
+            db: Arc::new(Mutex::new(mockdb)),
+            version: Arc::new(Mutex::new(0)),
+        })
+    }
+
+    async fn get(&self, key: &[u8; 32]) -> Result<Option<Vec<u8>>, String> {
+        let db = self.db.lock().await;
+        let tree: JellyfishMerkleTree<MockTreeStore, Sha256> = JellyfishMerkleTree::new(&db);
+        // get don't incr version
+        let version = {
+            let ver = self.version.lock().await;
+            *ver
+        };
+
+        println!("jmt get {:x?}, version: {}", key, version);
+        let val = tree
+            .get(KeyHash::with::<Sha256>(key), version)
+            .map_err(|e| format!("Failed to get key: {}", e))?; // 使用 map_err 转换错误类型
+
+        println!("jmt get {:x?}, version: {}, value:{:x?}", key, version, val);
+        Ok(val)
+    }
+
+    async fn put(&self, key: &[u8; 32], value: Vec<u8>) {
+        let db = self.db.lock().await;
+        let tree: JellyfishMerkleTree<MockTreeStore, Sha256> = JellyfishMerkleTree::new(&db);
+        let version = {
+            let mut ver = self.version.lock().await;
+            let current_version = *ver;
+            *ver += 1;
+            current_version
+        };
+        println!(
+            "jmt put key:{:x?}, value:{:x?}, version: {}",
+            key, value, version
+        );
+
+        let (_root0_hash, batch) = tree
+            .batch_put_value_sets(
+                vec![vec![(KeyHash::with::<Sha256>(key), value.clone())]],
+                None,
+                version,
+            )
+            .unwrap();
+        db.write_tree_update_batch(batch).unwrap();
+    }
+
+    async fn delete(&self, key: &[u8; 32]) {
+        let db = self.db.lock().await;
+        // TODO version
+        let tree: JellyfishMerkleTree<MockTreeStore, Sha256> = JellyfishMerkleTree::new(&db);
+        let version = {
+            let mut ver = self.version.lock().await;
+            let current_version = *ver;
+            *ver += 1;
+            current_version
+        };
+
+        println!("jmt delete {:x?}, version: {}", key, version);
+        let (_root0_hash, batch) = tree
+            .put_value_set(vec![(KeyHash::with::<Sha256>(key), None)], version)
+            .unwrap();
+
+        db.write_tree_update_batch(batch).unwrap();
+    }
+
+    async fn commit(&self) -> Result<(), String> {
+        println!("jmt commit");
+        Ok(())
+    }
+}

src/main.rs

@@ -2,11 +2,13 @@ use crate::database::Database;
 use clap::Parser;
 use eth_trie::MemoryMPT;
 use firewood::Firewood;
+use jmt::JMTDatabase;
 use runner::Runner;
 
 mod database;
 mod eth_trie;
 mod firewood;
+mod jmt;
 pub mod runner;
 mod stat;
 mod utils;
@@ -17,7 +19,7 @@ struct Args {
     #[arg(short, long)]
     engine: String,
 
-    #[arg(short, long, default_value = "./dataset")]
+    #[arg(short = 's', long, default_value = "./dataset")]
     datadir: String,
 
     #[arg(short = 'b', long = "bs", default_value_t = 3000)]
@@ -61,7 +63,16 @@ async fn main() {
                     println!("Open memory MPT success");
                     run(db, args).await;
                 }
-                Err(e) => println!("Failed to open firewood database: {}", e),
+                Err(e) => println!("Failed to open memory MPT database: {}", e),
+            };
+        }
+        "jmt" => {
+            match JMTDatabase::open(args.datadir.clone()).await {
+                Ok(db) => {
+                    println!("Open jmt MPT success");
+                    run(db, args).await;
+                }
+                Err(e) => println!("Failed to open jmt database: {}", e),
             };
         }
         _ => {