1use bigint::H256;
2use {DatabaseHandle, Change, insert, delete, build, get,
3 TrieMut, FixedTrieMut, FixedSecureTrieMut,
4 AnyTrieMut, AnySecureTrieMut, SecureTrieMut};
5
6use std::collections::HashMap;
7
8impl<'a> DatabaseHandle for &'a HashMap<H256, Vec<u8>> {
9 fn get(&self, hash: H256) -> &[u8] {
10 HashMap::get(self, &hash).unwrap()
11 }
12}
13
14#[derive(Clone, Debug)]
16pub struct MemoryTrieMut {
17 database: HashMap<H256, Vec<u8>>,
18 root: H256,
19}
20
21pub type FixedMemoryTrieMut<K, V> = FixedTrieMut<MemoryTrieMut, K, V>;
24pub type FixedSecureMemoryTrieMut<K, V> = FixedSecureTrieMut<MemoryTrieMut, K, V>;
27pub type SecureMemoryTrieMut = SecureTrieMut<MemoryTrieMut>;
29pub type AnyMemoryTrieMut = AnyTrieMut<MemoryTrieMut>;
32pub type AnySecureMemoryTrieMut = AnySecureTrieMut<MemoryTrieMut>;
35
36impl Default for MemoryTrieMut {
37 fn default() -> Self {
38 Self {
39 database: HashMap::new(),
40 root: empty_trie_hash!(),
41 }
42 }
43}
44
45impl Into<HashMap<H256, Vec<u8>>> for MemoryTrieMut {
46 fn into(self) -> HashMap<H256, Vec<u8>> {
47 self.database
48 }
49}
50
51impl TrieMut for MemoryTrieMut {
52 fn root(&self) -> H256 {
53 self.root
54 }
55
56 fn insert(&mut self, key: &[u8], value: &[u8]) {
57 let (new_root, change) = insert(self.root, &&self.database, key, value);
58
59 self.apply_change(change);
60 self.root = new_root;
61 }
62
63 fn delete(&mut self, key: &[u8]) {
64 let (new_root, change) = delete(self.root, &&self.database, key);
65
66 self.apply_change(change);
67 self.root = new_root;
68 }
69
70 fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
71 get(self.root, &&self.database, key).map(|v| v.into())
72 }
73}
74
75impl MemoryTrieMut {
76 fn apply_change(&mut self, change: Change) {
77 for add in change.adds {
78 self.database.insert(add.0, add.1);
79 }
80
81 for remove in change.removes {
82 self.database.remove(&remove);
83 }
84 }
85
86 pub fn build(map: &HashMap<Vec<u8>, Vec<u8>>) -> Self {
88 let (new_root, change) = build(map);
89
90 let mut ret = Self::default();
91 ret.apply_change(change);
92 ret.root = new_root;
93
94 ret
95 }
96}
97
98#[cfg(test)]
99mod tests {
100 use {TrieMut};
101 use super::MemoryTrieMut;
102 use merkle::MerkleNode;
103 use rlp::Rlp;
104
105 use std::collections::HashMap;
106 use std::str::FromStr;
107 use std::cell::UnsafeCell;
108 use bigint::H256;
109
110 #[test]
111 fn trie_middle_leaf() {
112 let mut map = HashMap::new();
113 map.insert("key1aa".as_bytes().into(), "0123456789012345678901234567890123456789xxx".as_bytes().into());
114 map.insert("key1".as_bytes().into(), "0123456789012345678901234567890123456789Very_Long".as_bytes().into());
115 map.insert("key2bb".as_bytes().into(), "aval3".as_bytes().into());
116 map.insert("key2".as_bytes().into(), "short".as_bytes().into());
117 map.insert("key3cc".as_bytes().into(), "aval3".as_bytes().into());
118 map.insert("key3".as_bytes().into(), "1234567890123456789012345678901".as_bytes().into());
119
120 let mut btrie = MemoryTrieMut::build(&map);
121
122 assert_eq!(btrie.root(), H256::from_str("0xcb65032e2f76c48b82b5c24b3db8f670ce73982869d38cd39a624f23d62a9e89").unwrap());
123 assert_eq!(btrie.get("key2bb".as_bytes()), Some("aval3".as_bytes().into()));
124 assert_eq!(btrie.get("key2bbb".as_bytes()), None);
125
126 let mut mtrie = MemoryTrieMut::default();
127 for (key, value) in &map {
128 mtrie.insert(key, value);
129 }
130
131 assert_eq!(btrie.database, mtrie.database);
132
133 mtrie.insert("key2bbb".as_bytes(), "aval4".as_bytes());
134 mtrie.delete("key2bbb".as_bytes());
135
136 assert_eq!(btrie.database, mtrie.database);
137
138 for (key, value) in &map {
139 mtrie.delete(key);
140 }
141
142 assert!(mtrie.database.len() == 0);
143 assert!(mtrie.root == empty_trie_hash!());
144 }
145
146 #[test]
147 fn trie_two_keys() {
148 let mut mtrie = MemoryTrieMut::default();
149 mtrie.insert("key1".as_bytes(), "aval1".as_bytes());
150 mtrie.insert("key2bb".as_bytes(), "aval3".as_bytes());
151 let db1 = mtrie.database.clone();
152
153 mtrie.insert("key2bbb".as_bytes(), "aval4".as_bytes());
154 mtrie.delete("key2bbb".as_bytes());
155
156 assert_eq!(db1, mtrie.database);
157 }
158}