paritytech · Dec 30, 2020 · Dec 29, 2020 · Dec 29, 2020 · Dec 29, 2020 · Dec 29, 2020
diff --git a/primitives/state-machine/src/in_memory_backend.rs b/primitives/state-machine/src/in_memory_backend.rs
@@ -17,56 +17,21 @@
 //! State machine in memory backend.
 
 use crate::{
-	StorageKey, StorageValue, StorageCollection,
-	trie_backend::TrieBackend,
+	StorageKey, StorageValue, StorageCollection, trie_backend::TrieBackend, backend::Backend,
 };
-use std::{collections::{BTreeMap, HashMap}};
+use std::collections::{BTreeMap, HashMap};
 use hash_db::Hasher;
-use sp_trie::{
-	MemoryDB, TrieMut,
-	trie_types::TrieDBMut,
-};
+use sp_trie::{MemoryDB, empty_trie_root, Layout};
 use codec::Codec;
 use sp_core::storage::{ChildInfo, Storage};
 
-/// Insert input pairs into memory db.
-fn insert_into_memory_db<H, I>(mut root: H::Out, mdb: &mut MemoryDB<H>, input: I) -> H::Out
-where
-	H: Hasher,
-	I: IntoIterator<Item=(StorageKey, Option<StorageValue>)>,
-{
-	{
-		let mut trie = if root == Default::default() {
-			TrieDBMut::<H>::new(mdb, &mut root)
-		} else {
-			TrieDBMut::<H>::from_existing(mdb, &mut root).unwrap()
-		};
-		for (key, value) in input {
-			if let Err(e) = match value {
-				Some(value) => {
-					trie.insert(&key, &value)
-				},
-				None => {
-					trie.remove(&key)
-				},
-			}  {
-				panic!("Failed to write to trie: {}", e);
-			}
-		}
-		trie.commit();
-	}
-	root
-}
-
 /// Create a new empty instance of in-memory backend.
 pub fn new_in_mem<H: Hasher>() -> TrieBackend<MemoryDB<H>, H>
 where
 	H::Out: Codec + Ord,
 {
 	let db = MemoryDB::default();
-	let mut backend = TrieBackend::new(db, Default::default());
-	backend.insert(std::iter::empty());
-	backend
+	TrieBackend::new(db, empty_trie_root::<Layout<H>>())
 }
 
 impl<H: Hasher> TrieBackend<MemoryDB<H>, H>
@@ -92,32 +57,17 @@ where
 		&mut self,
 		changes: T,
 	) {
-		let mut new_child_roots = Vec::new();
-		let mut root_map = None;
-		let root = self.root().clone();
-		for (child_info, map) in changes {
-			if let Some(child_info) = child_info.as_ref() {
-				let prefix_storage_key = child_info.prefixed_storage_key();
-				let ch = insert_into_memory_db::<H, _>(root, self.backend_storage_mut(), map.clone().into_iter());
-				new_child_roots.push((prefix_storage_key.into_inner(), Some(ch.as_ref().into())));
-			} else {
-				root_map = Some(map);
-			}
-		}
+		let (top, child) = changes.into_iter().partition::<Vec<_>, _>(|v| v.0.is_none());
+		let (root, transaction) = self.full_storage_root(
+			top.iter().map(|(_, v)| v).flatten().map(|(k, v)| (&k[..], v.as_deref())),
+			child.iter()
+				.filter_map(|v|
+					v.0.as_ref().map(|c| (c, v.1.iter().map(|(k, v)| (&k[..], v.as_deref()))))
+				),
+		);
 
-		let root = match root_map {
-			Some(map) => insert_into_memory_db::<H, _>(
-				root,
-				self.backend_storage_mut(),
-				map.into_iter().chain(new_child_roots.into_iter()),
-			),
-			None => insert_into_memory_db::<H, _>(
-				root,
-				self.backend_storage_mut(),
-				new_child_roots.into_iter(),
-			),
-		};
 		self.essence.set_root(root);
+		self.backend_storage_mut().consolidate(transaction);
 	}
 
 	/// Merge trie nodes into this backend.
