trie_test/
lib.rs

1extern crate bigint;
2extern crate rlp;
3extern crate sha3;
4#[cfg(test)] extern crate hexutil;
5
6use bigint::H256;
7use rlp::Rlp;
8use sha3::{Digest, Keccak256};
9use std::collections::HashMap;
10use merkle::{MerkleValue, MerkleNode};
11use merkle::nibble::{self, NibbleVec, NibbleSlice, Nibble};
12use std::ops::{Deref, DerefMut};
13use std::borrow::Borrow;
14use std::marker::PhantomData;
15use std::clone::Clone;
16
17use self::cache::Cache;
18use self::database::{Change, ChangeSet};
19
20pub use self::database::{Database, DatabaseOwned, DatabaseGuard, MemoryDatabase, MemoryDatabaseGuard};
21pub use self::iter::{FixedMerkleIterator, MerkleIterator};
22
23macro_rules! empty_nodes {
24    () => (
25        [MerkleValue::Empty, MerkleValue::Empty,
26         MerkleValue::Empty, MerkleValue::Empty,
27         MerkleValue::Empty, MerkleValue::Empty,
28         MerkleValue::Empty, MerkleValue::Empty,
29         MerkleValue::Empty, MerkleValue::Empty,
30         MerkleValue::Empty, MerkleValue::Empty,
31         MerkleValue::Empty, MerkleValue::Empty,
32         MerkleValue::Empty, MerkleValue::Empty]
33    )
34}
35
36macro_rules! empty_trie_hash {
37    () => {
38        {
39            use std::str::FromStr;
40
41            H256::from_str("0x56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421").unwrap()
42        }
43    }
44}
45
46pub mod merkle;
47mod cache;
48mod database;
49mod iter;
50
51pub type MemorySecureTrie = SecureTrie<HashMap<H256, Vec<u8>>>;
52pub type MemoryTrie = Trie<HashMap<H256, Vec<u8>>>;
53pub type FixedMemoryTrie<K, V> = FixedTrie<HashMap<H256, Vec<u8>>, K, V>;
54pub type FixedMemorySecureTrie<K, V> = FixedSecureTrie<HashMap<H256, Vec<u8>>, K, V>;
55
56#[derive(Clone, Debug)]
57pub struct FixedTrie<D: DatabaseGuard, K: rlp::Encodable + rlp::Decodable, V: rlp::Encodable + rlp::Decodable>(
58    Trie<D>, PhantomData<(K, V)>
59);
60
61impl<D: DatabaseGuard, K: rlp::Encodable + rlp::Decodable, V: rlp::Encodable + rlp::Decodable> FixedTrie<D, K, V> {
62    pub fn new(trie: Trie<D>) -> Self {
63        FixedTrie(trie, PhantomData)
64    }
65
66    pub fn empty(database: D) -> Self {
67        FixedTrie(Trie::empty(database), PhantomData)
68    }
69
70    pub fn existing(database: D, root: H256) -> Self {
71        FixedTrie(Trie::existing(database, root), PhantomData)
72    }
73
74    pub fn root(&self) -> H256 { self.0.root() }
75    pub fn is_empty(&self) -> bool { self.0.is_empty() }
76
77    pub fn get(&self, key: &K) -> Option<V> {
78        self.0.get(key)
79    }
80
81    pub fn insert(&mut self, key: K, value: V) {
82        self.0.insert(key, value)
83    }
84
85    pub fn remove(&mut self, key: &K) {
86        self.0.remove(key)
87    }
88
89    pub fn iter(&self) -> FixedMerkleIterator<D, K, V> {
90        FixedMerkleIterator::new(self.0.iter())
91    }
92}
93
94#[derive(Clone, Debug)]
95pub struct FixedSecureTrie<D: DatabaseGuard, K: AsRef<[u8]>, V: rlp::Encodable + rlp::Decodable>(
96    SecureTrie<D>, PhantomData<(K, V)>
97);
98
99impl<D: DatabaseGuard, K: AsRef<[u8]>, V: rlp::Encodable + rlp::Decodable> FixedSecureTrie<D, K, V> {
100    pub fn new(trie: SecureTrie<D>) -> Self {
101        FixedSecureTrie(trie, PhantomData)
102    }
103
104    pub fn empty(database: D) -> Self {
105        FixedSecureTrie(SecureTrie::empty(database), PhantomData)
106    }
107
108    pub fn existing(database: D, root: H256) -> Self {
109        FixedSecureTrie(SecureTrie::existing(database, root), PhantomData)
110    }
111
112    pub fn root(&self) -> H256 { self.0.root() }
113    pub fn is_empty(&self) -> bool { self.0.is_empty() }
114
115    pub fn get(&self, key: &K) -> Option<V> {
116        self.0.get(key)
117    }
118
119    pub fn insert(&mut self, key: K, value: V) {
120        self.0.insert(key, value)
121    }
122
123    pub fn remove(&mut self, key: &K) {
124        self.0.remove(key)
125    }
126}
127
128#[derive(Clone, Debug)]
129pub struct SecureTrie<D: DatabaseGuard>(Trie<D>);
130
131impl<D: DatabaseGuard> SecureTrie<D> {
132    pub fn new(trie: Trie<D>) -> Self {
133        SecureTrie(trie)
134    }
135
136    pub fn empty(database: D) -> Self {
137        SecureTrie(Trie::empty(database))
138    }
139
140    pub fn existing(database: D, root: H256) -> Self {
141        SecureTrie(Trie::existing(database, root))
142    }
143
144    pub fn root(&self) -> H256 { self.0.root() }
145    pub fn is_empty(&self) -> bool { self.0.is_empty() }
146
147    fn secure_key<K: AsRef<[u8]>>(key: &K) -> Vec<u8> {
148        Keccak256::digest(key.as_ref()).as_slice().into()
149    }
150
151    pub fn get<K: AsRef<[u8]>, V: rlp::Decodable>(&self, key: &K) -> Option<V> {
152        self.0.get_raw(&Self::secure_key(key)).map(|v| rlp::decode(v.as_slice()))
153    }
154
155    pub fn insert<K: AsRef<[u8]>, V: rlp::Encodable>(&mut self, key: K, value: V) {
156        self.0.insert_raw(Self::secure_key(&key), rlp::encode(&value).to_vec())
157    }
158
159    pub fn remove<K: AsRef<[u8]>>(&mut self, key: &K) {
160        self.0.remove_raw(&Self::secure_key(key))
161    }
162}
163
164#[derive(Clone, Debug)]
165pub struct Trie<D: DatabaseGuard> {
166    pub database: D,
167    pub root: H256,
168}
169
170impl<D: DatabaseGuard> Trie<D> {
171    pub fn empty(database: D) -> Self {
172        Self {
173            database,
174            root: empty_trie_hash!()
175        }
176    }
177
178    pub fn existing(database: D, root: H256) -> Self {
179        if root == empty_trie_hash!() {
180            return Self::empty(database);
181        }
182
183        assert!(database.get(root).is_some());
184        Self {
185            database,
186            root
187        }
188    }
189
190    pub fn iter(&self) -> MerkleIterator<D> {
191        if self.root == empty_trie_hash!() {
192            MerkleIterator::empty(&self.database)
193        } else {
194            let value = self.database.get(self.root).unwrap();
195            MerkleIterator::new(&self.database, value)
196        }
197    }
198
199    pub fn root(&self) -> H256 {
200        self.root
201    }
202
203    pub fn is_empty(&self) -> bool {
204        self.root() == empty_trie_hash!()
205    }
206
207    pub fn get<K: rlp::Encodable, V: rlp::Decodable>(&self, key: &K) -> Option<V> {
208        let key = rlp::encode(key).to_vec();
209
210        self.get_raw(&key).map(|v| rlp::decode(v.as_slice()))
211    }
212
213    pub fn insert<K: rlp::Encodable, V: rlp::Encodable>(&mut self, key: K, value: V) {
214        let key = rlp::encode(&key).to_vec();
215        let value = rlp::encode(&value).to_vec();
216
217        self.insert_raw(key, value);
218    }
219
220    pub fn remove<K: rlp::Encodable>(&mut self, key: &K) {
221        let key = rlp::encode(key).to_vec();
222
223        self.remove_raw(&key)
224    }
225
226    fn copy_nodes<'a, 'b>(old_nodes: &'a [MerkleValue<'b>]) -> [MerkleValue<'b>; 16] {
227        debug_assert!(old_nodes.len() == 16);
228        let mut nodes = empty_nodes!();
229        for i in 0..16 {
230            nodes[i] = old_nodes[i].clone();
231        }
232        nodes
233    }
234
235    fn build_value<'a, 'b>(database: &'a mut Change<'b, D>, node: MerkleNode<'b>) -> MerkleValue<'b> {
236        if node.inlinable() {
237            MerkleValue::Full(Box::new(node))
238        } else {
239            let subnode = rlp::encode(&node).to_vec();
240            let hash = H256::from(Keccak256::digest(&subnode).as_slice());
241            database.set(hash, subnode);
242            MerkleValue::Hash(hash)
243        }
244    }
245
246    fn build_submap<'a, 'b: 'a, T: Iterator<Item=(&'a NibbleVec, &'a &'b [u8])>>(
247        common_len: usize, map: T
248    ) -> HashMap<NibbleVec, &'b [u8]> {
249        let mut submap = HashMap::new();
250        for (key, value) in map {
251            submap.insert(key.split_at(common_len).1.into(), value.clone());
252        }
253        submap
254    }
255
256    fn build_node<'a, 'b>(database: &'a mut Change<'b, D>, map: &HashMap<NibbleVec, &'b [u8]>) -> MerkleNode<'b> {
257        if map.len() == 0 {
258            panic!();
259        }
260
261        if map.len() == 1 {
262            let key = map.keys().next().unwrap();
263            return MerkleNode::Leaf(key.clone(), map.get(key).unwrap().clone());
264        }
265
266        debug_assert!(map.len() > 1);
267
268        let common: NibbleSlice = nibble::common_all(map.keys().map(|v| v.as_ref()));
269
270        if common.len() >= 1 {
271            let submap = Self::build_submap(common.len(), map.iter());
272            debug_assert!(submap.len() > 0);
273            let node = Self::build_node(database, &submap);
274            let value = Self::build_value(database, node);
275            return MerkleNode::Extension(common.into(), value);
276        }
277
278        let mut nodes = empty_nodes!();
279
280        for i in 0..16 {
281            let nibble_index: Nibble = i.into();
282
283            let submap = Self::build_submap(1, map.iter().filter(|&(key, value)| {
284                key.len() > 0 && key[0] == nibble_index
285            }));
286            let value = if submap.len() == 0 {
287                MerkleValue::Empty
288            } else {
289                let node = Self::build_node(database, &submap);
290                Self::build_value(database, node)
291            };
292            nodes[i] = value;
293        }
294
295        let additional = map.iter()
296            .filter(|&(key, value)| key.len() == 0).next()
297            .map(|(key, value)| value.clone());
298
299        return MerkleNode::Branch(nodes, additional);
300    }
301
302    pub fn build(mut database: D, map: &HashMap<Vec<u8>, Vec<u8>>) -> Self {
303        if map.len() == 0 {
304            return Self::empty(database);
305        }
306
307        let mut node_map = HashMap::new();
308        for (key, value) in map {
309            node_map.insert(nibble::from_key(key.as_ref()), value.as_ref());
310        }
311
312        let (changeset, root_rlp) = {
313            let mut change = Change::new(&database);
314            let node = Self::build_node(&mut change, &node_map);
315            (ChangeSet::from(change), rlp::encode(&node).to_vec())
316        };
317        let hash = H256::from(Keccak256::digest(&root_rlp).as_slice());
318        changeset.drain(&mut database, true);
319        database.set(hash, root_rlp);
320
321        Trie {
322            database,
323            root: hash
324        }
325    }
326    
327    fn get_by_value<'a, 'b>(database: &'a mut Change<D>, cache: &'a Cache, nibble: NibbleVec, value: MerkleValue<'a>) -> Option<&'a [u8]> {
328        match value {
329            MerkleValue::Empty => None,
330            MerkleValue::Full(sub_node) => {
331                let sub_node: &MerkleNode<'a> = sub_node.borrow();
332                let sub_node: MerkleNode<'a> = (*sub_node).clone();
333                Self::get_by_node(database, cache, nibble, sub_node)
334            },
335            MerkleValue::Hash(h) => {
336                let dbv = match database.get(h) {
337                    Some(val) => val,
338                    None => return None,
339                };
340                let node = cache.insert(h, dbv);
341                Self::get_by_node(database, cache, nibble, node)
342            },
343        }
344    }
345
346    fn get_by_node<'a, 'b>(database: &'a mut Change<D>, cache: &'a Cache, nibble: NibbleVec, node: MerkleNode<'a>) -> Option<&'a [u8]> {
347        match node {
348            MerkleNode::Leaf(node_nibble, node_value) => {
349                if node_nibble == nibble {
350                    Some(node_value.into())
351                } else {
352                    None
353                }
354            },
355            MerkleNode::Extension(node_nibble, node_value) => {
356                if nibble.starts_with(&node_nibble) {
357                    Self::get_by_value(database, cache,
358                                       nibble.split_at(node_nibble.len()).1.into(),
359                                       node_value.clone())
360                } else {
361                    None
362                }
363            },
364            MerkleNode::Branch(nodes, additional) => {
365                if nibble.len() == 0 {
366                    additional.clone().map(|v| v.into())
367                } else {
368                    let nibble_index: usize = nibble[0].into();
369                    let node = &nodes[nibble_index];
370                    Self::get_by_value(database, cache,
371                                       nibble.split_at(1).1.into(), node.clone())
372                }
373            },
374        }
375    }
376
377    pub fn get_raw<'a, 'b>(&'a self, key: &'b [u8]) -> Option<Vec<u8>> {
378        if self.is_empty() {
379            return None;
380        }
381
382        let nibble = nibble::from_key(key);
383        let dbv = match self.database.get(self.root) {
384            Some(val) => val,
385            None => return None,
386        };
387        let node = MerkleNode::decode(&Rlp::new(&dbv));
388        let mut change = Change::new(&self.database);
389        let cache = Cache::new();
390        let ret = Self::get_by_node(&mut change, &cache, nibble, node).map(|v| v.into());
391        debug_assert!(change.inserted().len() == 0 && change.freed().len() == 0);
392        ret
393    }
394
395    fn insert_by_value<'a, 'b: 'a>(
396        database: &mut Change<'a, D>, cache: &'a Cache,
397        nibble: NibbleVec, merkle: MerkleValue<'a>, value: &'a [u8]
398    ) -> MerkleValue<'a> {
399        match merkle {
400            MerkleValue::Empty => {
401                let mut node_map = HashMap::new();
402                node_map.insert(nibble, value);
403
404                let new_node = Self::build_node(database, &node_map);
405                Self::build_value(database, new_node)
406            },
407            MerkleValue::Full(ref sub_node) => {
408                let sub_node: &MerkleNode<'a> = sub_node.borrow();
409                let sub_node: MerkleNode<'a> = (*sub_node).clone();
410
411                let new_node = Self::insert_by_node(database, cache, nibble, sub_node, value);
412                Self::build_value(database, new_node)
413            },
414            MerkleValue::Hash(h) => {
415                let dbv = match database.get(h) {
416                    Some(val) => val,
417                    None => panic!(),
418                };
419                let node = cache.insert(h, dbv);
420                let new_node = Self::insert_by_node(database, cache, nibble, node, value);
421                Self::build_value(database, new_node)
422            }
423        }
424    }
425
426    fn insert_by_node<'a, 'b: 'a>(
427        database: &mut Change<'a, D>, cache: &'a Cache,
428        nibble: NibbleVec, node: MerkleNode<'a>, value: &'a [u8]
429    ) -> MerkleNode<'a> {
430        match node {
431            MerkleNode::Leaf(ref node_nibble, ref node_value) => {
432                let mut node_map = HashMap::new();
433                node_map.insert(node_nibble.clone(), node_value.clone());
434                node_map.insert(nibble, value);
435
436                Self::build_node(database, &node_map)
437            },
438            MerkleNode::Extension(ref node_nibble, ref node_value) => {
439                if nibble.starts_with(node_nibble) {
440                    MerkleNode::Extension(
441                        node_nibble.clone(),
442                        Self::insert_by_value(
443                            database, cache, nibble.split_at(node_nibble.len()).1.into(),
444                            node_value.clone(), value))
445                } else {
446                    let common = nibble::common(&nibble, &node_nibble);
447                    let rest_len = node_nibble.len() - common.len() - 1;
448                    debug_assert!(node_nibble.len() - common.len() > 0);
449                    debug_assert!(nibble.len() - common.len() > 0);
450                    let rest_at: usize = node_nibble[common.len()].into();
451                    let insert_at: usize = nibble[common.len()].into();
452
453                    let rest = if rest_len > 0 {
454                        let new_node = MerkleNode::Extension(
455                            node_nibble.split_at(common.len() + 1).1.into(),
456                            node_value.clone());
457                        Self::build_value(database, new_node)
458                    } else /* if rest_len == 0 */ {
459                        node_value.clone()
460                    };
461
462                    let branched_node = {
463                        let mut nodes = empty_nodes!();
464                        nodes[rest_at] = rest;
465                        nodes[insert_at] = Self::insert_by_value(
466                            database, cache, nibble.split_at(common.len() + 1).1.into(),
467                            MerkleValue::Empty, value);
468                        MerkleNode::Branch(nodes, None)
469                    };
470                    let branched = Self::build_value(database, branched_node.clone());
471
472                    if common.len() >= 1 {
473                        MerkleNode::Extension(common.into(), branched)
474                    } else /* if common.len() == 0 */ {
475                        branched_node
476                    }
477                }
478            },
479            MerkleNode::Branch(ref node_nodes, ref node_additional) => {
480                let mut nodes = Self::copy_nodes(node_nodes);
481                if nibble.len() == 0 {
482                    MerkleNode::Branch(nodes, Some(value))
483                } else {
484                    let nibble_index: usize = nibble[0].into();
485                    let prev = nodes[nibble_index].clone();
486                    nodes[nibble_index] = Self::insert_by_value(
487                        database, cache, nibble.split_at(1).1.into(), prev, value);
488                    MerkleNode::Branch(nodes, node_additional.clone())
489                }
490            },
491        }
492    }
493
494    pub fn insert_raw(&mut self, key: Vec<u8>, value: Vec<u8>) {
495        let key: &[u8] = key.as_ref();
496        let value: &[u8] = value.as_ref();
497
498        let (changeset, root_rlp) = {
499            let cache = Cache::new();
500            let mut change = Change::new(&self.database);
501            let node = if self.is_empty() {
502                let mut node_map = HashMap::new();
503                node_map.insert(nibble::from_key(key), value.clone());
504
505                Self::build_node(&mut change, &node_map)
506            } else {
507                let nibble = nibble::from_key(key);
508                let dbv = match self.database.get(self.root) {
509                    Some(val) => val,
510                    None => panic!(),
511                };
512                let node = cache.insert(self.root, dbv);
513                Self::insert_by_node(&mut change, &cache, nibble, node, value)
514            };
515            (ChangeSet::from(change), rlp::encode(&node).to_vec())
516        };
517        let hash = H256::from(Keccak256::digest(&root_rlp).as_slice());
518        changeset.drain(&mut self.database, true);
519        self.database.set(hash, root_rlp);
520
521        self.root = hash;
522    }
523
524    fn remove_by_value<'a, 'b: 'a>(
525        database: &mut Change<'a, D>, cache: &'a Cache,
526        nibble: NibbleVec, merkle: MerkleValue<'a>
527    ) -> MerkleValue<'a> {
528        match merkle {
529            MerkleValue::Empty => {
530                MerkleValue::Empty
531            },
532            MerkleValue::Full(ref sub_node) => {
533                let sub_node: &MerkleNode<'a> = sub_node.borrow();
534                let sub_node: MerkleNode<'a> = (*sub_node).clone();
535
536                let new_node = Self::remove_by_node(database, cache, nibble, sub_node);
537                if new_node.is_none() {
538                    MerkleValue::Empty
539                } else {
540                    let new_node = new_node.unwrap();
541                    Self::build_value(database, new_node)
542                }
543            },
544            MerkleValue::Hash(h) => {
545                let dbv = match database.get(h) {
546                    Some(val) => val,
547                    None => panic!(),
548                };
549                let node = cache.insert(h, dbv);
550                let new_node = Self::remove_by_node(database, cache, nibble, node);
551                if new_node.is_none() {
552                    MerkleValue::Empty
553                } else {
554                    let new_node = new_node.unwrap();
555                    Self::build_value(database, new_node)
556                }
557            },
558        }
559    }
560
561    fn collapse<'a, 'b: 'a>(
562        database: &mut Change<'a, D>, cache: &'a Cache, node: MerkleNode<'a>
563    ) -> MerkleNode<'a> {
564        fn find_subnode<'a: 'b, 'b, D: DatabaseGuard>(
565            database: &mut Change<'a, D>, cache: &'a Cache, value: MerkleValue<'b>
566        ) -> MerkleNode<'b> {
567            match value {
568                MerkleValue::Empty => panic!(),
569                MerkleValue::Hash(h) => {
570                    let dbv = match database.get(h) {
571                        Some(val) => val,
572                        None => panic!(),
573                    };
574                    cache.insert(h, dbv)
575                },
576                MerkleValue::Full(f) => {
577                    let t: &MerkleNode = &f;
578                    t.clone()
579                },
580            }
581        }
582
583        match node {
584            MerkleNode::Leaf(_, _) => panic!(), // Leaf does not collapse.
585            MerkleNode::Extension(node_nibble, node_value) => {
586                let subnode = find_subnode(database, cache, node_value.clone());
587
588                match subnode {
589                    MerkleNode::Leaf(mut sub_nibble, sub_value) => {
590                        let mut new_sub_nibble = node_nibble.clone();
591                        new_sub_nibble.append(&mut sub_nibble);
592                        MerkleNode::Leaf(new_sub_nibble, sub_value)
593                    },
594                    MerkleNode::Extension(mut sub_nibble, sub_value) => {
595                        let mut new_sub_nibble = node_nibble.clone();
596                        new_sub_nibble.append(&mut sub_nibble);
597                        Self::collapse(database, cache,
598                                       MerkleNode::Extension(new_sub_nibble, sub_value))
599                    },
600                    _ => MerkleNode::Extension(node_nibble, node_value),
601                }
602            },
603            MerkleNode::Branch(node_nodes, node_additional) => {
604                let value_count = node_additional.iter().count() +
605                    node_nodes.iter().filter(|v| v != &&MerkleValue::Empty).count();
606
607                if value_count == 0 {
608                    panic!()
609                } else if value_count > 1 {
610                    MerkleNode::Branch(node_nodes, node_additional)
611                } else if node_additional.is_some() /* value_count == 1 */ {
612                    MerkleNode::Leaf(NibbleVec::new(), node_additional.unwrap())
613                } else /* value_count == 1, value in nodes */ {
614                    let (value_index, value) = node_nodes
615                        .iter().enumerate().filter(|&(_, value)| {
616                            value != &MerkleValue::Empty
617                        }).next()
618                        .map(|(value_index, value)| (value_index, value.clone())).unwrap();
619                    let value_nibble: Nibble = value_index.into();
620
621                    let subnode = find_subnode(database, cache, value.clone());
622                    match subnode {
623                        MerkleNode::Leaf(mut sub_nibble, sub_value) => {
624                            sub_nibble.insert(0, value_nibble);
625                            MerkleNode::Leaf(sub_nibble, sub_value)
626                        },
627                        MerkleNode::Extension(mut sub_nibble, sub_value) => {
628                            sub_nibble.insert(0, value_nibble);
629                            Self::collapse(database, cache,
630                                           MerkleNode::Extension(sub_nibble, sub_value))
631                        },
632                        MerkleNode::Branch(sub_nodes, sub_additional) => {
633                            Self::collapse(database, cache,
634                                           MerkleNode::Extension(vec![value_nibble], value))
635                        },
636                    }
637                }
638            },
639        }
640    }
641
642    fn remove_by_node<'a, 'b: 'a>(
643        database: &mut Change<'a, D>, cache: &'a Cache,
644        nibble: NibbleVec, node: MerkleNode<'a>
645    ) -> Option<MerkleNode<'a>> {
646        match node {
647            MerkleNode::Leaf(ref node_nibble, ref node_value) => {
648                if *node_nibble == nibble {
649                    None
650                } else {
651                    Some(MerkleNode::Leaf(node_nibble.clone(), node_value.clone()))
652                }
653            },
654            MerkleNode::Extension(ref node_nibble, ref node_value) => {
655                if nibble.starts_with(node_nibble) {
656                    let value = Self::remove_by_value(
657                        database, cache,
658                        nibble.split_at(node_nibble.len()).1.into(),
659                        node_value.clone());
660                    Some(Self::collapse(database, cache,
661                                        MerkleNode::Extension(node_nibble.clone(), value)))
662                } else {
663                    Some(MerkleNode::Extension(node_nibble.clone(), node_value.clone()))
664                }
665            },
666            MerkleNode::Branch(ref node_nodes, ref node_additional) => {
667                let mut nodes = Self::copy_nodes(node_nodes);
668                let mut additional = node_additional.clone();
669
670                if nibble.len() > 0 {
671                    let nibble_index: usize = nibble[0].into();
672                    nodes[nibble_index] = Self::remove_by_value(
673                        database, cache,
674                        nibble.split_at(1).1.into(),
675                        nodes[nibble_index].clone());
676                } else {
677                    additional = None;
678                }
679
680                let value_count = additional.iter().count() +
681                    nodes.iter().filter(|v| v != &&MerkleValue::Empty).count();
682
683                if value_count == 0 {
684                    None
685                } else {
686                    Some(Self::collapse(database, cache, MerkleNode::Branch(nodes, additional)))
687                }
688            },
689        }
690    }
691
692    pub fn remove_raw<'a, 'b: 'a>(&'a mut self, key: &'b [u8]) {
693        if self.is_empty() {
694            return;
695        }
696
697        let (changeset, root_rlp) = {
698            let cache = Cache::new();
699            let mut change = Change::new(&self.database);
700            let nibble = nibble::from_key(key);
701            let dbv = match self.database.get(self.root) {
702                Some(val) => val,
703                None => panic!(),
704            };
705            let node = cache.insert(self.root, dbv);
706            let root_rlp = Self::remove_by_node(&mut change, &cache, nibble, node).map(|v| rlp::encode(&v).to_vec());
707            (ChangeSet::from(change), root_rlp)
708        };
709        changeset.drain(&mut self.database, true);
710        if root_rlp.is_some() {
711            let root_rlp = root_rlp.unwrap();
712            let hash = H256::from(Keccak256::digest(&root_rlp).as_slice());
713            self.database.set(hash, root_rlp);
714            self.root = hash;
715        } else {
716            self.root = empty_trie_hash!();
717        }
718    }
719}
720
721#[cfg(test)]
722mod tests {
723    use super::{DatabaseGuard, Trie};
724    use std::collections::HashMap;
725    use std::str::FromStr;
726    use std::cell::UnsafeCell;
727    use bigint::H256;
728    use hexutil::read_hex;
729
730    #[test]
731    fn trie_middle_leaf() {
732        let mut map = HashMap::new();
733        map.insert("key1aa".as_bytes().into(), "0123456789012345678901234567890123456789xxx".as_bytes().into());
734        map.insert("key1".as_bytes().into(), "0123456789012345678901234567890123456789Very_Long".as_bytes().into());
735        map.insert("key2bb".as_bytes().into(), "aval3".as_bytes().into());
736        map.insert("key2".as_bytes().into(), "short".as_bytes().into());
737        map.insert("key3cc".as_bytes().into(), "aval3".as_bytes().into());
738        map.insert("key3".as_bytes().into(), "1234567890123456789012345678901".as_bytes().into());
739
740        let mut database: HashMap<H256, Vec<u8>> = HashMap::new();
741        let mut trie: Trie<HashMap<H256, Vec<u8>>> = Trie::build(database, &map);
742
743        assert_eq!(trie.root(), H256::from_str("0xcb65032e2f76c48b82b5c24b3db8f670ce73982869d38cd39a624f23d62a9e89").unwrap());
744        assert_eq!(trie.get_raw("key2bb".as_bytes()), Some("aval3".as_bytes().into()));
745        assert_eq!(trie.get_raw("key2bbb".as_bytes()), None);
746        let prev_hash = trie.root();
747        trie.insert_raw("key2bbb".as_bytes().into(), "aval4".as_bytes().into());
748        assert_eq!(trie.get_raw("key2bbb".as_bytes()), Some("aval4".as_bytes().into()));
749        trie.remove_raw("key2bbb".as_bytes());
750        assert_eq!(trie.get_raw("key2bbb".as_bytes()), None);
751        assert_eq!(prev_hash, trie.root());
752    }
753
754    #[test]
755    fn insert_middle_leaf() {
756        let mut database: HashMap<H256, Vec<u8>> = HashMap::new();
757        let mut trie = Trie::empty(database);
758
759        trie.insert_raw("key1aa".as_bytes().into(),
760                        "0123456789012345678901234567890123456789xxx".as_bytes().into());
761        trie.insert_raw("key1".as_bytes().into(),
762                        "0123456789012345678901234567890123456789Very_Long".as_bytes().into());
763        trie.insert_raw("key2bb".as_bytes().into(),
764                        "aval3".as_bytes().into());
765        trie.insert_raw("key2".as_bytes().into(),
766                        "short".as_bytes().into());
767        trie.insert_raw("key3cc".as_bytes().into(),
768                        "aval3".as_bytes().into());
769        trie.insert_raw("key3".as_bytes().into(),
770                        "1234567890123456789012345678901".as_bytes().into());
771        assert_eq!(trie.root(), H256::from_str("0xcb65032e2f76c48b82b5c24b3db8f670ce73982869d38cd39a624f23d62a9e89").unwrap());
772    }
773
774    #[test]
775    fn insert_animals() {
776        let mut database: HashMap<H256, Vec<u8>> = HashMap::new();
777        let mut trie = Trie::empty(database);
778
779        trie.insert_raw("doe".as_bytes().into(),
780                        "reindeer".as_bytes().into());
781        trie.insert_raw("dog".as_bytes().into(),
782                        "puppy".as_bytes().into());
783        trie.insert_raw("dogglesworth".as_bytes().into(),
784                        "cat".as_bytes().into());
785
786        let mut all_key_values: HashMap<Vec<u8>, Vec<u8>> = HashMap::new();
787        all_key_values.insert("doe".as_bytes().into(), "reindeer".as_bytes().into());
788        all_key_values.insert("dog".as_bytes().into(), "puppy".as_bytes().into());
789        all_key_values.insert("dogglesworth".as_bytes().into(), "cat".as_bytes().into());
790
791        for (key, value) in trie.iter() {
792            assert_eq!(all_key_values.get(&key), Some(&value));
793            all_key_values.remove(&key);
794        }
795
796        assert_eq!(trie.root(), H256::from_str("0x8aad789dff2f538bca5d8ea56e8abe10f4c7ba3a5dea95fea4cd6e7c3a1168d3").unwrap());
797    }
798
799    #[test]
800    fn insert_single_item() {
801        let mut database: HashMap<H256, Vec<u8>> = HashMap::new();
802        let mut trie = Trie::empty(database);
803
804        trie.insert_raw("A".as_bytes().into(),
805                    "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".as_bytes().into());
806
807        assert_eq!(trie.root(), H256::from_str("0xd23786fb4a010da3ce639d66d5e904a11dbc02746d1ce25029e53290cabf28ab").unwrap());
808    }
809
810    #[test]
811    fn testy() {
812        let mut database: HashMap<H256, Vec<u8>> = HashMap::new();
813        let mut trie = Trie::empty(database);
814
815        trie.insert_raw("test".as_bytes().into(),
816                        "test".as_bytes().into());
817        trie.insert_raw("te".as_bytes().into(),
818                        "testy".as_bytes().into());
819
820        assert_eq!(trie.root(), H256::from_str("0x8452568af70d8d140f58d941338542f645fcca50094b20f3c3d8c3df49337928").unwrap());
821    }
822
823    #[test]
824    fn sub_genesis() {
825        let mut database: HashMap<H256, Vec<u8>> = HashMap::new();
826        let mut trie = Trie::empty(database);
827
828        let k1 = read_hex("0x204188718653cd7e50f3fd51a820db66112517ca190c637e7cdd80782d56").unwrap();
829        let v1 = vec![248, 78, 128, 138, 21, 45, 2, 199, 225, 74, 246, 128, 0, 0, 160, 86, 232, 31, 23, 27, 204, 85, 166, 255, 131, 69, 230, 146, 192, 248, 110, 91, 72, 224, 27, 153, 108, 173, 192, 1, 98, 47, 181, 227, 99, 180, 33, 160, 197, 210, 70, 1, 134, 247, 35, 60, 146, 126, 125, 178, 220, 199, 3, 192, 229, 0, 182, 83, 202, 130, 39, 59, 123, 250, 216, 4, 93, 133, 164, 112];
830        let k2 = read_hex("0xa390953f116afb00f89fbedb2f8e77297e4e7e1749e2ef0e32e17808e4ad").unwrap();
831        let v2 = vec![248, 77, 128, 137, 108, 107, 147, 91, 139, 189, 64, 0, 0, 160, 86, 232, 31, 23, 27, 204, 85, 166, 255, 131, 69, 230, 146, 192, 248, 110, 91, 72, 224, 27, 153, 108, 173, 192, 1, 98, 47, 181, 227, 99, 180, 33, 160, 197, 210, 70, 1, 134, 247, 35, 60, 146, 126, 125, 178, 220, 199, 3, 192, 229, 0, 182, 83, 202, 130, 39, 59, 123, 250, 216, 4, 93, 133, 164, 112];
832
833        trie.insert_raw(k1, v1);
834        trie.insert_raw(k2, v2);
835
836        assert_eq!(trie.root(), H256::from_str("bcb5ffb5c6c3e43ef07550fa30af86d66b4015ee3f64aaf70cd0bf8fcc60a9c6").unwrap());
837    }
838
839    #[test]
840    fn trie_insert() {
841        let mut map = HashMap::new();
842
843        let mut database: HashMap<H256, Vec<u8>> = HashMap::new();
844        let mut trie: Trie<HashMap<H256, Vec<u8>>> = Trie::build(database, &map);
845
846        trie.insert_raw("foo".as_bytes().into(), "bar".as_bytes().into());
847        trie.insert_raw("food".as_bytes().into(), "bass".as_bytes().into());
848
849        assert_eq!(trie.root(), H256::from_str("0x17beaa1648bafa633cda809c90c04af50fc8aed3cb40d16efbddee6fdf63c4c3").unwrap());
850    }
851
852    #[test]
853    fn trie_delete() {
854        let mut map = HashMap::new();
855
856        let mut database: HashMap<H256, Vec<u8>> = HashMap::new();
857        let mut trie: Trie<HashMap<H256, Vec<u8>>> = Trie::build(database, &map);
858
859        trie.insert_raw("fooa".as_bytes().into(), "bar".as_bytes().into());
860        trie.insert_raw("food".as_bytes().into(), "bass".as_bytes().into());
861        let prev_hash = trie.root();
862        trie.insert_raw("fooc".as_bytes().into(), "basss".as_bytes().into());
863        trie.remove_raw("fooc".as_bytes());
864        assert_eq!(trie.root(), prev_hash);
865    }
866
867    #[test]
868    fn trie_empty() {
869        let mut map = HashMap::new();
870
871        let mut database: HashMap<H256, Vec<u8>> = HashMap::new();
872        let mut trie: Trie<HashMap<H256, Vec<u8>>> = Trie::build(database, &map);
873
874        assert_eq!(H256::from("0x56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421"),
875                   trie.root());
876    }
877}