Delete values immediately from DB overlay (#1631)

* Delete values immediately from DB overlay

* Match remove behaviour; test

* Warning fixed

util/src/journaldb/overlayrecentdb.rs

@@ -295,15 +295,14 @@ impl JournalDB for OverlayRecentDB {
 				}
 				// update the overlay
 				for k in overlay_deletions {
-					journal_overlay.backing_overlay.remove(&OverlayRecentDB::to_short_key(&k));
+					journal_overlay.backing_overlay.remove_and_purge(&OverlayRecentDB::to_short_key(&k));
 				}
 				// apply canon deletions
 				for k in canon_deletions {
 					if !journal_overlay.backing_overlay.contains(&OverlayRecentDB::to_short_key(&k)) {
 						try!(batch.delete(&k));
 					}
 				}
-				journal_overlay.backing_overlay.purge();
 			}
 			journal_overlay.journal.remove(&end_era);
 		}

util/src/memorydb.rs

@@ -24,6 +24,7 @@ use hashdb::*;
 use heapsize::*;
 use std::mem;
 use std::collections::HashMap;
+use std::collections::hash_map::Entry;
 use std::default::Default;
 
 #[derive(Debug,Clone)]
@@ -162,6 +163,24 @@ impl MemoryDB {
 	pub fn mem_used(&self) -> usize {
 		self.data.heap_size_of_children()
 	}
+
+	/// Remove an element and delete it from storage if reference count reaches zero.
+	pub fn remove_and_purge(&mut self, key: &H256) {
+		if key == &SHA3_NULL_RLP {
+			return;
+		}
+		match self.data.entry(key.clone()) {
+			Entry::Occupied(mut entry) =>
+				if entry.get().1 == 1 {
+					entry.remove();
+				} else {
+					entry.get_mut().1 -= 1;
+				},
+			Entry::Vacant(entry) => {
+				entry.insert((Bytes::new(), -1));
+			}
+		}
+	}
 }
 
 static NULL_RLP_STATIC: [u8; 1] = [0x80; 1];
@@ -269,3 +288,28 @@ fn memorydb_denote() {
 
 	assert_eq!(m.get(&hash).unwrap(), b"Hello world!");
 }
+
+#[test]
+fn memorydb_remove_and_purge() {
+	let hello_bytes = b"Hello world!";
+	let hello_key = hello_bytes.sha3();
+
+	let mut m = MemoryDB::new();
+	m.remove(&hello_key);
+	assert_eq!(m.raw(&hello_key).unwrap().1, -1);
+	m.purge();
+	assert_eq!(m.raw(&hello_key).unwrap().1, -1);
+	m.insert(hello_bytes);
+	assert_eq!(m.raw(&hello_key).unwrap().1, 0);
+	m.purge();
+	assert_eq!(m.raw(&hello_key), None);
+
+	let mut m = MemoryDB::new();
+	m.remove_and_purge(&hello_key);
+	assert_eq!(m.raw(&hello_key).unwrap().1, -1);
+	m.insert(hello_bytes);
+	m.insert(hello_bytes);
+	assert_eq!(m.raw(&hello_key).unwrap().1, 1);
+	m.remove_and_purge(&hello_key);
+	assert_eq!(m.raw(&hello_key), None);
+}