@@ -127,6 +77,12 @@ where
 		Self::new(clone, root)
 	}
 
+	/// Apply the given transaction to this backend and set the root to the given value.
+	pub fn apply_transaction(&mut self, root: H::Out, transaction: MemoryDB<H>) {
+		self.backend_storage_mut().consolidate(transaction);
+		self.essence.set_root(root);
+	}
+
 	/// Compare with another in-memory backend.
 	pub fn eq(&self, other: &Self) -> bool {
 		self.root() == other.root()
@@ -158,7 +114,9 @@ where
 {
 	fn from(inner: HashMap<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>>) -> Self {
 		let mut backend = new_in_mem();
-		backend.insert(inner.into_iter().map(|(k, m)| (k, m.into_iter().map(|(k, v)| (k, Some(v))).collect())));
+		backend.insert(
+			inner.into_iter().map(|(k, m)| (k, m.into_iter().map(|(k, v)| (k, Some(v))).collect())),
+		);
 		backend
 	}
 }
@@ -232,4 +190,16 @@ mod tests {
 		let storage_key = child_info.prefixed_storage_key();
 		assert!(trie_backend.storage(storage_key.as_slice()).unwrap().is_some());
 	}
+
+	#[test]
+	fn insert_multiple_times_child_data_works() {
+		let mut storage = new_in_mem::<BlakeTwo256>();
+		let child_info = ChildInfo::new_default(b"1");
+
+		storage.insert(vec![(Some(child_info.clone()), vec![(b"2".to_vec(), Some(b"3".to_vec()))])]);
+		storage.insert(vec![(Some(child_info.clone()), vec![(b"1".to_vec(), Some(b"3".to_vec()))])]);
+
+		assert_eq!(storage.child_storage(&child_info, &b"2"[..]), Ok(Some(b"3".to_vec())));
+		assert_eq!(storage.child_storage(&child_info, &b"1"[..]), Ok(Some(b"3".to_vec())));
+	}
 }
diff --git a/primitives/state-machine/src/testing.rs b/primitives/state-machine/src/testing.rs
@@ -153,8 +153,11 @@ impl<H: Hasher, N: ChangesTrieBlockNumber> TestExternalities<H, N>
 		&mut self.changes_trie_storage
 	}
 
-	/// Return a new backend with all pending value.
-	pub fn commit_all(&self) -> InMemoryBackend<H> {
+	/// Return a new backend with all pending changes.
+	///
+	/// In contrast to [`commit_all`](Self::commit_all) this will not panic if there are open
+	/// transactions.
+	fn as_backend(&self) -> InMemoryBackend<H> {
 		let top: Vec<_> = self.overlay.changes()
 			.map(|(k, v)| (k.clone(), v.value().cloned()))
 			.collect();
@@ -172,6 +175,23 @@ impl<H: Hasher, N: ChangesTrieBlockNumber> TestExternalities<H, N>
 		self.backend.update(transaction)
 	}
 
+	/// Commit all pending changes to the underlying backend.
+	///
+	/// # Panic
+	///
+	/// This will panic if there are still open transactions.
+	pub fn commit_all(&mut self) -> Result<(), String> {
+		let changes = self.overlay.drain_storage_changes(
+			&self.backend,
+			None,
+			Default::default(),
+			&mut Default::default(),
+		)?;
+
+		self.backend.apply_transaction(changes.transaction_storage_root, changes.transaction);
+		Ok(())
+	}
+
 	/// Execute the given closure while `self` is set as externalities.
 	///
 	/// Returns the result of the given closure.
@@ -209,7 +229,7 @@ impl<H: Hasher, N: ChangesTrieBlockNumber> PartialEq for TestExternalities<H, N>
 	/// This doesn't test if they are in the same state, only if they contains the
 	/// same data at this state
 	fn eq(&self, other: &TestExternalities<H, N>) -> bool {
-		self.commit_all().eq(&other.commit_all())
+		self.as_backend().eq(&other.as_backend())
 	}
 }
 
