Use prefixed keys for trie node. (#2130)

* Account for pending insertions when pruning

* Prefixed trie storage

* Comments

* Prefixed trie storage

* Fixed tests

* Fixed tests

* Bumped runtime version

* Bumped runtime version again

substrate/core/state-machine/src/backend.rs

@@ -112,9 +112,9 @@ impl Consolidate for Vec<(Option<Vec<u8>>, Vec<u8>, Option<Vec<u8>>)> {
 	}
 }
 
-impl<H: Hasher> Consolidate for MemoryDB<H> {
+impl<H: Hasher, KF: trie::KeyFunction<H>> Consolidate for trie::GenericMemoryDB<H, KF> {
 	fn consolidate(&mut self, other: Self) {
-		MemoryDB::consolidate(self, other)
+		trie::GenericMemoryDB::consolidate(self, other)
 	}
 }
 

substrate/core/state-machine/src/changes_trie/changes_iterator.rs

@@ -115,7 +115,7 @@ pub fn key_changes_proof_check<S: RootsStorage<H>, H: Hasher>(
 
 	let mut proof_db = MemoryDB::<H>::default();
 	for item in proof {
-		proof_db.insert(&item);
+		proof_db.insert(&[], &item);
 	}
 
 	let proof_db = InMemoryStorage::with_db(proof_db);

substrate/core/state-machine/src/changes_trie/mod.rs

@@ -77,7 +77,7 @@ pub trait RootsStorage<H: Hasher>: Send + Sync {
 /// Changes trie storage. Provides access to trie roots and trie nodes.
 pub trait Storage<H: Hasher>: RootsStorage<H> {
 	/// Get a trie node.
-	fn get(&self, key: &H::Out) -> Result<Option<DBValue>, String>;
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Result<Option<DBValue>, String>;
 }
 
 /// Changes trie configuration.

substrate/core/state-machine/src/changes_trie/storage.rs

@@ -92,7 +92,7 @@ impl<H: Hasher> InMemoryStorage<H> where H::Out: HeapSizeOf {
 	pub fn remove_from_storage(&self, keys: &HashSet<H::Out>) {
 		let mut data = self.data.write();
 		for key in keys {
-			data.mdb.remove_and_purge(key);
+			data.mdb.remove_and_purge(key, &[]);
 		}
 	}
 
@@ -116,8 +116,8 @@ impl<H: Hasher> RootsStorage<H> for InMemoryStorage<H> where H::Out: HeapSizeOf
 }
 
 impl<H: Hasher> Storage<H> for InMemoryStorage<H> where H::Out: HeapSizeOf {
-	fn get(&self, key: &H::Out) -> Result<Option<DBValue>, String> {
-		MemoryDB::<H>::get(&self.data.read().mdb, key)
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Result<Option<DBValue>, String> {
+		MemoryDB::<H>::get(&self.data.read().mdb, key, prefix)
 	}
 }
 
@@ -128,7 +128,9 @@ impl<'a, H: Hasher, S: 'a + Storage<H>> TrieBackendAdapter<'a, H, S> {
 }
 
 impl<'a, H: Hasher, S: 'a + Storage<H>> TrieBackendStorage<H> for TrieBackendAdapter<'a, H, S> {
-	fn get(&self, key: &H::Out) -> Result<Option<DBValue>, String> {
-		self.storage.get(key)
+	type Overlay = MemoryDB<H>;
+
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Result<Option<DBValue>, String> {
+		self.storage.get(key, prefix)
 	}
 }

substrate/core/state-machine/src/proving_backend.rs

@@ -21,7 +21,7 @@ use log::debug;
 use hash_db::Hasher;
 use heapsize::HeapSizeOf;
 use hash_db::HashDB;
-use trie::{Recorder, MemoryDB, TrieError, default_child_trie_root, read_trie_value_with, read_child_trie_value_with, record_all_keys};
+use trie::{Recorder, MemoryDB, PrefixedMemoryDB, TrieError, default_child_trie_root, read_trie_value_with, read_child_trie_value_with, record_all_keys};
 use crate::trie_backend::TrieBackend;
 use crate::trie_backend_essence::{Ephemeral, TrieBackendEssence, TrieBackendStorage};
 use crate::{Error, ExecutionError, Backend};
@@ -40,7 +40,7 @@ impl<'a, S, H> ProvingBackendEssence<'a, S, H>
 		H::Out: HeapSizeOf,
 {
 	pub fn storage(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>, String> {
-		let mut read_overlay = MemoryDB::default();
+		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral::new(
 			self.backend.backend_storage(),
 			&mut read_overlay,
@@ -54,7 +54,7 @@ impl<'a, S, H> ProvingBackendEssence<'a, S, H>
 	pub fn child_storage(&mut self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, String> {
 		let root = self.storage(storage_key)?.unwrap_or(default_child_trie_root::<H>(storage_key));
 
-		let mut read_overlay = MemoryDB::default();
+		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral::new(
 			self.backend.backend_storage(),
 			&mut read_overlay,
@@ -66,7 +66,7 @@ impl<'a, S, H> ProvingBackendEssence<'a, S, H>
 	}
 
 	pub fn record_all_keys(&mut self) {
-		let mut read_overlay = MemoryDB::default();
+		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral::new(
 			self.backend.backend_storage(),
 			&mut read_overlay,
@@ -116,8 +116,8 @@ impl<'a, S, H> Backend<H> for ProvingBackend<'a, S, H>
 		H::Out: Ord + HeapSizeOf,
 {
 	type Error = String;
-	type Transaction = MemoryDB<H>;
-	type TrieBackendStorage = MemoryDB<H>;
+	type Transaction = S::Overlay;
+	type TrieBackendStorage = PrefixedMemoryDB<H>;
 
 	fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
 		ProvingBackendEssence {
@@ -151,7 +151,7 @@ impl<'a, S, H> Backend<H> for ProvingBackend<'a, S, H>
 		self.backend.keys(prefix)
 	}
 
-	fn storage_root<I>(&self, delta: I) -> (H::Out, MemoryDB<H>)
+	fn storage_root<I>(&self, delta: I) -> (H::Out, Self::Transaction)
 		where I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>
 	{
 		self.backend.storage_root(delta)
@@ -181,7 +181,7 @@ where
 {
 	let db = create_proof_check_backend_storage(proof);
 
-	if !db.contains(&root) {
+	if !db.contains(&root, &[]) {
 		return Err(Box::new(ExecutionError::InvalidProof) as Box<Error>);
 	}
 
@@ -198,7 +198,7 @@ where
 {
 	let mut db = MemoryDB::default();
 	for item in proof {
-		db.insert(&item);
+		db.insert(&[], &item);
 	}
 	db
 }
@@ -210,7 +210,7 @@ mod tests {
 	use super::*;
 	use primitives::{Blake2Hasher};
 
-	fn test_proving<'a>(trie_backend: &'a TrieBackend<MemoryDB<Blake2Hasher>, Blake2Hasher>) -> ProvingBackend<'a, MemoryDB<Blake2Hasher>, Blake2Hasher> {
+	fn test_proving<'a>(trie_backend: &'a TrieBackend<PrefixedMemoryDB<Blake2Hasher>, Blake2Hasher>) -> ProvingBackend<'a, PrefixedMemoryDB<Blake2Hasher>, Blake2Hasher> {
 		ProvingBackend::new(trie_backend)
 	}
 

substrate/core/state-machine/src/trie_backend.rs

@@ -19,7 +19,7 @@
 use log::{warn, debug};
 use hash_db::Hasher;
 use heapsize::HeapSizeOf;
-use trie::{TrieDB, TrieError, Trie, MemoryDB, delta_trie_root, default_child_trie_root, child_delta_trie_root};
+use trie::{TrieDB, TrieError, Trie, delta_trie_root, default_child_trie_root, child_delta_trie_root};
 use crate::trie_backend_essence::{TrieBackendEssence, TrieBackendStorage, Ephemeral};
 use crate::Backend;
 
@@ -63,7 +63,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 	H::Out: Ord + HeapSizeOf,
 {
 	type Error = String;
-	type Transaction = MemoryDB<H>;
+	type Transaction = S::Overlay;
 	type TrieBackendStorage = S;
 
 	fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
@@ -83,7 +83,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 	}
 
 	fn pairs(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
-		let mut read_overlay = MemoryDB::default();
+		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral::new(self.essence.backend_storage(), &mut read_overlay);
 
 		let collect_all = || -> Result<_, Box<TrieError<H::Out>>> {
@@ -107,7 +107,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 	}
 
 	fn keys(&self, prefix: &Vec<u8>) -> Vec<Vec<u8>> {
-		let mut read_overlay = MemoryDB::default();
+		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral::new(self.essence.backend_storage(), &mut read_overlay);
 
 		let collect_all = || -> Result<_, Box<TrieError<H::Out>>> {
@@ -126,10 +126,10 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 		collect_all().map_err(|e| debug!(target: "trie", "Error extracting trie keys: {}", e)).unwrap_or_default()
 	}
 
-	fn storage_root<I>(&self, delta: I) -> (H::Out, MemoryDB<H>)
+	fn storage_root<I>(&self, delta: I) -> (H::Out, S::Overlay)
 		where I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>
 	{
-		let mut write_overlay = MemoryDB::default();
+		let mut write_overlay = S::Overlay::default();
 		let mut root = *self.essence.root();
 
 		{
@@ -154,7 +154,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 	{
 		let default_root = default_child_trie_root::<H>(storage_key);
 
-		let mut write_overlay = MemoryDB::default();
+		let mut write_overlay = S::Overlay::default();
 		let mut root = match self.storage(storage_key) {
 			Ok(value) => value.unwrap_or(default_child_trie_root::<H>(storage_key)),
 			Err(e) => {
@@ -189,12 +189,12 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 pub mod tests {
 	use std::collections::HashSet;
 	use primitives::{Blake2Hasher, H256};
-	use trie::{TrieMut, TrieDBMut};
+	use trie::{TrieMut, TrieDBMut, PrefixedMemoryDB};
 	use super::*;
 
-	fn test_db() -> (MemoryDB<Blake2Hasher>, H256) {
+	fn test_db() -> (PrefixedMemoryDB<Blake2Hasher>, H256) {
 		let mut root = H256::default();
-		let mut mdb = MemoryDB::<Blake2Hasher>::default();
+		let mut mdb = PrefixedMemoryDB::<Blake2Hasher>::default();
 		{
 			let mut trie = TrieDBMut::new(&mut mdb, &mut root);
 			trie.insert(b"key", b"value").expect("insert failed");
@@ -208,7 +208,7 @@ pub mod tests {
 		(mdb, root)
 	}
 
-	pub(crate) fn test_trie() -> TrieBackend<MemoryDB<Blake2Hasher>, Blake2Hasher> {
+	pub(crate) fn test_trie() -> TrieBackend<PrefixedMemoryDB<Blake2Hasher>, Blake2Hasher> {
 		let (mdb, root) = test_db();
 		TrieBackend::new(mdb, root)
 	}
@@ -230,8 +230,8 @@ pub mod tests {
 
 	#[test]
 	fn pairs_are_empty_on_empty_storage() {
-		assert!(TrieBackend::<MemoryDB<Blake2Hasher>, Blake2Hasher>::new(
-			MemoryDB::default(),
+		assert!(TrieBackend::<PrefixedMemoryDB<Blake2Hasher>, Blake2Hasher>::new(
+			PrefixedMemoryDB::default(),
 			Default::default(),
 		).pairs().is_empty());
 	}

substrate/core/state-machine/src/trie_backend_essence.rs

@@ -22,13 +22,14 @@ use std::sync::Arc;
 use log::{debug, warn};
 use hash_db::{self, Hasher};
 use heapsize::HeapSizeOf;
-use trie::{TrieDB, Trie, MemoryDB, DBValue, TrieError, default_child_trie_root, read_trie_value, read_child_trie_value, for_keys_in_child_trie};
+use trie::{TrieDB, Trie, MemoryDB, PrefixedMemoryDB, DBValue, TrieError, default_child_trie_root, read_trie_value, read_child_trie_value, for_keys_in_child_trie};
 use crate::changes_trie::Storage as ChangesTrieStorage;
+use crate::backend::Consolidate;
 
 /// Patricia trie-based storage trait.
 pub trait Storage<H: Hasher>: Send + Sync {
 	/// Get a trie node.
-	fn get(&self, key: &H::Out) -> Result<Option<DBValue>, String>;
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Result<Option<DBValue>, String>;
 }
 
 /// Patricia trie-based pairs storage essence.
@@ -63,7 +64,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 
 	/// Get the value of storage at given key.
 	pub fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, String> {
-		let mut read_overlay = MemoryDB::default();
+		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral {
 			storage: &self.storage,
 			overlay: &mut read_overlay,
@@ -78,7 +79,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 	pub fn child_storage(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, String> {
 		let root = self.storage(storage_key)?.unwrap_or(default_child_trie_root::<H>(storage_key));
 
-		let mut read_overlay = MemoryDB::default();
+		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral {
 			storage: &self.storage,
 			overlay: &mut read_overlay,
@@ -99,7 +100,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 			}
 		};
 
-		let mut read_overlay = MemoryDB::default();
+		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral {
 			storage: &self.storage,
 			overlay: &mut read_overlay,
@@ -112,7 +113,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 
 	/// Execute given closure for all keys starting with prefix.
 	pub fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], mut f: F) {
-		let mut read_overlay = MemoryDB::default();
+		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral {
 			storage: &self.storage,
 			overlay: &mut read_overlay,
@@ -145,7 +146,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 
 pub(crate) struct Ephemeral<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> {
 	storage: &'a S,
-	overlay: &'a mut MemoryDB<H>,
+	overlay: &'a mut S::Overlay,
 }
 
 impl<'a,
@@ -171,7 +172,7 @@ impl<'a,
 }
 
 impl<'a, S: TrieBackendStorage<H>, H: Hasher> Ephemeral<'a, S, H> {
-	pub fn new(storage: &'a S, overlay: &'a mut MemoryDB<H>) -> Self {
+	pub fn new(storage: &'a S, overlay: &'a mut S::Overlay) -> Self {
 		Ephemeral {
 			storage,
 			overlay,
@@ -187,10 +188,10 @@ impl<'a,
 	where H::Out: HeapSizeOf
 {
 	fn get(&self, key: &H::Out) -> Option<DBValue> {
-		if let Some(val) = hash_db::PlainDB::get(self.overlay, key) {
+		if let Some(val) = hash_db::HashDB::get(self.overlay, key, &[]) {
 			Some(val)
 		} else {
-			match self.storage.get(&key) {
+			match self.storage.get(&key, &[]) {
 				Ok(x) => x,
 				Err(e) => {
 					warn!(target: "trie", "Failed to read from DB: {}", e);
@@ -201,15 +202,15 @@ impl<'a,
 	}
 
 	fn contains(&self, key: &H::Out) -> bool {
-		hash_db::PlainDB::get(self, key).is_some()
+		hash_db::HashDB::get(self, key, &[]).is_some()
 	}
 
 	fn emplace(&mut self, key: H::Out, value: DBValue) {
-		hash_db::PlainDB::emplace(self.overlay, key, value)
+		hash_db::HashDB::emplace(self.overlay, key, &[], value)
 	}
 
 	fn remove(&mut self, key: &H::Out) {
-		hash_db::PlainDB::remove(self.overlay, key)
+		hash_db::HashDB::remove(self.overlay, key, &[])
 	}
 }
 
@@ -231,11 +232,11 @@ impl<'a,
 	for Ephemeral<'a, S, H>
 	where H::Out: HeapSizeOf
 {
-	fn get(&self, key: &H::Out) -> Option<DBValue> {
-		if let Some(val) = hash_db::HashDB::get(self.overlay, key) {
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Option<DBValue> {
+		if let Some(val) = hash_db::HashDB::get(self.overlay, key, prefix) {
 			Some(val)
 		} else {
-			match self.storage.get(&key) {
+			match self.storage.get(&key, prefix) {
 				Ok(x) => x,
 				Err(e) => {
 					warn!(target: "trie", "Failed to read from DB: {}", e);
@@ -245,20 +246,20 @@ impl<'a,
 		}
 	}
 
-	fn contains(&self, key: &H::Out) -> bool {
-		hash_db::HashDB::get(self, key).is_some()
+	fn contains(&self, key: &H::Out, prefix: &[u8]) -> bool {
+		hash_db::HashDB::get(self, key, prefix).is_some()
 	}
 
-	fn insert(&mut self, value: &[u8]) -> H::Out {
-		hash_db::HashDB::insert(self.overlay, value)
+	fn insert(&mut self, prefix: &[u8], value: &[u8]) -> H::Out {
+		hash_db::HashDB::insert(self.overlay, prefix, value)
 	}
 
-	fn emplace(&mut self, key: H::Out, value: DBValue) {
-		hash_db::HashDB::emplace(self.overlay, key, value)
+	fn emplace(&mut self, key: H::Out, prefix: &[u8], value: DBValue) {
+		hash_db::HashDB::emplace(self.overlay, key, prefix, value)
 	}
 
-	fn remove(&mut self, key: &H::Out) {
-		hash_db::HashDB::remove(self.overlay, key)
+	fn remove(&mut self, key: &H::Out, prefix: &[u8]) {
+		hash_db::HashDB::remove(self.overlay, key, prefix)
 	}
 }
 
@@ -269,33 +270,49 @@ impl<'a,
 	for Ephemeral<'a, S, H>
 	where H::Out: HeapSizeOf
 {
-	fn get(&self, key: &H::Out) -> Option<DBValue> { hash_db::HashDB::get(self, key) }
-	fn contains(&self, key: &H::Out) -> bool { hash_db::HashDB::contains(self, key) }
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Option<DBValue> { hash_db::HashDB::get(self, key, prefix) }
+	fn contains(&self, key: &H::Out, prefix: &[u8]) -> bool { hash_db::HashDB::contains(self, key, prefix) }
 }
 
 /// Key-value pairs storage that is used by trie backend essence.
 pub trait TrieBackendStorage<H: Hasher>: Send + Sync {
+	/// Type of in-memory overlay.
+	type Overlay: hash_db::HashDB<H, DBValue> + Default + Consolidate;
 	/// Get the value stored at key.
-	fn get(&self, key: &H::Out) -> Result<Option<DBValue>, String>;
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Result<Option<DBValue>, String>;
 }
 
 // This implementation is used by normal storage trie clients.
 impl<H: Hasher> TrieBackendStorage<H> for Arc<Storage<H>> {
-	fn get(&self, key: &H::Out) -> Result<Option<DBValue>, String> {
-		Storage::<H>::get(self.deref(), key)
+	type Overlay = PrefixedMemoryDB<H>;
+
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Result<Option<DBValue>, String> {
+		Storage::<H>::get(self.deref(), key, prefix)
 	}
 }
 
 // This implementation is used by test storage trie clients.
+impl<H: Hasher> TrieBackendStorage<H> for PrefixedMemoryDB<H> {
+	type Overlay = PrefixedMemoryDB<H>;
+
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Result<Option<DBValue>, String> {
+		Ok(hash_db::HashDB::get(self, key, prefix))
+	}
+}
+
 impl<H: Hasher> TrieBackendStorage<H> for MemoryDB<H> {
-	fn get(&self, key: &H::Out) -> Result<Option<DBValue>, String> {
-		Ok(hash_db::PlainDB::get(self, key))
+	type Overlay = MemoryDB<H>;
+
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Result<Option<DBValue>, String> {
+		Ok(hash_db::HashDB::get(self, key, prefix))
 	}
 }
 
 // This implementation is used by changes trie clients.
 impl<'a, S, H: Hasher> TrieBackendStorage<H> for &'a S where S: ChangesTrieStorage<H> {
-	fn get(&self, key: &H::Out) -> Result<Option<DBValue>, String> {
-		ChangesTrieStorage::<H>::get(*self, key)
+	type Overlay = MemoryDB<H>;
+
+	fn get(&self, key: &H::Out, prefix: &[u8]) -> Result<Option<DBValue>, String> {
+		ChangesTrieStorage::<H>::get(*self, key, prefix)
 	}
 }