trie/
memory.rs

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/// A memory-backed trie.
15#[derive(Clone, Debug)]
16pub struct MemoryTrieMut {
17    database: HashMap<H256, Vec<u8>>,
18    root: H256,
19}
20
21/// A memory-backed trie where the value is operated on a fixed RLP
22/// value type.
23pub type FixedMemoryTrieMut<K, V> = FixedTrieMut<MemoryTrieMut, K, V>;
24/// A memory-backed trie where the key is hashed and the value is
25/// operated on a fixed RLP value type.
26pub type FixedSecureMemoryTrieMut<K, V> = FixedSecureTrieMut<MemoryTrieMut, K, V>;
27/// A memory-backed trie where the key is hashed.
28pub type SecureMemoryTrieMut = SecureTrieMut<MemoryTrieMut>;
29/// A memory-backed trie where the value is operated on any RLP
30/// values.
31pub type AnyMemoryTrieMut = AnyTrieMut<MemoryTrieMut>;
32/// A memory-backed trie where the key is hashed and the value is
33/// operated on any RLP values.
34pub 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    /// Build a memory trie from a map.
87    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}