Support partial trie fetching and move mutable to trie-memory

Cargo.toml

@@ -7,5 +7,6 @@ members = [
   "./block-core",
   "./trie",
   "./trie/rocksdb",
+  "./trie/memory",
   "./bloom",
 ]

block/Cargo.toml

@@ -16,6 +16,7 @@ etcommon-bigint = { version = "0.2", path = "../bigint" }
 etcommon-rlp = { version = "0.2", path = "../rlp" }
 etcommon-bloom = { version = "0.2", path = "../bloom" }
 etcommon-trie = { version = "0.4", path = "../trie" }
+etcommon-trie-memory = { version = "0.4", path = "../trie/memory" }
 blockchain = "0.2"
 
 secp256k1-plus = { version = "0.5", optional = true }

block/src/block.rs

@@ -1,7 +1,7 @@
 use rlp::{self, Encodable, Decodable, RlpStream, DecoderError, UntrustedRlp};
 use bigint::{Address, Gas, H256, U256, B256, H64, H2048};
 use bloom::LogsBloom;
-use trie::FixedMemoryTrieMut;
+use trie_memory::FixedMemoryTrieMut;
 use sha3::{Keccak256, Digest};
 use std::collections::HashMap;
 use super::{Header, Transaction, Receipt, SignaturePatch};

block/src/lib.rs

@@ -5,6 +5,7 @@ extern crate secp256k1;
 extern crate sha3;
 extern crate blockchain;
 extern crate trie;
+extern crate trie_memory;
 extern crate block_core;
 #[cfg(test)] extern crate hexutil;
 #[cfg(test)] extern crate rand;

block/tests/genesis.rs

@@ -5,13 +5,14 @@ extern crate serde;
 extern crate serde_json;
 extern crate bigint;
 extern crate trie;
+extern crate trie_memory;
 extern crate block;
 extern crate sha3;
 extern crate rlp;
 extern crate rand;
 
 use bigint::{Address, H256, M256, U256};
-use trie::{FixedSecureMemoryTrieMut, MemoryTrieMut, TrieMut};
+use trie_memory::{FixedSecureMemoryTrieMut, MemoryTrieMut, TrieMut};
 use block::Account;
 use sha3::{Digest, Keccak256};
 use rand::Rng;

trie/memory/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "etcommon-trie-memory"
+version = "0.4.0"
+license = "Apache-2.0"
+authors = ["Wei Tang <hi@that.world>"]
+description = "Rocksdb adaptor for trie."
+repository = "https://github.com/ethereumproject/etcommon-rs"
+keywords = ["ethereum"]
+
+[lib]
+name = "trie_memory"
+
+[dependencies]
+etcommon-trie = { version = "0.4", path = ".." }
+etcommon-bigint = { version = "0.2", path = "../../bigint" }
+etcommon-rlp = { version = "0.2", path = "../../rlp" }
+sha3 = "0.6"
+
+[dev-dependencies]
+etcommon-hexutil = { version = "0.2", path = "../../hexutil" }

trie/memory/src/cache.rs

