diff --git a/src/database.rs b/src/database.rs index 329940b..6b82857 100644 --- a/src/database.rs +++ b/src/database.rs @@ -1,5 +1,7 @@ use crate::*; +use std::collections::HashMap; + /// Trait that must be implemented for a Database pub trait Database { /// Creates new instance of db @@ -17,4 +19,7 @@ pub trait Database { /// Puts the value to the db by the key fn put(&mut self, key: DBKey, value: Value) -> Result<()>; + + /// Batc + fn put_batch(&mut self, subtree: &HashMap) -> Result<()>; } diff --git a/src/hasher.rs b/src/hasher.rs index 10a544e..3769390 100644 --- a/src/hasher.rs +++ b/src/hasher.rs @@ -1,9 +1,11 @@ use crate::*; +use std::fmt::Debug; + /// Trait that must be implemented for Hash Function pub trait Hasher { /// Native type for the hash-function - type Fr: Copy + Eq + Default + Sync + Send; + type Fr: Copy + Eq + Default + Sync + Send + Debug; /// Serializes Self::Fr fn serialize(value: Self::Fr) -> Value; diff --git a/src/tree.rs b/src/tree.rs index 5ba94ae..772b70e 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -213,11 +213,23 @@ where let subtree = Arc::new(Mutex::new(subtree)); - let n = rayon::ThreadPoolBuilder::new() + let root_val = rayon::ThreadPoolBuilder::new() .num_threads(8) .build() .unwrap() - .install(|| Self::batch_recalculate(root_key, subtree, self.depth)); + .install(|| Self::batch_recalculate(root_key, Arc::clone(&subtree), self.depth)); + + let subtree = Mutex::into_inner(Arc::try_unwrap(subtree).unwrap()).unwrap(); + + // self.db.put_batch(&subtree.into_iter().)?; + + // Update root value and next_index in memory + self.root = root_val; + self.next_index = end; + + // Update next_index value in db + self.db + .put(NEXT_INDEX_KEY, self.next_index.to_be_bytes().to_vec())?; Ok(()) } @@ -262,20 +274,20 @@ where subtree: Arc>>, depth: usize, ) -> H::Fr { - // if key.0 == depth { - // return subtree.into_inner().unwrap().get(&key).u; - // } + if key.0 == depth { + return *subtree.lock().unwrap().get(&key).unwrap(); + } - // let (left, right) = rayon::join( - // || Self::batch_recalculate(key, subtree, depth), - // || Self::batch_recalculate(key, subtree, depth), - // ); + let (left, right) = rayon::join( + || Self::batch_recalculate(Key(key.0 + 1, key.1 * 2), Arc::clone(&subtree), depth), + || Self::batch_recalculate(Key(key.0 + 1, key.1 * 2 + 1), Arc::clone(&subtree), depth), + ); - H::default_leaf() - } + let result = H::hash(&[left, right]); - fn db_batch_insert(&mut self, subtree: HashMap) -> Result<()> { - Ok(()) + subtree.lock().unwrap().insert(key, result); + + result } /// Computes a Merkle proof for the leaf at the specified index diff --git a/tests/memory_keccak.rs b/tests/memory_keccak.rs index 9d52649..342e275 100644 --- a/tests/memory_keccak.rs +++ b/tests/memory_keccak.rs @@ -24,6 +24,10 @@ impl Database for MemoryDB { Ok(()) } + + fn put_batch(&mut self, subtree: &HashMap) -> Result<()> { + Ok(()) + } } impl Hasher for MyKeccak { diff --git a/tests/poseidon.rs b/tests/poseidon.rs index 977040e..cab941e 100644 --- a/tests/poseidon.rs +++ b/tests/poseidon.rs @@ -61,6 +61,10 @@ impl Database for MemoryDB { Ok(()) } + + fn put_batch(&mut self, subtree: &HashMap) -> Result<()> { + Ok(()) + } } impl Database for MySled { @@ -95,6 +99,10 @@ impl Database for MySled { Ok(()) } + + fn put_batch(&mut self, subtree: &HashMap) -> Result<()> { + Ok(()) + } } #[test] diff --git a/tests/sled_keccak.rs b/tests/sled_keccak.rs index 21b1366..fbe49aa 100644 --- a/tests/sled_keccak.rs +++ b/tests/sled_keccak.rs @@ -1,5 +1,6 @@ use hex_literal::hex; use pmtree::*; +use std::collections::HashMap; use std::fs; use tiny_keccak::{Hasher as _, Keccak}; @@ -38,6 +39,10 @@ impl Database for MySled { Ok(()) } + + fn put_batch(&mut self, subtree: &HashMap) -> Result<()> { + Ok(()) + } } impl Hasher for MyKeccak {