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 {
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 {
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!(), 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() {
612 MerkleNode::Leaf(NibbleVec::new(), node_additional.unwrap())
613 } else {
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}