@@ -0,0 +1,108 @@
+use trie::merkle::MerkleNode;
+use bigint::H256;
+use rlp::Rlp;
+use std::ptr;
+use std::collections::HashMap;
+use std::cell::{RefCell, Cell};
+
+pub struct Cache {
+    map: RefCell<HashMap<H256, usize>>,
+
+    cache_head: Cell<*mut Node>,
+    cache_len: Cell<usize>,
+}
+
+struct Node {
+    next: *mut Node,
+    value: Vec<u8>,
+}
+
+impl Drop for Cache {
+    fn drop(&mut self) {
+        if self.cache_head.get().is_null() {
+            return;
+        }
+
+        let mut all_ptrs = Vec::new();
+        all_ptrs.push(self.cache_head.get());
+
+        let mut cur_node = unsafe { &*self.cache_head.get() };
+
+        loop {
+            if cur_node.next.is_null() {
+                break;
+            }
+
+            all_ptrs.push(cur_node.next);
+            cur_node = unsafe { &*cur_node.next };
+        }
+
+        for ptr in all_ptrs {
+            unsafe { Box::from_raw(ptr); }
+        }
+    }
+}
+
+impl Cache {
+    fn at<'a>(&'a self, index: usize) -> Option<&'a [u8]> {
+        if self.cache_head.get().is_null() {
+            return None;
+        }
+
+        let mut cur_index = self.cache_len.get() - 1;
+        let mut cur_node = unsafe { &*self.cache_head.get() };
+
+        loop {
+            if cur_index < index {
+                return None;
+            }
+
+            if cur_index == index {
+                return Some(cur_node.value.as_ref());
+            }
+
+            if cur_node.next.is_null() {
+                return None;
+            }
+
+            cur_index -= 1;
+            cur_node = unsafe { &*cur_node.next };
+        }
+    }
+
+    pub fn new() -> Cache {
+        Cache {
+            map: RefCell::new(HashMap::new()),
+
+            cache_head: Cell::new(ptr::null_mut()),
+            cache_len: Cell::new(0),
+        }
+    }
+
+    pub fn insert<'a>(&'a self, key: H256, value: Vec<u8>) -> &'a [u8] {
+        let index = self.cache_len.get();
+        self.cache_len.set(self.cache_len.get() + 1);
+
+        self.map.borrow_mut().insert(key, index);
+        let node_ptr = Box::into_raw(Box::new(Node {
+            next: self.cache_head.get(),
+            value: value,
+        }));
+        self.cache_head.set(node_ptr);
+
+        self.at(index).unwrap()
+    }
+
+    pub fn get<'a>(&'a self, key: H256) -> Option<&'a [u8]> {
+        let mut map = self.map.borrow_mut();
+        match map.get(&key) {
+            Some(index) => Some(self.at(*index).unwrap()),
+            None => None,
+        }
+    }
+
+    pub fn contains_key(&self, key: H256) -> bool {
+        let mut map = self.map.borrow_mut();
+        map.contains_key(&key)
+    }
+}

trie/src/gc.rs → trie/memory/src/gc.rs

@@ -1,5 +1,6 @@
 use bigint::H256;