@@ -258,7 +278,7 @@ impl<H, N> sp_externalities::ExtensionStore for TestExternalities<H, N> where
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use sp_core::{H256, traits::Externalities};
+	use sp_core::{H256, traits::Externalities, storage::ChildInfo};
 	use sp_runtime::traits::BlakeTwo256;
 	use hex_literal::hex;
 
@@ -289,4 +309,45 @@ mod tests {
 		fn assert_send<T: Send>() {}
 		assert_send::<TestExternalities::<BlakeTwo256, u64>>();
 	}
+
+	#[test]
+	fn commit_all_and_kill_child_storage() {
+		let mut ext = TestExternalities::<BlakeTwo256, u64>::default();
+		let child_info = ChildInfo::new_default(&b"test_child"[..]);
+
+		{
+			let mut ext = ext.ext();
+			ext.place_child_storage(&child_info, b"doe".to_vec(), Some(b"reindeer".to_vec()));
+			ext.place_child_storage(&child_info, b"dog".to_vec(), Some(b"puppy".to_vec()));
+			ext.place_child_storage(&child_info, b"dog2".to_vec(), Some(b"puppy2".to_vec()));
+		}
+
+		ext.commit_all().unwrap();
+
+		{
+			let mut ext = ext.ext();
+
+			assert!(!ext.kill_child_storage(&child_info, Some(2)), "Should not delete all keys");
+
+			assert!(ext.child_storage(&child_info, &b"doe"[..]).is_none());
+			assert!(ext.child_storage(&child_info, &b"dog"[..]).is_none());
+			assert!(ext.child_storage(&child_info, &b"dog2"[..]).is_some());
+		}
+	}
+
+	#[test]
+	fn as_backend_generates_same_backend_as_commit_all() {
+		let mut ext = TestExternalities::<BlakeTwo256, u64>::default();
+		{
+			let mut ext = ext.ext();
+			ext.set_storage(b"doe".to_vec(), b"reindeer".to_vec());
+			ext.set_storage(b"dog".to_vec(), b"puppy".to_vec());
+			ext.set_storage(b"dogglesworth".to_vec(), b"cat".to_vec());
+		}
+
+		let backend = ext.as_backend();
+
+		ext.commit_all().unwrap();
+		assert!(ext.backend.eq(&backend), "Both backend should be equal.");
+	}
 }
diff --git a/primitives/trie/src/lib.rs b/primitives/trie/src/lib.rs
@@ -184,7 +184,7 @@ pub fn delta_trie_root<L: TrieConfiguration, I, A, B, DB, V>(
 	DB: hash_db::HashDB<L::Hash, trie_db::DBValue>,
 {
 	{
-		let mut trie = TrieDBMut::<L>::from_existing(&mut *db, &mut root)?;
+		let mut trie = TrieDBMut::<L>::from_existing(db, &mut root)?;
 
 		let mut delta = delta.into_iter().collect::<Vec<_>>();
 		delta.sort_by(|l, r| l.0.borrow().cmp(r.0.borrow()));
@@ -223,9 +223,13 @@ pub fn read_trie_value_with<
 	Ok(TrieDB::<L>::new(&*db, root)?.get_with(key, query).map(|x| x.map(|val| val.to_vec()))?)
 }
 
+/// Determine the empty trie root.
+pub fn empty_trie_root<L: TrieConfiguration>() -> <L::Hash as Hasher>::Out {
+	L::trie_root::<_, Vec<u8>, Vec<u8>>(core::iter::empty())
+}
+
 /// Determine the empty child trie root.
-pub fn empty_child_trie_root<L: TrieConfiguration>(
-) -> <L::Hash as Hasher>::Out {
+pub fn empty_child_trie_root<L: TrieConfiguration>() -> <L::Hash as Hasher>::Out {
 	L::trie_root::<_, Vec<u8>, Vec<u8>>(core::iter::empty())
 }