Try to speed up calcute speed

proof-tool/Cargo.toml

@@ -13,4 +13,7 @@ generic-array = "0.14"
 serde_json = "1"
 hex = "0.4"
 rand = { version = "0.8", features = ["min_const_gen"] }
-serde = { version = "1.0.208", features = ["derive"] }
+serde = { version = "1.0.208", features = ["derive"] }
+tokio = { version = "1.39", features = ["full"] }
+bytes = "1.7.1"
+flate2 = "1.0.33"

proof-tool/src/ercmerkle_tree.rs

@@ -1,3 +1,6 @@
+use std::io::Write;
+use std::sync::Arc;
+use bytes::{Bytes};
 use clap::ValueEnum;
 use generic_array::GenericArray;
 use generic_array::typenum::{U16, U32};
@@ -49,48 +52,52 @@ pub struct MerkleTreeData {
 }
 
 pub struct MerkleTree {
-    leaf_hash: Vec<HalfHash>,
     tree: Vec<Vec<HalfHash>>,
     root: Hash,
     hash_type: HashType,
 }
 
-impl MerkleTree {
+pub struct MerkleTreeBuilder {
+    leaf_hash: Vec<HalfHash>,
+    hash_type: HashType,
+}
+
+type MerkleTreeStable = Arc<MerkleTree>;
+
+impl MerkleTreeBuilder {
     pub fn new(hash_type: HashType) -> Self {
-        MerkleTree {
+        Self {
             leaf_hash: vec![],
-            tree: vec![],
-            root: Default::default(),
             hash_type,
         }
     }
 
-    pub fn load(data: MerkleTreeData) -> Self {
-        Self {
-            leaf_hash: vec![],
-            tree: data.tree.iter().map(|v|v.iter().map(|s|decode_half_hash(s)).collect()).collect(),
-            root: decode_hash(&data.root),
-            hash_type: data.hash_type,
-        }
+    pub fn add_leaf(&mut self, leaf: &[u8]) {
+        let hash32 = calc_hash(&self.hash_type, leaf);
+        self.leaf_hash.push(HalfHash::from_slice(&hash32.as_slice()[16..]).clone());
     }
 
-    pub fn save(&self) -> MerkleTreeData {
-        MerkleTreeData {
-            hash_type: self.hash_type.clone(),
-            tree: self.tree.iter().map(|v|v.iter().map(|h|hex(h.as_slice())).collect()).collect(),
-            root: hex(self.root.as_slice()),
+    pub async fn add_leafs(&mut self, leafs: Vec<Bytes>) {
+        let hash_type = self.hash_type.clone();
+        let mut handles = Vec::new();
+        for leaf in leafs {
+            handles.push(tokio::spawn(async move {
+                let hash_type = hash_type.clone();
+                let hash32 = calc_hash(&hash_type, &leaf);
+                HalfHash::from_slice(&hash32.as_slice()[16..]).clone()
+            }))
         }
-    }
 
-    pub fn add_leaf(&mut self, leaf: &[u8]) {
-        let hash32 = calc_hash(&self.hash_type, leaf);
-        self.leaf_hash.push(HalfHash::from_slice(&hash32.as_slice()[16..]).clone());
+        for handle in handles {
+            self.leaf_hash.push(handle.await.unwrap())
+        }
     }
 
-    pub fn calc_tree(&mut self) {
+    pub fn calc_tree(self, file_size: usize) -> MerkleTree {
+        let mut tree = MerkleTree::new(self.hash_type);
         let mut cur_layer = self.leaf_hash.clone();
         let mut next_layer= Vec::new();
-        self.tree.push(cur_layer.clone());
+        tree.tree.push(cur_layer.clone());
         let mut hash = Hash::default();
         while cur_layer.len() > 1 {
             for chunk in cur_layer.chunks(2) {
@@ -101,12 +108,56 @@ impl MerkleTree {
                     next_layer.push(HalfHash::from_slice(&hash.as_slice()[16..]).clone())
                 }
             }
-            self.tree.push(next_layer.clone());
+            tree.tree.push(next_layer.clone());
             cur_layer = next_layer.clone();
             next_layer.clear();
         }
 
-        self.root = hash;
+        hash.as_mut_slice().write(&file_size.to_be_bytes()).unwrap();
+        hash.as_mut_slice()[0] &= (1 << 6) - 1;
+
+        match self.hash_type {
+            HashType::Sha256 => {
+                // do nothing
+            }
+            HashType::Keccak256 => {
+                hash.as_mut_slice()[0] |= 1 << 7;
+            }
+        }
+
+        tree.root = hash;
+
+        tree
+    }
+}
+
+impl MerkleTree {
+    pub fn new(hash_type: HashType) -> Self {
+        Self {
+            tree: vec![],
+            root: Default::default(),
+            hash_type,
+        }
+    }
+
+    pub fn load(data: MerkleTreeData) -> MerkleTreeStable {
+        Arc::new(Self {
+            tree: data.tree.iter().map(|v|v.iter().map(|s|decode_half_hash(s)).collect()).collect(),
+            root: decode_hash(&data.root),
+            hash_type: data.hash_type,
+        })
+    }
+
+    pub fn save(&self) -> MerkleTreeData {
+        MerkleTreeData {
+            hash_type: self.hash_type.clone(),
+            tree: self.tree.iter().map(|v|v.iter().map(|h|hex(h.as_slice())).collect()).collect(),
+            root: hex(self.root.as_slice()),
+        }
+    }
+
+    pub fn leaf_size(&self) -> usize {
+        self.tree[0].len()
     }
 
     pub fn get_root(&self) -> &Hash {

proof-tool/src/main.rs

@@ -1,11 +1,16 @@
 mod ercmerkle_tree;
 
-use std::env::join_paths;
-use std::io::{Read, Write};
+use std::fs::File;
+use std::io::{BufReader, BufWriter, Read, Write};
 use std::os::windows::prelude::FileExt;
 use std::path::PathBuf;
+use std::time::Instant;
+use bytes::{Buf, BufMut, BytesMut};
 use clap::{Parser, Subcommand};
-use crate::ercmerkle_tree::{calc_hash, HashType, hex, MerkleTree};
+use flate2::Compression;
+use flate2::read::DeflateDecoder;
+use flate2::write::DeflateEncoder;
+use crate::ercmerkle_tree::{HashType, hex, MerkleTree, MerkleTreeBuilder};
 
 fn compare_bytes(a: &[u8], b: &[u8]) -> i32 {
     let n = a.len().min(b.len());
@@ -22,6 +27,12 @@ fn compare_bytes(a: &[u8], b: &[u8]) -> i32 {
 
 #[derive(Parser)]
 struct App {
+    #[arg(short, long)]
+    benchmark: bool,
+    #[arg(short, long, default_value = "64")]
+    task_size: u64,
+    #[arg(short, long)]
+    compress: bool,
     #[command(subcommand)]
     command: Subcommands
 }
@@ -38,107 +49,145 @@ enum Subcommands {
         #[arg(value_name="FILE")]
         file_path: PathBuf,
         nonce_hash: String,
-
+        #[arg(help = "only for debug")]
         leaf_index: Option<u64>,
     }
 }
 
-fn main() {
-    let cli = App::parse();
+#[tokio::main]
+async fn main() {
+    let mut cli = App::parse();
+    let mut start = Instant::now();
     match cli.command {
         Subcommands::Create { file_path, hash_type } => {
-            let mut file = std::fs::File::open(&file_path).unwrap();
+            let mut file = File::open(&file_path).unwrap();
             let size = file.metadata().unwrap().len();
-            println!("calcuting merkle tree...");
-            let mut merkle_tree = MerkleTree::new(hash_type);
-            let mut read_buf = Vec::with_capacity(1024);
-            read_buf.resize(1024, 0);
+            println!("calcuting merkle leaf...");
+            let mut merkle_builder = MerkleTreeBuilder::new(hash_type);
             loop {
-                read_buf.fill(0);
+                let mut read_buf = BytesMut::zeroed((1024 * cli.task_size) as usize);
                 let readed = file.read(&mut read_buf).unwrap();
                 if readed == 0 { break; }
-                merkle_tree.add_leaf(&read_buf);
-            }
-            merkle_tree.calc_tree();
+                let mut bufs = Vec::new();
 
-            let mut root_hash = merkle_tree.get_root().clone();
-            root_hash.as_mut_slice().write(&size.to_be_bytes()).unwrap();
-            root_hash.as_mut_slice()[0] &= (1 << 6) - 1;
+                let leafs = (readed as f64 / 1024f64).ceil() as usize;
 
-            match hash_type {
-                HashType::Sha256 => {
-                    // do nothing
-                }
-                HashType::Keccak256 => {
-                    root_hash.as_mut_slice()[0] |= 1 << 7;
+                for _ in 0..leafs {
+                    bufs.push(read_buf.split_to(1024).freeze());
                 }
+
+                merkle_builder.add_leafs(bufs).await;
+            }
+            if cli.benchmark {
+                let dur = Instant::now().duration_since(start).as_secs();
+                println!("calcuting leaf speed: {:.2} MB/sec, total use {dur} secs", (size as f64) / 1024f64 / 1024f64 / (dur as f64));
             }
 
-            merkle_tree.update_root(root_hash);
+            println!("calcuting merkle tree...");
+            start = Instant::now();
+            let tree = merkle_builder.calc_tree(size as usize);
+
+            if cli.benchmark {
+                let dur = Instant::now().duration_since(start).as_secs();
+                let leafs = tree.leaf_size();
+                println!("calcuting tree speed: {:.2} leafs/sec, total use {dur} secs", leafs as f64 / (dur as f64));
+            }
+            println!("save merkle tree...");
+            let mut merkle_data = BytesMut::new().writer();
+            serde_json::to_writer(&mut merkle_data, &tree.save()).unwrap();
+            if cli.compress {
+                let file = File::create(file_path.with_extension("merkle.lzma")).unwrap();
+                let mut encoder = DeflateEncoder::new(file, Compression::default());
+                encoder.write(merkle_data.get_ref()).unwrap();
+                encoder.finish().unwrap().flush().unwrap();
+            } else {
+                std::fs::write(file_path.with_extension("merkle"), merkle_data.get_ref()).unwrap();
+            }
 
-            let merkle_data = merkle_tree.save();
-            std::fs::write(file_path.with_extension("merkle"), &serde_json::to_vec(&merkle_data).unwrap()).unwrap();
+            println!("create file root hash 0x{}", hex::encode(tree.get_root()));
 
-            println!("create file root hash 0x{}", hex::encode(&root_hash));
         }
         Subcommands::Proof { file_path, nonce_hash , leaf_index } => {
-            let merkle_data = serde_json::from_slice(&std::fs::read(file_path.with_extension("merkle")).unwrap()).unwrap();
+            println!("reading merkle tree...");
+            let merkle_data = if cli.compress {
+
+                let mut decoder = DeflateDecoder::new(File::open(file_path.with_extension("merkle")).unwrap());
+                let mut data_buf = Vec::new();
+                decoder.read_to_end(&mut data_buf).unwrap();
+                serde_json::from_slice(&data_buf).unwrap()
+            } else {
+                serde_json::from_slice(&std::fs::read(file_path.with_extension("merkle")).unwrap()).unwrap()
+            };
+
             let merkle_tree = MerkleTree::load(merkle_data);
             let hash = hex::decode(&nonce_hash.as_str()[2..]).unwrap();
 
-            let file = std::fs::File::open(file_path).unwrap();
+            let mut file = File::open(file_path).unwrap();
             let length = file.metadata().unwrap().len();
             let total_leaf_size = (length as f64 / 1024f64).ceil() as u64;
 
             let mut min_root: Option<ercmerkle_tree::Hash> = None;
             let mut min_index = None;
-
-            let mut read_buf = Vec::with_capacity(1024);
-            read_buf.resize(1024, 0);
+
             if let Some(i) = leaf_index {
-                read_buf.fill(0);
+                let mut read_buf = BytesMut::zeroed(1024);
                 file.seek_read(&mut read_buf, i * 1024).unwrap();
                 let path = merkle_tree.get_path(i);
                 let new_leaf: Vec<u8> = read_buf.iter().chain(hash.iter()).map(|v|*v).collect();
                 let new_root = merkle_tree.proof_by_path(path, i, &new_leaf);
                 //println!("index {} new root {}", i, hex(new_root.as_slice()));
-                if let Some(root) = min_root {
-                    if compare_bytes(new_root.as_slice(), root.as_slice()) < 0 {
-                        min_root.insert(new_root);
-                        min_index.insert(i);
-                    }
-                } else {
-                    min_root.insert(new_root);
-                    min_index.insert(i);
-                }
+                let _ = min_root.insert(new_root);
+                let _ = min_index.insert(i);
             } else {
-                for i in 0..total_leaf_size {
-                    read_buf.fill(0);
-                    file.seek_read(&mut read_buf, i*1024).unwrap();
-                    let path = merkle_tree.get_path(i);
-                    let new_leaf: Vec<u8> = read_buf.iter().chain(hash.iter()).map(|v|*v).collect();
-                    let new_root = merkle_tree.proof_by_path(path, i, &new_leaf);
-                    //println!("index {} new root {}", i, hex(new_root.as_slice()));
-                    if let Some(root) = min_root {
-                        if compare_bytes(new_root.as_slice(), root.as_slice()) < 0 {
-                            min_root.insert(new_root);
-                            min_index.insert(i);
+                println!("finding min index...");
+                //let mut read_buf = BytesMut::zeroed((1024 * cli.task_size) as usize);
+                let mut i = 0;
+                loop {
+                    let mut read_buf = BytesMut::zeroed((1024 * cli.task_size) as usize);
+                    let readed = file.read(&mut read_buf).unwrap();
+                    if readed == 0 { break; }
+
+                    let mut handles = Vec::new();
+
+                    let leafs = (readed as f64 / 1024f64).ceil() as usize;
+                    for _ in 0..leafs {
+                        let tree = merkle_tree.clone();
+                        let mut buf = read_buf.split_to(1024);
+                        buf.extend_from_slice(&hash);
+                        handles.push(tokio::spawn(async move {
+                            (tree.proof_by_path(tree.get_path(i), i, &buf.freeze()), i)
+                        }));
+                        i += 1;
+                    }
+
+                    for handle in handles {
+                        let (new_root, index) = handle.await.unwrap();
+                        if let Some(root) = min_root {
+                            if compare_bytes(new_root.as_slice(), root.as_slice()) < 0 {
+                                let _ = min_root.insert(new_root);
+                                let _ = min_index.insert(index);
+                            }
+                        } else {
+                            let _ = min_root.insert(new_root);
+                            let _ = min_index.insert(index);
                         }
-                    } else {
-                        min_root.insert(new_root);
-                        min_index.insert(i);
                     }
+
                 }
             }
 
-
             println!("found min index {}, min root {}", min_index.unwrap(), hex(min_root.unwrap().as_slice()));
             let paths: Vec<String> = merkle_tree.get_path(min_index.unwrap()).iter().map(|v|hex(v.as_slice())).collect();
             println!("min path [");
             for path in paths {
                 println!("\t{},", path)
             }
-            println!("]")
+            println!("]");
+
+            if cli.benchmark {
+                let dur = Instant::now().duration_since(start).as_secs();
+                println!("create speed: {:.2} MB/sec, total use {dur} secs", (length as f64) / 1024f64 / 1024f64 / (dur as f64));
+            }
         }
     }