-use {Change, TrieMut, DatabaseHandle, get, insert, delete};
+use trie::{Change, DatabaseHandle, get, insert, delete};
+use TrieMut;
 
 pub trait ItemCounter {
     fn increase(&mut self, key: H256) -> usize;
@@ -11,9 +12,11 @@ pub trait DatabaseMut {
     fn set(&mut self, key: H256, value: Option<&[u8]>);
 }
 
-impl<'a, D: DatabaseMut> DatabaseHandle for &'a D {
-    fn get(&self, key: H256) -> &[u8] {
-        DatabaseMut::get(*self, key)
+struct DatabaseMutHandle<'a, D: DatabaseMut + 'a>(&'a D);
+
+impl<'a, D: DatabaseMut> DatabaseHandle for DatabaseMutHandle<'a, D> {
+    fn get(&self, key: H256) -> Option<&[u8]> {
+        Some(DatabaseMut::get(self.0, key))
     }
 }
 
@@ -62,20 +65,20 @@ impl<'a, D: DatabaseMut> TrieMut for DatabaseTrieMut<'a, D> {
     }
 
     fn insert(&mut self, key: &[u8], value: &[u8]) {
-        let (new_root, change) = insert(self.root, &self.database, key, value);
+        let (new_root, change) = insert(self.root, &DatabaseMutHandle(self.database), key, value).unwrap();
 
         self.change.merge(&change);
         self.root = new_root;
     }
 
     fn delete(&mut self, key: &[u8]) {
-        let (new_root, change) = delete(self.root, &self.database, key);
+        let (new_root, change) = delete(self.root, &DatabaseMutHandle(self.database), key).unwrap();
 
         self.change.merge(&change);
         self.root = new_root;
     }
 
     fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
-        get(self.root, &self.database, key).map(|v| v.into())
+        get(self.root, &DatabaseMutHandle(self.database), key).unwrap().map(|v| v.into())
     }
 }

trie/memory/src/lib.rs

@@ -0,0 +1,46 @@
+extern crate bigint;
+#[macro_use]
+extern crate trie;
+extern crate rlp;
+extern crate sha3;
+#[cfg(test)] extern crate hexutil;
+
+pub mod gc;
+mod memory;
+mod mutable;
+mod cache;
+
+use cache::Cache;
+use bigint::H256;
+use trie::DatabaseHandle;
+
+pub use memory::*;
+pub use mutable::*;
+
+pub trait CachedDatabaseHandle {
+    fn get(&self, key: H256) -> Vec<u8>;
+}
+
+pub struct CachedHandle<D: CachedDatabaseHandle> {
+    db: D,
+    cache: Cache,
+}
+
+impl<D: CachedDatabaseHandle> CachedHandle<D> {
+    pub fn new(db: D) -> Self {
+        Self {
+            db,
+            cache: Cache::new(),
+        }
+    }
+}
+
+impl<D: CachedDatabaseHandle> DatabaseHandle for CachedHandle<D> {
+    fn get(&self, key: H256) -> Option<&[u8]> {
+        if !self.cache.contains_key(key) {
+            Some(self.cache.insert(key, self.db.get(key)))
+        } else {
+            Some(self.cache.get(key).unwrap())
+        }
+    }
+}

trie/src/memory.rs → trie/memory/src/memory.rs

@@ -1,16 +1,10 @@
 use bigint::H256;
-use {DatabaseHandle, Change, insert, delete, build, get,
-     TrieMut, FixedTrieMut, FixedSecureTrieMut,
+use trie::{DatabaseHandle, Change, insert, delete, build, get, EMPTY_TRIE_HASH};
+use {TrieMut, FixedTrieMut, FixedSecureTrieMut,
      AnyTrieMut, AnySecureTrieMut, SecureTrieMut};
 
 use std::collections::HashMap;
 
-impl<'a> DatabaseHandle for &'a HashMap<H256, Vec<u8>> {
-    fn get(&self, hash: H256) -> &[u8] {
-        HashMap::get(self, &hash).unwrap()
-    }
-}
-
 /// A memory-backed trie.
 #[derive(Clone, Debug)]
 pub struct MemoryTrieMut {
@@ -37,7 +31,7 @@ impl Default for MemoryTrieMut {
     fn default() -> Self {
         Self {
             database: HashMap::new(),
-            root: empty_trie_hash!(),
+            root: EMPTY_TRIE_HASH,
         }
     }
 }
@@ -54,21 +48,21 @@ impl TrieMut for MemoryTrieMut {
     }
 
     fn insert(&mut self, key: &[u8], value: &[u8]) {
-        let (new_root, change) = insert(self.root, &&self.database, key, value);
+        let (new_root, change) = insert(self.root, &&self.database, key, value).unwrap();
 
         self.apply_change(change);
         self.root = new_root;
     }
 
     fn delete(&mut self, key: &[u8]) {
-        let (new_root, change) = delete(self.root, &&self.database, key);
+        let (new_root, change) = delete(self.root, &&self.database, key).unwrap();
 
         self.apply_change(change);
         self.root = new_root;
     }
 
     fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
-        get(self.root, &&self.database, key).map(|v| v.into())
+        get(self.root, &&self.database, key).unwrap().map(|v| v.into())
     }
 }
 
@@ -99,7 +93,8 @@ impl MemoryTrieMut {
 mod tests {
     use {TrieMut};
     use super::MemoryTrieMut;
-    use merkle::MerkleNode;
+    use trie::EMPTY_TRIE_HASH;
+    use trie::merkle::MerkleNode;
     use rlp::Rlp;
 
     use std::collections::HashMap;
@@ -140,7 +135,7 @@ mod tests {
         }
 
         assert!(mtrie.database.len() == 0);
-        assert!(mtrie.root == empty_trie_hash!());
+        assert!(mtrie.root == EMPTY_TRIE_HASH);
     }
 
     #[test]

trie/src/mutable.rs → trie/memory/src/mutable.rs

@@ -1,6 +1,6 @@
 use bigint::H256;
-use rlp::{self, Rlp};
 use sha3::{Digest, Keccak256};
+use rlp::{self, Rlp};
 
 use std::marker::PhantomData;
 

trie/rocksdb/Cargo.toml

@@ -13,4 +13,5 @@ name = "trie_rocksdb"
 [dependencies]
 etcommon-trie = { version = "0.4", path = ".." }
 etcommon-bigint = { version = "0.2", path = "../../bigint" }
+etcommon-trie-memory = { version = "0.4", path = "../memory" }
 rocksdb = { git = "https://github.com/paritytech/rust-rocksdb" }