trie_memory/
memory.rs

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