Implement nested storage transactions (#6269)

* Add transactional storage functionality to OverlayChanges

A collection already has a natural None state. No need to
wrap it with an option.

* Add storage transactions runtime interface

* Add frame support for transactions

* Fix committed typo

* Rename 'changes' variable to 'overlay'

* Fix renaming change

* Fixed strange line break

* Rename clear to clear_where

* Add comment regarding delete value on mutation

* Add comment which changes are covered by a transaction

* Do force the arg to with_transaction return a Result

* Use rust doc comments on every documentable place

* Fix wording of insert_diry doc

* Improve doc on start_transaction

* Rename value to overlayed in close_transaction

* Inline negation

* Improve wording of close_transaction comments

* Get rid of an expect by using get_or_insert_with

* Remove trailing whitespace

* Rename should to expected in tests

* Rolling back a transaction must mark the overlay as dirty

* Protect client initiated storage tx from being droped by runtime

* Review nits

* Return Err when entering or exiting runtime fails

* Documentation fixup

* Remove close type

* Move enter/exit runtime to excute_aux in the state-machine

* Rename Discard -> Rollback

* Move child changeset creation to constructor

* Move child spawning into the closure

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Fixup for code suggestion

* Unify re-exports

* Rename overlay_changes to mod.rs and move into subdir

* Change proof wording

* Adapt a new test from master to storage-tx

* Suggestions from the latest round of review

* Fix warning message

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/primitives/state-machine/src/basic.rs

@@ -307,6 +307,18 @@ impl Externalities for BasicExternalities {
 		Ok(None)
 	}
 
+	fn storage_start_transaction(&mut self) {
+		unimplemented!("Transactions are not supported by BasicExternalities");
+	}
+
+	fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
+		unimplemented!("Transactions are not supported by BasicExternalities");
+	}
+
+	fn storage_commit_transaction(&mut self) -> Result<(), ()> {
+		unimplemented!("Transactions are not supported by BasicExternalities");
+	}
+
 	fn wipe(&mut self) {}
 
 	fn commit(&mut self) {}

substrate/primitives/state-machine/src/changes_trie/build.rs

@@ -25,7 +25,7 @@ use num_traits::One;
 use crate::{
 	StorageKey,
 	backend::Backend,
-	overlayed_changes::OverlayedChanges,
+	overlayed_changes::{OverlayedChanges, OverlayedValue},
 	trie_backend_essence::TrieBackendEssence,
 	changes_trie::{
 		AnchorBlockId, ConfigurationRange, Storage, BlockNumber,
@@ -43,7 +43,7 @@ pub(crate) fn prepare_input<'a, B, H, Number>(
 	backend: &'a B,
 	storage: &'a dyn Storage<H, Number>,
 	config: ConfigurationRange<'a, Number>,
-	changes: &'a OverlayedChanges,
+	overlay: &'a OverlayedChanges,
 	parent: &'a AnchorBlockId<H::Out, Number>,
 ) -> Result<(
 		impl Iterator<Item=InputPair<Number>> + 'a,
@@ -60,7 +60,7 @@ pub(crate) fn prepare_input<'a, B, H, Number>(
 	let (extrinsics_input, children_extrinsics_input) = prepare_extrinsics_input(
 		backend,
 		&number,
-		changes,
+		overlay,
 	)?;
 	let (digest_input, mut children_digest_input, digest_input_blocks) = prepare_digest_input::<H, Number>(
 		parent,
@@ -96,7 +96,7 @@ pub(crate) fn prepare_input<'a, B, H, Number>(
 fn prepare_extrinsics_input<'a, B, H, Number>(
 	backend: &'a B,
 	block: &Number,
-	changes: &'a OverlayedChanges,
+	overlay: &'a OverlayedChanges,
 ) -> Result<(
 		impl Iterator<Item=InputPair<Number>> + 'a,
 		BTreeMap<ChildIndex<Number>, impl Iterator<Item=InputPair<Number>> + 'a>,
@@ -108,20 +108,21 @@ fn prepare_extrinsics_input<'a, B, H, Number>(
 {
 	let mut children_result = BTreeMap::new();
 
-	for child_info in changes.child_infos() {
+	for (child_changes, child_info) in overlay.children() {
 		let child_index = ChildIndex::<Number> {
 			block: block.clone(),
 			storage_key: child_info.prefixed_storage_key(),
 		};
 
 		let iter = prepare_extrinsics_input_inner(
-			backend, block, changes,
-			Some(child_info.clone())
+			backend, block, overlay,
+			Some(child_info.clone()),
+			child_changes,
 		)?;
 		children_result.insert(child_index, iter);
 	}
 
-	let top = prepare_extrinsics_input_inner(backend, block, changes, None)?;
+	let top = prepare_extrinsics_input_inner(backend, block, overlay, None, overlay.changes())?;
 
 	Ok((top, children_result))
 }
@@ -129,40 +130,38 @@ fn prepare_extrinsics_input<'a, B, H, Number>(
 fn prepare_extrinsics_input_inner<'a, B, H, Number>(
 	backend: &'a B,
 	block: &Number,
-	changes: &'a OverlayedChanges,
+	overlay: &'a OverlayedChanges,
 	child_info: Option<ChildInfo>,
+	changes: impl Iterator<Item=(&'a StorageKey, &'a OverlayedValue)>
 ) -> Result<impl Iterator<Item=InputPair<Number>> + 'a, String>
 	where
 		B: Backend<H>,
 		H: Hasher,
 		Number: BlockNumber,
 {
-	changes.changes(child_info.as_ref())
-		.filter(|( _, v)| v.extrinsics().is_some())
+	changes
+		.filter(|( _, v)| v.extrinsics().next().is_some())
 		.try_fold(BTreeMap::new(), |mut map: BTreeMap<&[u8], (ExtrinsicIndex<Number>, Vec<u32>)>, (k, v)| {
 			match map.entry(k) {
 				Entry::Vacant(entry) => {
 					// ignore temporary values (values that have null value at the end of operation
 					// AND are not in storage at the beginning of operation
 					if let Some(child_info) = child_info.as_ref() {
-						if !changes.child_storage(child_info, k).map(|v| v.is_some()).unwrap_or_default() {
+						if !overlay.child_storage(child_info, k).map(|v| v.is_some()).unwrap_or_default() {
 							if !backend.exists_child_storage(&child_info, k)
 								.map_err(|e| format!("{}", e))? {
 								return Ok(map);
 							}
 						}
 					} else {
-						if !changes.storage(k).map(|v| v.is_some()).unwrap_or_default() {
+						if !overlay.storage(k).map(|v| v.is_some()).unwrap_or_default() {
 							if !backend.exists_storage(k).map_err(|e| format!("{}", e))? {
 								return Ok(map);
 							}
 						}
 					};
 
-					let extrinsics = v.extrinsics()
-						.expect("filtered by filter() call above; qed")
-						.cloned()
-						.collect();
+					let extrinsics = v.extrinsics().cloned().collect();
 					entry.insert((ExtrinsicIndex {
 						block: block.clone(),
 						key: k.to_vec(),
@@ -173,9 +172,7 @@ fn prepare_extrinsics_input_inner<'a, B, H, Number>(
 					// AND we are checking it before insertion
 					let extrinsics = &mut entry.get_mut().1;
 					extrinsics.extend(
-						v.extrinsics()
-							.expect("filtered by filter() call above; qed")
-							.cloned()
+						v.extrinsics().cloned()
 					);
 					extrinsics.sort_unstable();
 				},
@@ -404,6 +401,8 @@ mod test {
 		let mut changes = OverlayedChanges::default();
 		changes.set_collect_extrinsics(true);
 
+		changes.start_transaction();
+
 		changes.set_extrinsic_index(1);
 		changes.set_storage(vec![101], Some(vec![203]));
 
@@ -411,7 +410,7 @@ mod test {
 		changes.set_storage(vec![100], Some(vec![202]));
 		changes.set_child_storage(&child_info_1, vec![100], Some(vec![202]));
 
-		changes.commit_prospective();
+		changes.commit_transaction().unwrap();
 
 		changes.set_extrinsic_index(0);
 		changes.set_storage(vec![100], Some(vec![0]));

substrate/primitives/state-machine/src/ext.rs

@@ -37,6 +37,10 @@ use std::{error, fmt, any::{Any, TypeId}};
 use log::{warn, trace};
 
 const EXT_NOT_ALLOWED_TO_FAIL: &str = "Externalities not allowed to fail within runtime";
+const BENCHMARKING_FN: &str = "\
+	This is a special fn only for benchmarking where a database commit happens from the runtime.
+	For that reason client started transactions before calling into runtime are not allowed.
+	Without client transactions the loop condition garantuees the success of the tx close.";
 
 /// Errors that can occur when interacting with the externalities.
 #[derive(Debug, Copy, Clone)]
@@ -147,7 +151,7 @@ where
 
 		self.backend.pairs().iter()
 			.map(|&(ref k, ref v)| (k.to_vec(), Some(v.to_vec())))
-			.chain(self.overlay.changes(None).map(|(k, v)| (k.clone(), v.value().cloned())))
+			.chain(self.overlay.changes().map(|(k, v)| (k.clone(), v.value().cloned())))
 			.collect::<HashMap<_, _>>()
 			.into_iter()
 			.filter_map(|(k, maybe_val)| maybe_val.map(|val| (k, val)))
@@ -477,15 +481,14 @@ where
 			);
 			root.encode()
 		} else {
+			let root = if let Some((changes, info)) = self.overlay.child_changes(storage_key) {
+				let delta = changes.map(|(k, v)| (k.as_ref(), v.value().map(AsRef::as_ref)));
+				Some(self.backend.child_storage_root(info, delta))
+			} else {
+				None
+			};
 
-			if let Some(child_info) = self.overlay.default_child_info(storage_key) {
-				let (root, is_empty, _) = {
-					let delta = self.overlay.changes(Some(child_info))
-						.map(|(k, v)| (k.as_ref(), v.value().map(AsRef::as_ref)));
-
-					self.backend.child_storage_root(child_info, delta)
-				};
-
+			if let Some((root, is_empty, _)) = root {
 				let root = root.encode();
 				// We store update in the overlay in order to be able to use 'self.storage_transaction'
 				// cache. This is brittle as it rely on Ext only querying the trie backend for
@@ -547,20 +550,37 @@ where
 		root.map(|r| r.map(|o| o.encode()))
 	}
 
+	fn storage_start_transaction(&mut self) {
+		self.overlay.start_transaction()
+	}
+
+	fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
+		self.mark_dirty();
+		self.overlay.rollback_transaction().map_err(|_| ())
+	}
+
+	fn storage_commit_transaction(&mut self) -> Result<(), ()> {
+		self.overlay.commit_transaction().map_err(|_| ())
+	}
+
 	fn wipe(&mut self) {
-		self.overlay.discard_prospective();
+		for _ in 0..self.overlay.transaction_depth() {
+			self.overlay.rollback_transaction().expect(BENCHMARKING_FN);
+		}
 		self.overlay.drain_storage_changes(
 			&self.backend,
 			None,
 			Default::default(),
 			self.storage_transaction_cache,
 		).expect(EXT_NOT_ALLOWED_TO_FAIL);
-		self.storage_transaction_cache.reset();
-		self.backend.wipe().expect(EXT_NOT_ALLOWED_TO_FAIL)
+		self.backend.wipe().expect(EXT_NOT_ALLOWED_TO_FAIL);
+		self.mark_dirty();
 	}
 
 	fn commit(&mut self) {
-		self.overlay.commit_prospective();
+		for _ in 0..self.overlay.transaction_depth() {
+			self.overlay.commit_transaction().expect(BENCHMARKING_FN);
+		}
 		let changes = self.overlay.drain_storage_changes(
 			&self.backend,
 			None,
@@ -571,7 +591,7 @@ where
 			changes.transaction_storage_root,
 			changes.transaction,
 		).expect(EXT_NOT_ALLOWED_TO_FAIL);
-		self.storage_transaction_cache.reset();
+		self.mark_dirty();
 	}
 }
 

substrate/primitives/state-machine/src/lib.rs

@@ -79,6 +79,10 @@ pub use in_memory_backend::new_in_mem;
 pub use stats::{UsageInfo, UsageUnit, StateMachineStats};
 pub use sp_core::traits::CloneableSpawn;
 
+const PROOF_CLOSE_TRANSACTION: &str = "\
+	Closing a transaction that was started in this function. Client initiated transactions
+	are protected from being closed by the runtime. qed";
+
 type CallResult<R, E> = Result<NativeOrEncoded<R>, E>;
 
 /// Default handler of the execution manager.
@@ -297,6 +301,8 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 			None => &mut cache,
 		};
 
+		self.overlay.enter_runtime().expect("StateMachine is never called from the runtime; qed");
+
 		let mut ext = Ext::new(
 			self.overlay,
 			self.offchain_overlay,
@@ -324,6 +330,9 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 			native_call,
 		);
 
+		self.overlay.exit_runtime()
+			.expect("Runtime is not able to call this function in the overlay; qed");
+
 		trace!(
 			target: "state", "{:04x}: Return. Native={:?}, Result={:?}",
 			id,
@@ -347,11 +356,11 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 				CallResult<R, Exec::Error>,
 			) -> CallResult<R, Exec::Error>
 	{
-		let pending_changes = self.overlay.clone_pending();
+		self.overlay.start_transaction();
 		let (result, was_native) = self.execute_aux(true, native_call.take());
 
 		if was_native {
-			self.overlay.replace_pending(pending_changes);
+			self.overlay.rollback_transaction().expect(PROOF_CLOSE_TRANSACTION);
 			let (wasm_result, _) = self.execute_aux(
 				false,
 				native_call,
@@ -366,6 +375,7 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 				on_consensus_failure(wasm_result, result)
 			}
 		} else {
+			self.overlay.commit_transaction().expect(PROOF_CLOSE_TRANSACTION);
 			result
 		}
 	}
@@ -378,16 +388,17 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 			R: Decode + Encode + PartialEq,
 			NC: FnOnce() -> result::Result<R, String> + UnwindSafe,
 	{
-		let pending_changes = self.overlay.clone_pending();
+		self.overlay.start_transaction();
 		let (result, was_native) = self.execute_aux(
 			true,
 			native_call.take(),
 		);
 
 		if !was_native || result.is_ok() {
+			self.overlay.commit_transaction().expect(PROOF_CLOSE_TRANSACTION);
 			result
 		} else {
-			self.overlay.replace_pending(pending_changes);
+			self.overlay.rollback_transaction().expect(PROOF_CLOSE_TRANSACTION);
 			let (wasm_result, _) = self.execute_aux(
 				false,
 				native_call,
@@ -977,7 +988,7 @@ mod tests {
 		let mut overlay = OverlayedChanges::default();
 		overlay.set_storage(b"aba".to_vec(), Some(b"1312".to_vec()));
 		overlay.set_storage(b"bab".to_vec(), Some(b"228".to_vec()));
-		overlay.commit_prospective();
+		overlay.start_transaction();
 		overlay.set_storage(b"abd".to_vec(), Some(b"69".to_vec()));
 		overlay.set_storage(b"bbd".to_vec(), Some(b"42".to_vec()));
 
@@ -994,10 +1005,10 @@ mod tests {
 			);
 			ext.clear_prefix(b"ab");
 		}
-		overlay.commit_prospective();
+		overlay.commit_transaction().unwrap();
 
 		assert_eq!(
-			overlay.changes(None).map(|(k, v)| (k.clone(), v.value().cloned()))
+			overlay.changes().map(|(k, v)| (k.clone(), v.value().cloned()))
 				.collect::<HashMap<_, _>>(),
 			map![
 				b"abc".to_vec() => None.into(),
@@ -1083,7 +1094,7 @@ mod tests {
 				Some(vec![reference_data[0].clone()].encode()),
 			);
 		}
-		overlay.commit_prospective();
+		overlay.start_transaction();
 		{
 			let mut ext = Ext::new(
 				&mut overlay,
@@ -1102,7 +1113,7 @@ mod tests {
 				Some(reference_data.encode()),
 			);
 		}
-		overlay.discard_prospective();
+		overlay.rollback_transaction().unwrap();
 		{
 			let ext = Ext::new(
 				&mut overlay,
@@ -1145,7 +1156,7 @@ mod tests {
 			ext.clear_storage(key.as_slice());
 			ext.storage_append(key.clone(), Item::InitializationItem.encode());
 		}
-		overlay.commit_prospective();
+		overlay.start_transaction();
 
 		// For example, first transaction resulted in panic during block building
 		{
@@ -1170,7 +1181,7 @@ mod tests {
 				Some(vec![Item::InitializationItem, Item::DiscardedItem].encode()),
 			);
 		}
-		overlay.discard_prospective();
+		overlay.rollback_transaction().unwrap();
 
 		// Then we apply next transaction which is valid this time.
 		{
@@ -1196,7 +1207,7 @@ mod tests {
 			);
 
 		}
-		overlay.commit_prospective();
+		overlay.start_transaction();
 
 		// Then only initlaization item and second (commited) item should persist.
 		{
@@ -1317,9 +1328,11 @@ mod tests {
 		let backend = state.as_trie_backend().unwrap();
 
 		let mut overlay = OverlayedChanges::default();
+		overlay.start_transaction();
 		overlay.set_storage(b"ccc".to_vec(), Some(b"".to_vec()));
 		assert_eq!(overlay.storage(b"ccc"), Some(Some(&[][..])));
-		overlay.commit_prospective();
+		overlay.commit_transaction().unwrap();
+		overlay.start_transaction();
 		assert_eq!(overlay.storage(b"ccc"), Some(Some(&[][..])));
 		assert_eq!(overlay.storage(b"bbb"), None);
 
@@ -1339,7 +1352,7 @@ mod tests {
 			ext.clear_storage(b"ccc");
 			assert_eq!(ext.storage(b"ccc"), None);
 		}
-		overlay.commit_prospective();
+		overlay.commit_transaction().unwrap();
 		assert_eq!(overlay.storage(b"ccc"), Some(None));
 	}
 }

substrate/primitives/state-machine/src/overlayed_changes/changeset.rs

@@ -0,0 +1,752 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2020 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License
+
+//! Houses the code that implements the transactional overlay storage.
+
+use super::{StorageKey, StorageValue};
+
+use itertools::Itertools;
+use std::collections::{HashSet, BTreeMap, BTreeSet};
+use smallvec::SmallVec;
+use log::warn;
+
+const PROOF_OVERLAY_NON_EMPTY: &str = "\
+	An OverlayValue is always created with at least one transaction and dropped as soon
+	as the last transaction is removed; qed";
+
+type DirtyKeysSets = SmallVec<[HashSet<StorageKey>; 5]>;
+type Transactions = SmallVec<[InnerValue; 5]>;
+
+/// Error returned when trying to commit or rollback while no transaction is open or
+/// when the runtime is trying to close a transaction started by the client.
+#[derive(Debug)]
+#[cfg_attr(test, derive(PartialEq))]
+pub struct NoOpenTransaction;
+
+/// Error when calling `enter_runtime` when already being in runtime execution mode.
+#[derive(Debug)]
+#[cfg_attr(test, derive(PartialEq))]
+pub struct AlreadyInRuntime;
+
+/// Error when calling `exit_runtime` when not being in runtime exection mdde.
+#[derive(Debug)]
+#[cfg_attr(test, derive(PartialEq))]
+pub struct NotInRuntime;
+
+/// Describes in which mode the node is currently executing.
+#[derive(Debug, Clone, Copy)]
+pub enum ExecutionMode {
+	/// Exeuting in client mode: Removal of all transactions possible.
+	Client,
+	/// Executing in runtime mode: Transactions started by the client are protected.
+	Runtime,
+}
+
+#[derive(Debug, Default, Clone)]
+#[cfg_attr(test, derive(PartialEq))]
+struct InnerValue {
+	/// Current value. None if value has been deleted.
+	value: Option<StorageValue>,
+	/// The set of extrinsic indices where the values has been changed.
+	/// Is filled only if runtime has announced changes trie support.
+	extrinsics: BTreeSet<u32>,
+}
+
+/// An overlay that contains all versions of a value for a specific key.
+#[derive(Debug, Default, Clone)]
+#[cfg_attr(test, derive(PartialEq))]
+pub struct OverlayedValue {
+	/// The individual versions of that value.
+	/// One entry per transactions during that the value was actually written.
+	transactions: Transactions,
+}
+
+/// Holds a set of changes with the ability modify them using nested transactions.
+#[derive(Debug, Default, Clone)]
+pub struct OverlayedChangeSet {
+	/// Stores the changes that this overlay constitutes.
+	changes: BTreeMap<StorageKey, OverlayedValue>,
+	/// Stores which keys are dirty per transaction. Needed in order to determine which
+	/// values to merge into the parent transaction on commit. The length of this vector
+	/// therefore determines how many nested transactions are currently open (depth).
+	dirty_keys: DirtyKeysSets,
+	/// The number of how many transactions beginning from the first transactions are started
+	/// by the client. Those transactions are protected against close (commit, rollback)
+	/// when in runtime mode.
+	num_client_transactions: usize,
+	/// Determines whether the node is using the overlay from the client or the runtime.
+	execution_mode: ExecutionMode,
+}
+
+impl Default for ExecutionMode {
+	fn default() -> Self {
+		Self::Client
+	}
+}
+
+impl OverlayedValue {
+	/// The value as seen by the current transaction.
+	pub fn value(&self) -> Option<&StorageValue> {
+		self.transactions.last().expect(PROOF_OVERLAY_NON_EMPTY).value.as_ref()
+	}
+
+	/// Unique list of extrinsic indices which modified the value.
+	pub fn extrinsics(&self) -> impl Iterator<Item=&u32> {
+		self.transactions.iter().flat_map(|t| t.extrinsics.iter()).unique()
+	}
+
+	/// Mutable reference to the most recent version.
+	fn value_mut(&mut self) -> &mut Option<StorageValue> {
+		&mut self.transactions.last_mut().expect(PROOF_OVERLAY_NON_EMPTY).value
+	}
+
+	/// Remove the last version and return it.
+	fn pop_transaction(&mut self) -> InnerValue {
+		self.transactions.pop().expect(PROOF_OVERLAY_NON_EMPTY)
+	}
+
+	/// Mutable reference to the set which holds the indices for the **current transaction only**.
+	fn transaction_extrinsics_mut(&mut self) -> &mut BTreeSet<u32> {
+		&mut self.transactions.last_mut().expect(PROOF_OVERLAY_NON_EMPTY).extrinsics
+	}
+
+	/// Writes a new version of a value.
+	///
+	/// This makes sure that the old version is not overwritten and can be properly
+	/// rolled back when required.
+	fn set(
+		&mut self,
+		value: Option<StorageValue>,
+		first_write_in_tx: bool,
+		at_extrinsic: Option<u32>,
+	) {
+		if first_write_in_tx || self.transactions.is_empty() {
+			self.transactions.push(InnerValue {
+				value,
+				.. Default::default()
+			});
+		} else {
+			*self.value_mut() = value;
+		}
+
+		if let Some(extrinsic) = at_extrinsic {
+			self.transaction_extrinsics_mut().insert(extrinsic);
+		}
+	}
+}
+
+/// Inserts a key into the dirty set.
+///
+/// Returns true iff we are currently have at least one open transaction and if this
+/// is the first write to the given key that transaction.
+fn insert_dirty(set: &mut DirtyKeysSets, key: StorageKey) -> bool {
+	set.last_mut().map(|dk| dk.insert(key)).unwrap_or_default()
+}
+
+impl OverlayedChangeSet {
+	/// Create a new changeset at the same transaction state but without any contents.
+	///
+	/// This changeset might be created when there are already open transactions.
+	/// We need to catch up here so that the child is at the same transaction depth.
+	pub fn spawn_child(&self) -> Self {
+		use std::iter::repeat;
+		Self {
+			dirty_keys: repeat(HashSet::new()).take(self.transaction_depth()).collect(),
+			num_client_transactions: self.num_client_transactions,
+			execution_mode: self.execution_mode,
+			.. Default::default()
+		}
+	}
+
+	/// True if no changes at all are contained in the change set.
+	pub fn is_empty(&self) -> bool {
+		self.changes.is_empty()
+	}
+
+	/// Get an optional reference to the value stored for the specified key.
+	pub fn get(&self, key: &[u8]) -> Option<&OverlayedValue> {
+		self.changes.get(key)
+	}
+
+	/// Set a new value for the specified key.
+	///
+	/// Can be rolled back or committed when called inside a transaction.
+	pub fn set(
+		&mut self,
+		key: StorageKey,
+		value: Option<StorageValue>,
+		at_extrinsic: Option<u32>,
+	) {
+		let overlayed = self.changes.entry(key.clone()).or_default();
+		overlayed.set(value, insert_dirty(&mut self.dirty_keys, key), at_extrinsic);
+	}
+
+	/// Get a mutable reference for a value.
+	///
+	/// Can be rolled back or committed when called inside a transaction.
+	#[must_use = "A change was registered, so this value MUST be modified."]
+	pub fn modify(
+		&mut self,
+		key: StorageKey,
+		init: impl Fn() -> StorageValue,
+		at_extrinsic: Option<u32>,
+	) -> &mut Option<StorageValue> {
+		let overlayed = self.changes.entry(key.clone()).or_default();
+		let first_write_in_tx = insert_dirty(&mut self.dirty_keys, key);
+		let clone_into_new_tx = if let Some(tx) = overlayed.transactions.last() {
+			if first_write_in_tx {
+				Some(tx.value.clone())
+			} else {
+				None
+			}
+		} else {
+			Some(Some(init()))
+		};
+
+		if let Some(cloned) = clone_into_new_tx {
+			overlayed.set(cloned, first_write_in_tx, at_extrinsic);
+		}
+		overlayed.value_mut()
+	}
+
+	/// Set all values to deleted which are matched by the predicate.
+	///
+	/// Can be rolled back or committed when called inside a transaction.
+	pub fn clear_where(
+		&mut self,
+		predicate: impl Fn(&[u8], &OverlayedValue) -> bool,
+		at_extrinsic: Option<u32>,
+	) {
+		for (key, val) in self.changes.iter_mut().filter(|(k, v)| predicate(k, v)) {
+			val.set(None, insert_dirty(&mut self.dirty_keys, key.to_owned()), at_extrinsic);
+		}
+	}
+
+	/// Get a list of all changes as seen by current transaction.
+	pub fn changes(&self) -> impl Iterator<Item=(&StorageKey, &OverlayedValue)> {
+		self.changes.iter()
+	}
+
+	/// Get the change that is next to the supplied key.
+	pub fn next_change(&self, key: &[u8]) -> Option<(&[u8], &OverlayedValue)> {
+		use std::ops::Bound;
+		let range = (Bound::Excluded(key), Bound::Unbounded);
+		self.changes.range::<[u8], _>(range).next().map(|(k, v)| (&k[..], v))
+	}
+
+	/// Consume this changeset and return all committed changes.
+	///
+	/// Panics:
+	/// Panics if there are open transactions: `transaction_depth() > 0`
+	pub fn drain_commited(self) -> impl Iterator<Item=(StorageKey, Option<StorageValue>)> {
+		assert!(self.transaction_depth() == 0, "Drain is not allowed with open transactions.");
+		self.changes.into_iter().map(|(k, mut v)| (k, v.pop_transaction().value))
+	}
+
+	/// Returns the current nesting depth of the transaction stack.
+	///
+	/// A value of zero means that no transaction is open and changes are committed on write.
+	pub fn transaction_depth(&self) -> usize {
+		self.dirty_keys.len()
+	}
+
+	/// Call this before transfering control to the runtime.
+	///
+	/// This protects all existing transactions from being removed by the runtime.
+	/// Calling this while already inside the runtime will return an error.
+	pub fn enter_runtime(&mut self) -> Result<(), AlreadyInRuntime> {
+		if let ExecutionMode::Runtime = self.execution_mode {
+			return Err(AlreadyInRuntime);
+		}
+		self.execution_mode = ExecutionMode::Runtime;
+		self.num_client_transactions = self.transaction_depth();
+		Ok(())
+	}
+
+	/// Call this when control returns from the runtime.
+	///
+	/// This commits all dangling transaction left open by the runtime.
+	/// Calling this while already outside the runtime will return an error.
+	pub fn exit_runtime(&mut self) -> Result<(), NotInRuntime> {
+		if let ExecutionMode::Client = self.execution_mode {
+			return Err(NotInRuntime);
+		}
+		self.execution_mode = ExecutionMode::Client;
+		if self.has_open_runtime_transactions() {
+			warn!(
+				"{} storage transactions are left open by the runtime. Those will be rolled back.",
+				self.transaction_depth() - self.num_client_transactions,
+			);
+		}
+		while self.has_open_runtime_transactions() {
+			self.rollback_transaction()
+				.expect("The loop condition checks that the transaction depth is > 0; qed");
+		}
+		Ok(())
+	}
+
+	/// Start a new nested transaction.
+	///
+	/// This allows to either commit or roll back all changes that were made while this
+	/// transaction was open. Any transaction must be closed by either `commit_transaction`
+	/// or `rollback_transaction` before this overlay can be converted into storage changes.
+	///
+	/// Changes made without any open transaction are committed immediately.
+	pub fn start_transaction(&mut self) {
+		self.dirty_keys.push(Default::default());
+	}
+
+	/// Rollback the last transaction started by `start_transaction`.
+	///
+	/// Any changes made during that transaction are discarded. Returns an error if
+	/// there is no open transaction that can be rolled back.
+	pub fn rollback_transaction(&mut self) -> Result<(), NoOpenTransaction> {
+		self.close_transaction(true)
+	}
+
+	/// Commit the last transaction started by `start_transaction`.
+	///
+	/// Any changes made during that transaction are committed. Returns an error if
+	/// there is no open transaction that can be committed.
+	pub fn commit_transaction(&mut self) -> Result<(), NoOpenTransaction> {
+		self.close_transaction(false)
+	}
+
+	fn close_transaction(&mut self, rollback: bool) -> Result<(), NoOpenTransaction> {
+		// runtime is not allowed to close transactions started by the client
+		if let ExecutionMode::Runtime = self.execution_mode {
+			if !self.has_open_runtime_transactions() {
+				return Err(NoOpenTransaction)
+			}
+		}
+
+		for key in self.dirty_keys.pop().ok_or(NoOpenTransaction)? {
+			let overlayed = self.changes.get_mut(&key).expect("\
+				A write to an OverlayedValue is recorded in the dirty key set. Before an
+				OverlayedValue is removed, its containing dirty set is removed. This
+				function is only called for keys that are in the dirty set. qed\
+			");
+
+			if rollback {
+				overlayed.pop_transaction();
+
+				// We need to remove the key as an `OverlayValue` with no transactions
+				// violates its invariant of always having at least one transaction.
+				if overlayed.transactions.is_empty() {
+					self.changes.remove(&key);
+				}
+			} else {
+				let has_predecessor = if let Some(dirty_keys) = self.dirty_keys.last_mut() {
+					// Not the last tx: Did the previous tx write to this key?
+					!dirty_keys.insert(key)
+				} else {
+					// Last tx: Is there already a value in the committed set?
+					// Check against one rather than empty because the current tx is still
+					// in the list as it is popped later in this function.
+					overlayed.transactions.len() > 1
+				};
+
+				// We only need to merge if there is an pre-existing value. It may be a value from
+				// the previous transaction or a value committed without any open transaction.
+				if has_predecessor {
+					let dropped_tx = overlayed.pop_transaction();
+					*overlayed.value_mut() = dropped_tx.value;
+					overlayed.transaction_extrinsics_mut().extend(dropped_tx.extrinsics);
+				}
+			}
+		}
+
+		Ok(())
+	}
+
+	fn has_open_runtime_transactions(&self) -> bool {
+		self.transaction_depth() > self.num_client_transactions
+	}
+}
+
+#[cfg(test)]
+mod test {
+	use super::*;
+	use pretty_assertions::assert_eq;
+
+	type Changes<'a> = Vec<(&'a [u8], (Option<&'a [u8]>, Vec<u32>))>;
+	type Drained<'a> = Vec<(&'a [u8], Option<&'a [u8]>)>;
+
+	fn assert_changes(is: &OverlayedChangeSet, expected: &Changes) {
+		let is: Changes = is.changes().map(|(k, v)| {
+			(k.as_ref(), (v.value().map(AsRef::as_ref), v.extrinsics().cloned().collect()))
+		}).collect();
+		assert_eq!(&is, expected);
+	}
+
+	fn assert_drained_changes(is: OverlayedChangeSet, expected: Changes) {
+		let is = is.drain_commited().collect::<Vec<_>>();
+		let expected = expected
+			.iter()
+			.map(|(k, v)| (k.to_vec(), v.0.map(From::from))).collect::<Vec<_>>();
+		assert_eq!(is, expected);
+	}
+
+	fn assert_drained(is: OverlayedChangeSet, expected: Drained) {
+		let is = is.drain_commited().collect::<Vec<_>>();
+		let expected = expected
+			.iter()
+			.map(|(k, v)| (k.to_vec(), v.map(From::from))).collect::<Vec<_>>();
+		assert_eq!(is, expected);
+	}
+
+	#[test]
+	fn no_transaction_works() {
+		let mut changeset = OverlayedChangeSet::default();
+		assert_eq!(changeset.transaction_depth(), 0);
+
+		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
+		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(2));
+		changeset.set(b"key0".to_vec(), Some(b"val0-1".to_vec()), Some(9));
+
+		assert_drained(changeset, vec![
+			(b"key0", Some(b"val0-1")),
+			(b"key1", Some(b"val1")),
+		]);
+	}
+
+	#[test]
+	fn transaction_works() {
+		let mut changeset = OverlayedChangeSet::default();
+		assert_eq!(changeset.transaction_depth(), 0);
+
+		// no transaction: committed on set
+		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
+		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(1));
+		changeset.set(b"key0".to_vec(), Some(b"val0-1".to_vec()), Some(10));
+
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 1);
+
+		// we will commit that later
+		changeset.set(b"key42".to_vec(), Some(b"val42".to_vec()), Some(42));
+		changeset.set(b"key99".to_vec(), Some(b"val99".to_vec()), Some(99));
+
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 2);
+
+		// we will roll that back
+		changeset.set(b"key42".to_vec(), Some(b"val42-rolled".to_vec()), Some(421));
+		changeset.set(b"key7".to_vec(), Some(b"val7-rolled".to_vec()), Some(77));
+		changeset.set(b"key0".to_vec(), Some(b"val0-rolled".to_vec()), Some(1000));
+		changeset.set(b"key5".to_vec(), Some(b"val5-rolled".to_vec()), None);
+
+		// changes contain all changes not only the commmited ones.
+		let all_changes: Changes = vec![
+			(b"key0", (Some(b"val0-rolled"), vec![1, 10, 1000])),
+			(b"key1", (Some(b"val1"), vec![1])),
+			(b"key42", (Some(b"val42-rolled"), vec![42, 421])),
+			(b"key5", (Some(b"val5-rolled"), vec![])),
+			(b"key7", (Some(b"val7-rolled"), vec![77])),
+			(b"key99", (Some(b"val99"), vec![99])),
+		];
+		assert_changes(&changeset, &all_changes);
+
+		// this should be no-op
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 3);
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 4);
+		changeset.rollback_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 3);
+		changeset.commit_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 2);
+		assert_changes(&changeset, &all_changes);
+
+		// roll back our first transactions that actually contains something
+		changeset.rollback_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 1);
+
+		let rolled_back: Changes = vec![
+			(b"key0", (Some(b"val0-1"), vec![1, 10])),
+			(b"key1", (Some(b"val1"), vec![1])),
+			(b"key42", (Some(b"val42"), vec![42])),
+			(b"key99", (Some(b"val99"), vec![99])),
+		];
+		assert_changes(&changeset, &rolled_back);
+
+		changeset.commit_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 0);
+		assert_changes(&changeset, &rolled_back);
+
+		assert_drained_changes(changeset, rolled_back);
+	}
+
+	#[test]
+	fn transaction_commit_then_rollback_works() {
+		let mut changeset = OverlayedChangeSet::default();
+		assert_eq!(changeset.transaction_depth(), 0);
+
+		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
+		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(1));
+		changeset.set(b"key0".to_vec(), Some(b"val0-1".to_vec()), Some(10));
+
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 1);
+
+		changeset.set(b"key42".to_vec(), Some(b"val42".to_vec()), Some(42));
+		changeset.set(b"key99".to_vec(), Some(b"val99".to_vec()), Some(99));
+
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 2);
+
+		changeset.set(b"key42".to_vec(), Some(b"val42-rolled".to_vec()), Some(421));
+		changeset.set(b"key7".to_vec(), Some(b"val7-rolled".to_vec()), Some(77));
+		changeset.set(b"key0".to_vec(), Some(b"val0-rolled".to_vec()), Some(1000));
+		changeset.set(b"key5".to_vec(), Some(b"val5-rolled".to_vec()), None);
+
+		let all_changes: Changes = vec![
+			(b"key0", (Some(b"val0-rolled"), vec![1, 10, 1000])),
+			(b"key1", (Some(b"val1"), vec![1])),
+			(b"key42", (Some(b"val42-rolled"), vec![42, 421])),
+			(b"key5", (Some(b"val5-rolled"), vec![])),
+			(b"key7", (Some(b"val7-rolled"), vec![77])),
+			(b"key99", (Some(b"val99"), vec![99])),
+		];
+		assert_changes(&changeset, &all_changes);
+
+		// this should be no-op
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 3);
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 4);
+		changeset.rollback_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 3);
+		changeset.commit_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 2);
+		assert_changes(&changeset, &all_changes);
+
+		changeset.commit_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 1);
+
+		assert_changes(&changeset, &all_changes);
+
+		changeset.rollback_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 0);
+
+		let rolled_back: Changes = vec![
+			(b"key0", (Some(b"val0-1"), vec![1, 10])),
+			(b"key1", (Some(b"val1"), vec![1])),
+		];
+		assert_changes(&changeset, &rolled_back);
+
+		assert_drained_changes(changeset, rolled_back);
+	}
+
+	#[test]
+	fn modify_works() {
+		let mut changeset = OverlayedChangeSet::default();
+		assert_eq!(changeset.transaction_depth(), 0);
+		let init = || b"valinit".to_vec();
+
+		// committed set
+		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(0));
+		changeset.set(b"key1".to_vec(), None, Some(1));
+		let val = changeset.modify(b"key3".to_vec(), init, Some(3));
+		assert_eq!(val, &Some(b"valinit".to_vec()));
+		val.as_mut().unwrap().extend_from_slice(b"-modified");
+
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 1);
+		changeset.start_transaction();
+		assert_eq!(changeset.transaction_depth(), 2);
+
+		// non existing value -> init value should be returned
+		let val = changeset.modify(b"key2".to_vec(), init, Some(2));
+		assert_eq!(val, &Some(b"valinit".to_vec()));
+		val.as_mut().unwrap().extend_from_slice(b"-modified");
+
+		// existing value should be returned by modify
+		let val = changeset.modify(b"key0".to_vec(), init, Some(10));
+		assert_eq!(val, &Some(b"val0".to_vec()));
+		val.as_mut().unwrap().extend_from_slice(b"-modified");
+
+		// should work for deleted keys
+		let val = changeset.modify(b"key1".to_vec(), init, Some(20));
+		assert_eq!(val, &None);
+		*val = Some(b"deleted-modified".to_vec());
+
+		let all_changes: Changes = vec![
+			(b"key0", (Some(b"val0-modified"), vec![0, 10])),
+			(b"key1", (Some(b"deleted-modified"), vec![1, 20])),
+			(b"key2", (Some(b"valinit-modified"), vec![2])),
+			(b"key3", (Some(b"valinit-modified"), vec![3])),
+		];
+		assert_changes(&changeset, &all_changes);
+		changeset.commit_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 1);
+		assert_changes(&changeset, &all_changes);
+
+		changeset.rollback_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 0);
+		let rolled_back: Changes = vec![
+			(b"key0", (Some(b"val0"), vec![0])),
+			(b"key1", (None, vec![1])),
+			(b"key3", (Some(b"valinit-modified"), vec![3])),
+		];
+		assert_changes(&changeset, &rolled_back);
+		assert_drained_changes(changeset, rolled_back);
+	}
+
+	#[test]
+	fn clear_works() {
+		let mut changeset = OverlayedChangeSet::default();
+
+		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
+		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(2));
+		changeset.set(b"del1".to_vec(), Some(b"delval1".to_vec()), Some(3));
+		changeset.set(b"del2".to_vec(), Some(b"delval2".to_vec()), Some(4));
+
+		changeset.start_transaction();
+
+		changeset.clear_where(|k, _| k.starts_with(b"del"), Some(5));
+
+		assert_changes(&changeset, &vec![
+			(b"del1", (None, vec![3, 5])),
+			(b"del2", (None, vec![4, 5])),
+			(b"key0", (Some(b"val0"), vec![1])),
+			(b"key1", (Some(b"val1"), vec![2])),
+		]);
+
+		changeset.rollback_transaction().unwrap();
+
+		assert_changes(&changeset, &vec![
+			(b"del1", (Some(b"delval1"), vec![3])),
+			(b"del2", (Some(b"delval2"), vec![4])),
+			(b"key0", (Some(b"val0"), vec![1])),
+			(b"key1", (Some(b"val1"), vec![2])),
+		]);
+	}
+
+	#[test]
+	fn next_change_works() {
+		let mut changeset = OverlayedChangeSet::default();
+
+		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(0));
+		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(1));
+		changeset.set(b"key2".to_vec(), Some(b"val2".to_vec()), Some(2));
+
+		changeset.start_transaction();
+
+		changeset.set(b"key3".to_vec(), Some(b"val3".to_vec()), Some(3));
+		changeset.set(b"key4".to_vec(), Some(b"val4".to_vec()), Some(4));
+		changeset.set(b"key11".to_vec(), Some(b"val11".to_vec()), Some(11));
+
+		assert_eq!(changeset.next_change(b"key0").unwrap().0, b"key1");
+		assert_eq!(changeset.next_change(b"key0").unwrap().1.value(), Some(&b"val1".to_vec()));
+		assert_eq!(changeset.next_change(b"key1").unwrap().0, b"key11");
+		assert_eq!(changeset.next_change(b"key1").unwrap().1.value(), Some(&b"val11".to_vec()));
+		assert_eq!(changeset.next_change(b"key11").unwrap().0, b"key2");
+		assert_eq!(changeset.next_change(b"key11").unwrap().1.value(), Some(&b"val2".to_vec()));
+		assert_eq!(changeset.next_change(b"key2").unwrap().0, b"key3");
+		assert_eq!(changeset.next_change(b"key2").unwrap().1.value(), Some(&b"val3".to_vec()));
+		assert_eq!(changeset.next_change(b"key3").unwrap().0, b"key4");
+		assert_eq!(changeset.next_change(b"key3").unwrap().1.value(), Some(&b"val4".to_vec()));
+		assert_eq!(changeset.next_change(b"key4"), None);
+
+		changeset.rollback_transaction().unwrap();
+
+		assert_eq!(changeset.next_change(b"key0").unwrap().0, b"key1");
+		assert_eq!(changeset.next_change(b"key0").unwrap().1.value(), Some(&b"val1".to_vec()));
+		assert_eq!(changeset.next_change(b"key1").unwrap().0, b"key2");
+		assert_eq!(changeset.next_change(b"key1").unwrap().1.value(), Some(&b"val2".to_vec()));
+		assert_eq!(changeset.next_change(b"key11").unwrap().0, b"key2");
+		assert_eq!(changeset.next_change(b"key11").unwrap().1.value(), Some(&b"val2".to_vec()));
+		assert_eq!(changeset.next_change(b"key2"), None);
+		assert_eq!(changeset.next_change(b"key3"), None);
+		assert_eq!(changeset.next_change(b"key4"), None);
+
+	}
+
+	#[test]
+	fn no_open_tx_commit_errors() {
+		let mut changeset = OverlayedChangeSet::default();
+		assert_eq!(changeset.transaction_depth(), 0);
+		assert_eq!(changeset.commit_transaction(), Err(NoOpenTransaction));
+	}
+
+	#[test]
+	fn no_open_tx_rollback_errors() {
+		let mut changeset = OverlayedChangeSet::default();
+		assert_eq!(changeset.transaction_depth(), 0);
+		assert_eq!(changeset.rollback_transaction(), Err(NoOpenTransaction));
+	}
+
+	#[test]
+	fn unbalanced_transactions_errors() {
+		let mut changeset = OverlayedChangeSet::default();
+		changeset.start_transaction();
+		changeset.commit_transaction().unwrap();
+		assert_eq!(changeset.commit_transaction(), Err(NoOpenTransaction));
+	}
+
+	#[test]
+	#[should_panic]
+	fn drain_with_open_transaction_panics() {
+		let mut changeset = OverlayedChangeSet::default();
+		changeset.start_transaction();
+		let _ = changeset.drain_commited();
+	}
+
+	#[test]
+	fn runtime_cannot_close_client_tx() {
+		let mut changeset = OverlayedChangeSet::default();
+		changeset.start_transaction();
+		changeset.enter_runtime().unwrap();
+		changeset.start_transaction();
+		changeset.commit_transaction().unwrap();
+		assert_eq!(changeset.commit_transaction(), Err(NoOpenTransaction));
+		assert_eq!(changeset.rollback_transaction(), Err(NoOpenTransaction));
+	}
+
+	#[test]
+	fn exit_runtime_closes_runtime_tx() {
+		let mut changeset = OverlayedChangeSet::default();
+
+		changeset.start_transaction();
+
+		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
+
+		changeset.enter_runtime().unwrap();
+		changeset.start_transaction();
+		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(2));
+		changeset.exit_runtime().unwrap();
+
+		changeset.commit_transaction().unwrap();
+		assert_eq!(changeset.transaction_depth(), 0);
+
+		assert_drained(changeset, vec![
+			(b"key0", Some(b"val0")),
+		]);
+	}
+
+	#[test]
+	fn enter_exit_runtime_fails_when_already_in_requested_mode() {
+		let mut changeset = OverlayedChangeSet::default();
+
+		assert_eq!(changeset.exit_runtime(), Err(NotInRuntime));
+		assert_eq!(changeset.enter_runtime(), Ok(()));
+		assert_eq!(changeset.enter_runtime(), Err(AlreadyInRuntime));
+		assert_eq!(changeset.exit_runtime(), Ok(()));
+		assert_eq!(changeset.exit_runtime(), Err(NotInRuntime));
+	}
+}

substrate/primitives/state-machine/src/overlayed_changes/mod.rs

@@ -17,6 +17,8 @@
 
 //! The overlayed changes to state.
 
+mod changeset;
+
 use crate::{
 	backend::Backend, ChangesTrieTransaction,
 	changes_trie::{
@@ -25,14 +27,16 @@ use crate::{
 	},
 	stats::StateMachineStats,
 };
+use self::changeset::OverlayedChangeSet;
 
-use std::{mem, ops, collections::{HashMap, BTreeMap, BTreeSet}};
+use std::collections::HashMap;
 use codec::{Decode, Encode};
-use sp_core::storage::{well_known_keys::EXTRINSIC_INDEX, ChildInfo, ChildType};
+use sp_core::storage::{well_known_keys::EXTRINSIC_INDEX, ChildInfo};
 use sp_core::offchain::storage::OffchainOverlayedChanges;
-
 use hash_db::Hasher;
 
+pub use self::changeset::{OverlayedValue, NoOpenTransaction, AlreadyInRuntime, NotInRuntime};
+
 /// Storage key.
 pub type StorageKey = Vec<u8>;
 
@@ -45,43 +49,21 @@ pub type StorageCollection = Vec<(StorageKey, Option<StorageValue>)>;
 /// In memory arrays of storage values for multiple child tries.
 pub type ChildStorageCollection = Vec<(StorageKey, StorageCollection)>;
 
-/// The overlayed changes to state to be queried on top of the backend.
+/// The set of changes that are overlaid onto the backend.
 ///
-/// A transaction shares all prospective changes within an inner overlay
-/// that can be cleared.
+/// It allows changes to be modified using nestable transactions.
 #[derive(Debug, Default, Clone)]
 pub struct OverlayedChanges {
-	/// Changes that are not yet committed.
-	prospective: OverlayedChangeSet,
-	/// Committed changes.
-	committed: OverlayedChangeSet,
+	/// Top level storage changes.
+	top: OverlayedChangeSet,
+	/// Child storage changes. The map key is the child storage key without the common prefix.
+	children: HashMap<StorageKey, (OverlayedChangeSet, ChildInfo)>,
 	/// True if extrinsics stats must be collected.
 	collect_extrinsics: bool,
 	/// Collect statistic on this execution.
 	stats: StateMachineStats,
 }
 
-/// The storage value, used inside OverlayedChanges.
-#[derive(Debug, Default, Clone)]
-#[cfg_attr(test, derive(PartialEq))]
-pub struct OverlayedValue {
-	/// Current value. None if value has been deleted.
-	value: Option<StorageValue>,
-	/// The set of extrinsic indices where the values has been changed.
-	/// Is filled only if runtime has announced changes trie support.
-	extrinsics: Option<BTreeSet<u32>>,
-}
-
-/// Prospective or committed overlayed change set.
-#[derive(Debug, Default, Clone)]
-#[cfg_attr(test, derive(PartialEq))]
-pub struct OverlayedChangeSet {
-	/// Top level storage changes.
-	top: BTreeMap<StorageKey, OverlayedValue>,
-	/// Child storage changes. The map key is the child storage key without the common prefix.
-	children_default: HashMap<StorageKey, (BTreeMap<StorageKey, OverlayedValue>, ChildInfo)>,
-}
-
 /// A storage changes structure that can be generated by the data collected in [`OverlayedChanges`].
 ///
 /// This contains all the changes to the storage and transactions to apply theses changes to the
@@ -174,45 +156,10 @@ impl<Transaction: Default, H: Hasher, N: BlockNumber> Default for StorageChanges
 	}
 }
 
-#[cfg(test)]
-impl std::iter::FromIterator<(StorageKey, OverlayedValue)> for OverlayedChangeSet {
-	fn from_iter<T: IntoIterator<Item = (StorageKey, OverlayedValue)>>(iter: T) -> Self {
-		Self {
-			top: iter.into_iter().collect(),
-			children_default: Default::default(),
-		}
-	}
-}
-
-impl OverlayedValue {
-	/// The most recent value contained in this overlay.
-	pub fn value(&self) -> Option<&StorageValue> {
-		self.value.as_ref()
-	}
-
-	/// List of indices of extrinsics which modified the value using this overlay.
-	pub fn extrinsics(&self) -> Option<impl Iterator<Item=&u32>> {
-		self.extrinsics.as_ref().map(|v| v.iter())
-	}
-}
-
-impl OverlayedChangeSet {
-	/// Whether the change set is empty.
-	pub fn is_empty(&self) -> bool {
-		self.top.is_empty() && self.children_default.is_empty()
-	}
-
-	/// Clear the change set.
-	pub fn clear(&mut self) {
-		self.top.clear();
-		self.children_default.clear();
-	}
-}
-
 impl OverlayedChanges {
-	/// Whether the overlayed changes are empty.
+	/// Whether no changes are contained in the top nor in any of the child changes.
 	pub fn is_empty(&self) -> bool {
-		self.prospective.is_empty() && self.committed.is_empty()
+		self.top.is_empty() && self.children.is_empty()
 	}
 
 	/// Ask to collect/not to collect extrinsics indices where key(s) has been changed.
@@ -224,326 +171,241 @@ impl OverlayedChanges {
 	/// to the backend); Some(None) if the key has been deleted. Some(Some(...)) for a key whose
 	/// value has been set.
 	pub fn storage(&self, key: &[u8]) -> Option<Option<&[u8]>> {
-		self.prospective.top.get(key)
-			.or_else(|| self.committed.top.get(key))
-			.map(|x| {
-				let size_read = x.value.as_ref().map(|x| x.len() as u64).unwrap_or(0);
-				self.stats.tally_read_modified(size_read);
-				x.value.as_ref().map(AsRef::as_ref)
-			})
+		self.top.get(key).map(|x| {
+			let value = x.value();
+			let size_read = value.map(|x| x.len() as u64).unwrap_or(0);
+			self.stats.tally_read_modified(size_read);
+			value.map(AsRef::as_ref)
+		})
 	}
 
-	/// Returns mutable reference to current changed value (prospective).
-	/// If there is no value in the overlay, the default callback is used to initiate
-	/// the value.
-	/// Warning this function register a change, so the mutable reference MUST be modified.
+	/// Returns mutable reference to current value.
+	/// If there is no value in the overlay, the given callback is used to initiate the value.
+	/// Warning this function registers a change, so the mutable reference MUST be modified.
+	///
+	/// Can be rolled back or committed when called inside a transaction.
 	#[must_use = "A change was registered, so this value MUST be modified."]
 	pub fn value_mut_or_insert_with(
 		&mut self,
 		key: &[u8],
 		init: impl Fn() -> StorageValue,
 	) -> &mut StorageValue {
-		let extrinsic_index = self.extrinsic_index();
-		let committed = &self.committed.top;
+		let value = self.top.modify(key.to_owned(), init, self.extrinsic_index());
 
-		let mut entry = self.prospective.top.entry(key.to_vec())
-			.or_insert_with(|| {
-				if let Some(overlay_state) = committed.get(key).cloned() {
-					overlay_state
-				} else {
-					OverlayedValue { value: Some(init()), ..Default::default() }
-				}
-			});
-
-		//if was deleted initialise back with empty vec
-		if entry.value.is_none() {
-			entry.value = Some(Default::default());
-		}
-		if let Some(extrinsic) = extrinsic_index {
-			entry.extrinsics.get_or_insert_with(Default::default)
-				.insert(extrinsic);
-		}
-		entry.value.as_mut().expect("Initialized above; qed")
+		// if the value was deleted initialise it back with an empty vec
+		value.get_or_insert_with(StorageValue::default)
 	}
 
 	/// Returns a double-Option: None if the key is unknown (i.e. and the query should be referred
 	/// to the backend); Some(None) if the key has been deleted. Some(Some(...)) for a key whose
 	/// value has been set.
 	pub fn child_storage(&self, child_info: &ChildInfo, key: &[u8]) -> Option<Option<&[u8]>> {
-		if let Some(map) = self.prospective.children_default.get(child_info.storage_key()) {
-			if let Some(val) = map.0.get(key) {
-				let size_read = val.value.as_ref().map(|x| x.len() as u64).unwrap_or(0);
-				self.stats.tally_read_modified(size_read);
-				return Some(val.value.as_ref().map(AsRef::as_ref));
-			}
-		}
-
-		if let Some(map) = self.committed.children_default.get(child_info.storage_key()) {
-			if let Some(val) = map.0.get(key) {
-				let size_read = val.value.as_ref().map(|x| x.len() as u64).unwrap_or(0);
-				self.stats.tally_read_modified(size_read);
-				return Some(val.value.as_ref().map(AsRef::as_ref));
-			}
-		}
-
-		None
+		let map = self.children.get(child_info.storage_key())?;
+		let value = map.0.get(key)?.value();
+		let size_read = value.map(|x| x.len() as u64).unwrap_or(0);
+		self.stats.tally_read_modified(size_read);
+		Some(value.map(AsRef::as_ref))
 	}
 
-	/// Inserts the given key-value pair into the prospective change set.
+	/// Set a new value for the specified key.
 	///
-	/// `None` can be used to delete a value specified by the given key.
+	/// Can be rolled back or committed when called inside a transaction.
 	pub(crate) fn set_storage(&mut self, key: StorageKey, val: Option<StorageValue>) {
 		let size_write = val.as_ref().map(|x| x.len() as u64).unwrap_or(0);
 		self.stats.tally_write_overlay(size_write);
-		let extrinsic_index = self.extrinsic_index();
-		let entry = self.prospective.top.entry(key).or_default();
-		entry.value = val;
-
-		if let Some(extrinsic) = extrinsic_index {
-			entry.extrinsics.get_or_insert_with(Default::default)
-				.insert(extrinsic);
-		}
+		self.top.set(key, val, self.extrinsic_index());
 	}
 
-	/// Inserts the given key-value pair into the prospective child change set.
+	/// Set a new value for the specified key and child.
 	///
 	/// `None` can be used to delete a value specified by the given key.
+	///
+	/// Can be rolled back or committed when called inside a transaction.
 	pub(crate) fn set_child_storage(
 		&mut self,
 		child_info: &ChildInfo,
 		key: StorageKey,
 		val: Option<StorageValue>,
 	) {
+		let extrinsic_index = self.extrinsic_index();
 		let size_write = val.as_ref().map(|x| x.len() as u64).unwrap_or(0);
 		self.stats.tally_write_overlay(size_write);
-		let extrinsic_index = self.extrinsic_index();
 		let storage_key = child_info.storage_key().to_vec();
-		let map_entry = self.prospective.children_default.entry(storage_key)
-			.or_insert_with(|| (Default::default(), child_info.to_owned()));
-		let updatable = map_entry.1.try_update(child_info);
+		let top = &self.top;
+		let (changeset, info) = self.children.entry(storage_key).or_insert_with(||
+			(
+				top.spawn_child(),
+				child_info.to_owned()
+			)
+		);
+		let updatable = info.try_update(child_info);
 		debug_assert!(updatable);
-
-		let entry = map_entry.0.entry(key).or_default();
-		entry.value = val;
-
-		if let Some(extrinsic) = extrinsic_index {
-			entry.extrinsics.get_or_insert_with(Default::default)
-				.insert(extrinsic);
-		}
+		changeset.set(key, val, extrinsic_index);
 	}
 
 	/// Clear child storage of given storage key.
 	///
-	/// NOTE that this doesn't take place immediately but written into the prospective
-	/// change set, and still can be reverted by [`discard_prospective`].
-	///
-	/// [`discard_prospective`]: #method.discard_prospective
+	/// Can be rolled back or committed when called inside a transaction.
 	pub(crate) fn clear_child_storage(
 		&mut self,
 		child_info: &ChildInfo,
 	) {
 		let extrinsic_index = self.extrinsic_index();
-		let storage_key = child_info.storage_key();
-		let map_entry = self.prospective.children_default.entry(storage_key.to_vec())
-			.or_insert_with(|| (Default::default(), child_info.to_owned()));
-		let updatable = map_entry.1.try_update(child_info);
+		let storage_key = child_info.storage_key().to_vec();
+		let top = &self.top;
+		let (changeset, info) = self.children.entry(storage_key).or_insert_with(||
+			(
+				top.spawn_child(),
+				child_info.to_owned()
+			)
+		);
+		let updatable = info.try_update(child_info);
 		debug_assert!(updatable);
-
-		map_entry.0.values_mut().for_each(|e| {
-			if let Some(extrinsic) = extrinsic_index {
-				e.extrinsics.get_or_insert_with(Default::default)
-					.insert(extrinsic);
-			}
-
-			e.value = None;
-		});
-
-		if let Some((committed_map, _child_info)) = self.committed.children_default.get(storage_key) {
-			for (key, value) in committed_map.iter() {
-				if !map_entry.0.contains_key(key) {
-					map_entry.0.insert(key.clone(), OverlayedValue {
-						value: None,
-						extrinsics: extrinsic_index.map(|i| {
-							let mut e = value.extrinsics.clone()
-								.unwrap_or_else(|| BTreeSet::default());
-							e.insert(i);
-							e
-						}),
-					});
-				}
-			}
-		}
+		changeset.clear_where(|_, _| true, extrinsic_index);
 	}
 
 	/// Removes all key-value pairs which keys share the given prefix.
 	///
-	/// NOTE that this doesn't take place immediately but written into the prospective
-	/// change set, and still can be reverted by [`discard_prospective`].
-	///
-	/// [`discard_prospective`]: #method.discard_prospective
+	/// Can be rolled back or committed when called inside a transaction.
 	pub(crate) fn clear_prefix(&mut self, prefix: &[u8]) {
-		let extrinsic_index = self.extrinsic_index();
-
-		// Iterate over all prospective and mark all keys that share
-		// the given prefix as removed (None).
-		for (key, entry) in self.prospective.top.iter_mut() {
-			if key.starts_with(prefix) {
-				entry.value = None;
-
-				if let Some(extrinsic) = extrinsic_index {
-					entry.extrinsics.get_or_insert_with(Default::default)
-						.insert(extrinsic);
-				}
-			}
-		}
-
-		// Then do the same with keys from committed changes.
-		// NOTE that we are making changes in the prospective change set.
-		for key in self.committed.top.keys() {
-			if key.starts_with(prefix) {
-				let entry = self.prospective.top.entry(key.clone()).or_default();
-				entry.value = None;
-
-				if let Some(extrinsic) = extrinsic_index {
-					entry.extrinsics.get_or_insert_with(Default::default)
-						.insert(extrinsic);
-				}
-			}
-		}
+		self.top.clear_where(|key, _| key.starts_with(prefix), self.extrinsic_index());
 	}
 
+	/// Removes all key-value pairs which keys share the given prefix.
+	///
+	/// Can be rolled back or committed when called inside a transaction
 	pub(crate) fn clear_child_prefix(
 		&mut self,
 		child_info: &ChildInfo,
 		prefix: &[u8],
 	) {
 		let extrinsic_index = self.extrinsic_index();
-		let storage_key = child_info.storage_key();
-		let map_entry = self.prospective.children_default.entry(storage_key.to_vec())
-			.or_insert_with(|| (Default::default(), child_info.to_owned()));
-		let updatable = map_entry.1.try_update(child_info);
+		let storage_key = child_info.storage_key().to_vec();
+		let top = &self.top;
+		let (changeset, info) = self.children.entry(storage_key).or_insert_with(||
+			(
+				top.spawn_child(),
+				child_info.to_owned()
+			)
+		);
+		let updatable = info.try_update(child_info);
 		debug_assert!(updatable);
+		changeset.clear_where(|key, _| key.starts_with(prefix), extrinsic_index);
+	}
 
-		for (key, entry) in map_entry.0.iter_mut() {
-			if key.starts_with(prefix) {
-				entry.value = None;
+	/// Returns the current nesting depth of the transaction stack.
+	///
+	/// A value of zero means that no transaction is open and changes are committed on write.
+	pub fn transaction_depth(&self) -> usize {
+		// The top changeset and all child changesets transact in lockstep and are
+		// therefore always at the same transaction depth.
+		self.top.transaction_depth()
+	}
 
-				if let Some(extrinsic) = extrinsic_index {
-					entry.extrinsics.get_or_insert_with(Default::default)
-						.insert(extrinsic);
-				}
-			}
-		}
-
-		if let Some((child_committed, _child_info)) = self.committed.children_default.get(storage_key) {
-			// Then do the same with keys from committed changes.
-			// NOTE that we are making changes in the prospective change set.
-			for key in child_committed.keys() {
-				if key.starts_with(prefix) {
-					let entry = map_entry.0.entry(key.clone()).or_default();
-					entry.value = None;
-
-					if let Some(extrinsic) = extrinsic_index {
-						entry.extrinsics.get_or_insert_with(Default::default)
-							.insert(extrinsic);
-					}
-				}
-			}
+	/// Start a new nested transaction.
+	///
+	/// This allows to either commit or roll back all changes that where made while this
+	/// transaction was open. Any transaction must be closed by either `rollback_transaction` or
+	/// `commit_transaction` before this overlay can be converted into storage changes.
+	///
+	/// Changes made without any open transaction are committed immediatly.
+	pub fn start_transaction(&mut self) {
+		self.top.start_transaction();
+		for (_, (changeset, _)) in self.children.iter_mut() {
+			changeset.start_transaction();
 		}
 	}
 
-	/// Discard prospective changes to state.
-	pub fn discard_prospective(&mut self) {
-		self.prospective.clear();
+	/// Rollback the last transaction started by `start_transaction`.
+	///
+	/// Any changes made during that transaction are discarded. Returns an error if
+	/// there is no open transaction that can be rolled back.
+	pub fn rollback_transaction(&mut self) -> Result<(), NoOpenTransaction> {
+		self.top.rollback_transaction()?;
+		self.children.retain(|_, (changeset, _)| {
+			changeset.rollback_transaction()
+				.expect("Top and children changesets are started in lockstep; qed");
+			!changeset.is_empty()
+		});
+		Ok(())
 	}
 
-	/// Commit prospective changes to state.
-	pub fn commit_prospective(&mut self) {
-		if self.committed.is_empty() {
-			mem::swap(&mut self.prospective, &mut self.committed);
-		} else {
-			let top_to_commit = mem::replace(&mut self.prospective.top, BTreeMap::new());
-			for (key, val) in top_to_commit.into_iter() {
-				let entry = self.committed.top.entry(key).or_default();
-				entry.value = val.value;
-
-				if let Some(prospective_extrinsics) = val.extrinsics {
-					entry.extrinsics.get_or_insert_with(Default::default)
-						.extend(prospective_extrinsics);
-				}
-			}
-			for (storage_key, (map, child_info)) in self.prospective.children_default.drain() {
-				let child_content = self.committed.children_default.entry(storage_key)
-					.or_insert_with(|| (Default::default(), child_info));
-				// No update to child info at this point (will be needed for deletion).
-				for (key, val) in map.into_iter() {
-					let entry = child_content.0.entry(key).or_default();
-					entry.value = val.value;
-
-					if let Some(prospective_extrinsics) = val.extrinsics {
-						entry.extrinsics.get_or_insert_with(Default::default)
-							.extend(prospective_extrinsics);
-					}
-				}
-			}
+	/// Commit the last transaction started by `start_transaction`.
+	///
+	/// Any changes made during that transaction are committed. Returns an error if there
+	/// is no open transaction that can be committed.
+	pub fn commit_transaction(&mut self) -> Result<(), NoOpenTransaction> {
+		self.top.commit_transaction()?;
+		for (_, (changeset, _)) in self.children.iter_mut() {
+			changeset.commit_transaction()
+				.expect("Top and children changesets are started in lockstep; qed");
 		}
+		Ok(())
 	}
 
-	/// Consume `OverlayedChanges` and take committed set.
+	/// Call this before transfering control to the runtime.
+	///
+	/// This protects all existing transactions from being removed by the runtime.
+	/// Calling this while already inside the runtime will return an error.
+	pub fn enter_runtime(&mut self) -> Result<(), AlreadyInRuntime> {
+		self.top.enter_runtime()?;
+		for (_, (changeset, _)) in self.children.iter_mut() {
+			changeset.enter_runtime()
+				.expect("Top and children changesets are entering runtime in lockstep; qed")
+		}
+		Ok(())
+	}
+
+	/// Call this when control returns from the runtime.
+	///
+	/// This commits all dangling transaction left open by the runtime.
+	/// Calling this while outside the runtime will return an error.
+	pub fn exit_runtime(&mut self) -> Result<(), NotInRuntime> {
+		self.top.exit_runtime()?;
+		for (_, (changeset, _)) in self.children.iter_mut() {
+			changeset.exit_runtime()
+				.expect("Top and children changesets are entering runtime in lockstep; qed");
+		}
+		Ok(())
+	}
+
+	/// Consume all changes (top + children) and return them.
+	///
+	/// After calling this function no more changes are contained in this changeset.
 	///
 	/// Panics:
-	/// Will panic if there are any uncommitted prospective changes.
+	/// Panics if `transaction_depth() > 0`
 	fn drain_committed(&mut self) -> (
 		impl Iterator<Item=(StorageKey, Option<StorageValue>)>,
 		impl Iterator<Item=(StorageKey, (impl Iterator<Item=(StorageKey, Option<StorageValue>)>, ChildInfo))>,
 	) {
-		assert!(self.prospective.is_empty());
+		use std::mem::take;
 		(
-			std::mem::take(&mut self.committed.top)
-				.into_iter()
-				.map(|(k, v)| (k, v.value)),
-			std::mem::take(&mut self.committed.children_default)
-				.into_iter()
-				.map(|(sk, (v, ci))| (sk, (v.into_iter().map(|(k, v)| (k, v.value)), ci))),
+			take(&mut self.top).drain_commited(),
+			take(&mut self.children).into_iter()
+				.map(|(key, (val, info))| (
+						key,
+						(val.drain_commited(), info)
+					)
+				),
 		)
 	}
 
-	/// Get an iterator over all pending and committed child tries in the overlay.
-	pub fn child_infos(&self) -> impl IntoIterator<Item=&ChildInfo> {
-		self.committed.children_default.iter()
-			.chain(self.prospective.children_default.iter())
-			.map(|(_, v)| &v.1).collect::<BTreeSet<_>>()
+	/// Get an iterator over all child changes as seen by the current transaction.
+	pub fn children(&self)
+		-> impl Iterator<Item=(impl Iterator<Item=(&StorageKey, &OverlayedValue)>, &ChildInfo)> {
+		self.children.iter().map(|(_, v)| (v.0.changes(), &v.1))
 	}
 
-	/// Get an iterator over all pending and committed changes.
-	///
-	/// Supplying `None` for `child_info` will only return changes that are in the top
-	/// trie. Specifying some `child_info` will return only the changes in that
-	/// child trie.
-	pub fn changes(&self, child_info: Option<&ChildInfo>)
-		-> impl Iterator<Item=(&StorageKey, &OverlayedValue)>
-	{
-		let (committed, prospective) = if let Some(child_info) = child_info {
-			match child_info.child_type() {
-				ChildType::ParentKeyId => (
-					self.committed.children_default.get(child_info.storage_key()).map(|c| &c.0),
-					self.prospective.children_default.get(child_info.storage_key()).map(|c| &c.0),
-				),
-			}
-		} else {
-			(Some(&self.committed.top), Some(&self.prospective.top))
-		};
-		committed.into_iter().flatten().chain(prospective.into_iter().flatten())
+	/// Get an iterator over all top changes as been by the current transaction.
+	pub fn changes(&self) -> impl Iterator<Item=(&StorageKey, &OverlayedValue)> {
+		self.top.changes()
 	}
 
-	/// Return a clone of the currently pending changes.
-	pub fn clone_pending(&self) -> OverlayedChangeSet {
-		self.prospective.clone()
-	}
-
-	/// Replace the currently pending changes.
-	pub fn replace_pending(&mut self, pending: OverlayedChangeSet) {
-		self.prospective = pending;
+	/// Get an optional iterator over all child changes stored under the supplied key.
+	pub fn child_changes(&self, key: &[u8])
+		-> Option<(impl Iterator<Item=(&StorageKey, &OverlayedValue)>, &ChildInfo)> {
+		self.children.get(key).map(|(overlay, info)| (overlay.changes(), info))
 	}
 
 	/// Convert this instance with all changes into a [`StorageChanges`] instance.
@@ -607,10 +469,7 @@ impl OverlayedChanges {
 	/// Inserts storage entry responsible for current extrinsic index.
 	#[cfg(test)]
 	pub(crate) fn set_extrinsic_index(&mut self, extrinsic_index: u32) {
-		self.prospective.top.insert(EXTRINSIC_INDEX.to_vec(), OverlayedValue {
-			value: Some(extrinsic_index.encode()),
-			extrinsics: None,
-		});
+		self.top.set(EXTRINSIC_INDEX.to_vec(), Some(extrinsic_index.encode()), None);
 	}
 
 	/// Returns current extrinsic index to use in changes trie construction.
@@ -629,7 +488,8 @@ impl OverlayedChanges {
 		}
 	}
 
-	/// Generate the storage root using `backend` and all changes from `prospective` and `committed`.
+	/// Generate the storage root using `backend` and all changes
+	/// as seen by the current transaction.
 	///
 	/// Returns the storage root and caches storage transaction in the given `cache`.
 	pub fn storage_root<H: Hasher, N: BlockNumber, B: Backend<H>>(
@@ -639,35 +499,13 @@ impl OverlayedChanges {
 	) -> H::Out
 		where H::Out: Ord + Encode,
 	{
-		let child_storage_keys = self.prospective.children_default.keys()
-				.chain(self.committed.children_default.keys());
-		let child_delta_iter = child_storage_keys.map(|storage_key|
-			(
-				self.default_child_info(storage_key)
-					.expect("child info initialized in either committed or prospective"),
-				self.committed.children_default.get(storage_key)
-					.into_iter()
-					.flat_map(|(map, _)|
-						map.iter().map(|(k, v)| (&k[..], v.value().map(|v| &v[..])))
-					)
-					.chain(
-						self.prospective.children_default.get(storage_key)
-							.into_iter()
-							.flat_map(|(map, _)|
-								map.iter().map(|(k, v)| (&k[..], v.value().map(|v| &v[..])))
-							)
-					),
-			)
-		);
+		let delta = self.changes().map(|(k, v)| (&k[..], v.value().map(|v| &v[..])));
+		let child_delta = self.children()
+			.map(|(changes, info)| (info, changes.map(
+				|(k, v)| (&k[..], v.value().map(|v| &v[..]))
+			)));
 
-		// compute and memoize
-		let delta = self.committed
-			.top
-			.iter()
-			.map(|(k, v)| (&k[..], v.value().map(|v| &v[..])))
-			.chain(self.prospective.top.iter().map(|(k, v)| (&k[..], v.value().map(|v| &v[..]))));
-
-		let (root, transaction) = backend.full_storage_root(delta, child_delta_iter);
+		let (root, transaction) = backend.full_storage_root(delta, child_delta);
 
 		cache.transaction = Some(transaction);
 		cache.transaction_storage_root = Some(root);
@@ -704,41 +542,10 @@ impl OverlayedChanges {
 		})
 	}
 
-	/// Get child info for a storage key.
-	/// Take the latest value so prospective first.
-	pub fn default_child_info(&self, storage_key: &[u8]) -> Option<&ChildInfo> {
-		if let Some((_, ci)) = self.prospective.children_default.get(storage_key) {
-			return Some(&ci);
-		}
-		if let Some((_, ci)) = self.committed.children_default.get(storage_key) {
-			return Some(&ci);
-		}
-		None
-	}
-
 	/// Returns the next (in lexicographic order) storage key in the overlayed alongside its value.
 	/// If no value is next then `None` is returned.
 	pub fn next_storage_key_change(&self, key: &[u8]) -> Option<(&[u8], &OverlayedValue)> {
-		let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded);
-
-		let next_prospective_key = self.prospective.top
-			.range::<[u8], _>(range)
-			.next()
-			.map(|(k, v)| (&k[..], v));
-
-		let next_committed_key = self.committed.top
-			.range::<[u8], _>(range)
-			.next()
-			.map(|(k, v)| (&k[..], v));
-
-		match (next_committed_key, next_prospective_key) {
-			// Committed is strictly less than prospective
-			(Some(committed_key), Some(prospective_key)) if committed_key.0 < prospective_key.0 =>
-				Some(committed_key),
-			(committed_key, None) => committed_key,
-			// Prospective key is less or equal to committed or committed doesn't exist
-			(_, prospective_key) => prospective_key,
-		}
+		self.top.next_change(key)
 	}
 
 	/// Returns the next (in lexicographic order) child storage key in the overlayed alongside its
@@ -748,48 +555,32 @@ impl OverlayedChanges {
 		storage_key: &[u8],
 		key: &[u8]
 	) -> Option<(&[u8], &OverlayedValue)> {
-		let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded);
-
-		let next_prospective_key = self.prospective.children_default.get(storage_key)
-			.and_then(|(map, _)| map.range::<[u8], _>(range).next().map(|(k, v)| (&k[..], v)));
-
-		let next_committed_key = self.committed.children_default.get(storage_key)
-			.and_then(|(map, _)| map.range::<[u8], _>(range).next().map(|(k, v)| (&k[..], v)));
-
-		match (next_committed_key, next_prospective_key) {
-			// Committed is strictly less than prospective
-			(Some(committed_key), Some(prospective_key)) if committed_key.0 < prospective_key.0 =>
-				Some(committed_key),
-			(committed_key, None) => committed_key,
-			// Prospective key is less or equal to committed or committed doesn't exist
-			(_, prospective_key) => prospective_key,
-		}
-	}
-}
-
-#[cfg(test)]
-impl From<Option<StorageValue>> for OverlayedValue {
-	fn from(value: Option<StorageValue>) -> OverlayedValue {
-		OverlayedValue { value, ..Default::default() }
+		self.children
+			.get(storage_key)
+			.and_then(|(overlay, _)|
+				overlay.next_change(key)
+			)
 	}
 }
 
 #[cfg(test)]
 mod tests {
 	use hex_literal::hex;
-	use sp_core::{
-		Blake2Hasher, traits::Externalities, storage::well_known_keys::EXTRINSIC_INDEX,
-	};
+	use sp_core::{Blake2Hasher, traits::Externalities};
 	use crate::InMemoryBackend;
 	use crate::ext::Ext;
 	use super::*;
+	use std::collections::BTreeMap;
 
-	fn strip_extrinsic_index(map: &BTreeMap<StorageKey, OverlayedValue>)
-		-> BTreeMap<StorageKey, OverlayedValue>
-	{
-		let mut clone = map.clone();
-		clone.remove(&EXTRINSIC_INDEX.to_vec());
-		clone
+	fn assert_extrinsics(
+		overlay: &OverlayedChangeSet,
+		key: impl AsRef<[u8]>,
+		expected: Vec<u32>,
+	) {
+		assert_eq!(
+			overlay.get(key.as_ref()).unwrap().extrinsics().cloned().collect::<Vec<_>>(),
+			expected
+		)
 	}
 
 	#[test]
@@ -800,23 +591,28 @@ mod tests {
 
 		assert!(overlayed.storage(&key).is_none());
 
+		overlayed.start_transaction();
+
 		overlayed.set_storage(key.clone(), Some(vec![1, 2, 3]));
 		assert_eq!(overlayed.storage(&key).unwrap(), Some(&[1, 2, 3][..]));
 
-		overlayed.commit_prospective();
+		overlayed.commit_transaction().unwrap();
+
 		assert_eq!(overlayed.storage(&key).unwrap(), Some(&[1, 2, 3][..]));
 
+		overlayed.start_transaction();
+
 		overlayed.set_storage(key.clone(), Some(vec![]));
 		assert_eq!(overlayed.storage(&key).unwrap(), Some(&[][..]));
 
 		overlayed.set_storage(key.clone(), None);
 		assert!(overlayed.storage(&key).unwrap().is_none());
 
-		overlayed.discard_prospective();
+		overlayed.rollback_transaction().unwrap();
+
 		assert_eq!(overlayed.storage(&key).unwrap(), Some(&[1, 2, 3][..]));
 
 		overlayed.set_storage(key.clone(), None);
-		overlayed.commit_prospective();
 		assert!(overlayed.storage(&key).unwrap().is_none());
 	}
 
@@ -829,18 +625,18 @@ mod tests {
 			(b"doug".to_vec(), b"notadog".to_vec()),
 		].into_iter().collect();
 		let backend = InMemoryBackend::<Blake2Hasher>::from(initial);
-		let mut overlay = OverlayedChanges {
-			committed: vec![
-				(b"dog".to_vec(), Some(b"puppy".to_vec()).into()),
-				(b"dogglesworth".to_vec(), Some(b"catYYY".to_vec()).into()),
-				(b"doug".to_vec(), Some(vec![]).into()),
-			].into_iter().collect(),
-			prospective: vec![
-				(b"dogglesworth".to_vec(), Some(b"cat".to_vec()).into()),
-				(b"doug".to_vec(), None.into()),
-			].into_iter().collect(),
-			..Default::default()
-		};
+		let mut overlay = OverlayedChanges::default();
+		overlay.set_collect_extrinsics(false);
+
+		overlay.start_transaction();
+		overlay.set_storage(b"dog".to_vec(), Some(b"puppy".to_vec()));
+		overlay.set_storage(b"dogglesworth".to_vec(), Some(b"catYYY".to_vec()));
+		overlay.set_storage(b"doug".to_vec(), Some(vec![]));
+		overlay.commit_transaction().unwrap();
+
+		overlay.start_transaction();
+		overlay.set_storage(b"dogglesworth".to_vec(), Some(b"cat".to_vec()));
+		overlay.set_storage(b"doug".to_vec(), None);
 
 		let mut offchain_overlay = Default::default();
 		let mut cache = StorageTransactionCache::default();
@@ -862,6 +658,8 @@ mod tests {
 		let mut overlay = OverlayedChanges::default();
 		overlay.set_collect_extrinsics(true);
 
+		overlay.start_transaction();
+
 		overlay.set_storage(vec![100], Some(vec![101]));
 
 		overlay.set_extrinsic_index(0);
@@ -873,17 +671,11 @@ mod tests {
 		overlay.set_extrinsic_index(2);
 		overlay.set_storage(vec![1], Some(vec![6]));
 
-		assert_eq!(strip_extrinsic_index(&overlay.prospective.top),
-			vec![
-				(vec![1], OverlayedValue { value: Some(vec![6]),
-				 extrinsics: Some(vec![0, 2].into_iter().collect()) }),
-				(vec![3], OverlayedValue { value: Some(vec![4]),
-				 extrinsics: Some(vec![1].into_iter().collect()) }),
-				(vec![100], OverlayedValue { value: Some(vec![101]),
-				 extrinsics: Some(vec![NO_EXTRINSIC_INDEX].into_iter().collect()) }),
-			].into_iter().collect());
+		assert_extrinsics(&overlay.top, vec![1], vec![0, 2]);
+		assert_extrinsics(&overlay.top, vec![3], vec![1]);
+		assert_extrinsics(&overlay.top, vec![100], vec![NO_EXTRINSIC_INDEX]);
 
-		overlay.commit_prospective();
+		overlay.start_transaction();
 
 		overlay.set_extrinsic_index(3);
 		overlay.set_storage(vec![3], Some(vec![7]));
@@ -891,75 +683,53 @@ mod tests {
 		overlay.set_extrinsic_index(4);
 		overlay.set_storage(vec![1], Some(vec![8]));
 
-		assert_eq!(strip_extrinsic_index(&overlay.committed.top),
-			vec![
-				(vec![1], OverlayedValue { value: Some(vec![6]),
-				 extrinsics: Some(vec![0, 2].into_iter().collect()) }),
-				(vec![3], OverlayedValue { value: Some(vec![4]),
-				 extrinsics: Some(vec![1].into_iter().collect()) }),
-				(vec![100], OverlayedValue { value: Some(vec![101]),
-				 extrinsics: Some(vec![NO_EXTRINSIC_INDEX].into_iter().collect()) }),
-			].into_iter().collect());
+		assert_extrinsics(&overlay.top, vec![1], vec![0, 2, 4]);
+		assert_extrinsics(&overlay.top, vec![3], vec![1, 3]);
+		assert_extrinsics(&overlay.top, vec![100], vec![NO_EXTRINSIC_INDEX]);
 
-		assert_eq!(strip_extrinsic_index(&overlay.prospective.top),
-			vec![
-				(vec![1], OverlayedValue { value: Some(vec![8]),
-				 extrinsics: Some(vec![4].into_iter().collect()) }),
-				(vec![3], OverlayedValue { value: Some(vec![7]),
-				 extrinsics: Some(vec![3].into_iter().collect()) }),
-			].into_iter().collect());
+		overlay.rollback_transaction().unwrap();
 
-		overlay.commit_prospective();
-
-		assert_eq!(strip_extrinsic_index(&overlay.committed.top),
-			vec![
-				(vec![1], OverlayedValue { value: Some(vec![8]),
-				 extrinsics: Some(vec![0, 2, 4].into_iter().collect()) }),
-				(vec![3], OverlayedValue { value: Some(vec![7]),
-				 extrinsics: Some(vec![1, 3].into_iter().collect()) }),
-				(vec![100], OverlayedValue { value: Some(vec![101]),
-				 extrinsics: Some(vec![NO_EXTRINSIC_INDEX].into_iter().collect()) }),
-			].into_iter().collect());
-
-		assert_eq!(overlay.prospective,
-			Default::default());
+		assert_extrinsics(&overlay.top, vec![1], vec![0, 2]);
+		assert_extrinsics(&overlay.top, vec![3], vec![1]);
+		assert_extrinsics(&overlay.top, vec![100], vec![NO_EXTRINSIC_INDEX]);
 	}
 
 	#[test]
 	fn next_storage_key_change_works() {
 		let mut overlay = OverlayedChanges::default();
+		overlay.start_transaction();
 		overlay.set_storage(vec![20], Some(vec![20]));
 		overlay.set_storage(vec![30], Some(vec![30]));
 		overlay.set_storage(vec![40], Some(vec![40]));
-		overlay.commit_prospective();
+		overlay.commit_transaction().unwrap();
 		overlay.set_storage(vec![10], Some(vec![10]));
 		overlay.set_storage(vec![30], None);
 
 		// next_prospective < next_committed
 		let next_to_5 = overlay.next_storage_key_change(&[5]).unwrap();
 		assert_eq!(next_to_5.0.to_vec(), vec![10]);
-		assert_eq!(next_to_5.1.value, Some(vec![10]));
+		assert_eq!(next_to_5.1.value(), Some(&vec![10]));
 
 		// next_committed < next_prospective
 		let next_to_10 = overlay.next_storage_key_change(&[10]).unwrap();
 		assert_eq!(next_to_10.0.to_vec(), vec![20]);
-		assert_eq!(next_to_10.1.value, Some(vec![20]));
+		assert_eq!(next_to_10.1.value(), Some(&vec![20]));
 
 		// next_committed == next_prospective
 		let next_to_20 = overlay.next_storage_key_change(&[20]).unwrap();
 		assert_eq!(next_to_20.0.to_vec(), vec![30]);
-		assert_eq!(next_to_20.1.value, None);
+		assert_eq!(next_to_20.1.value(), None);
 
 		// next_committed, no next_prospective
 		let next_to_30 = overlay.next_storage_key_change(&[30]).unwrap();
 		assert_eq!(next_to_30.0.to_vec(), vec![40]);
-		assert_eq!(next_to_30.1.value, Some(vec![40]));
+		assert_eq!(next_to_30.1.value(), Some(&vec![40]));
 
 		overlay.set_storage(vec![50], Some(vec![50]));
 		// next_prospective, no next_committed
 		let next_to_40 = overlay.next_storage_key_change(&[40]).unwrap();
 		assert_eq!(next_to_40.0.to_vec(), vec![50]);
-		assert_eq!(next_to_40.1.value, Some(vec![50]));
+		assert_eq!(next_to_40.1.value(), Some(&vec![50]));
 	}
 
 	#[test]
@@ -968,37 +738,38 @@ mod tests {
 		let child_info = &child_info;
 		let child = child_info.storage_key();
 		let mut overlay = OverlayedChanges::default();
+		overlay.start_transaction();
 		overlay.set_child_storage(child_info, vec![20], Some(vec![20]));
 		overlay.set_child_storage(child_info, vec![30], Some(vec![30]));
 		overlay.set_child_storage(child_info, vec![40], Some(vec![40]));
-		overlay.commit_prospective();
+		overlay.commit_transaction().unwrap();
 		overlay.set_child_storage(child_info, vec![10], Some(vec![10]));
 		overlay.set_child_storage(child_info, vec![30], None);
 
 		// next_prospective < next_committed
 		let next_to_5 = overlay.next_child_storage_key_change(child, &[5]).unwrap();
 		assert_eq!(next_to_5.0.to_vec(), vec![10]);
-		assert_eq!(next_to_5.1.value, Some(vec![10]));
+		assert_eq!(next_to_5.1.value(), Some(&vec![10]));
 
 		// next_committed < next_prospective
 		let next_to_10 = overlay.next_child_storage_key_change(child, &[10]).unwrap();
 		assert_eq!(next_to_10.0.to_vec(), vec![20]);
-		assert_eq!(next_to_10.1.value, Some(vec![20]));
+		assert_eq!(next_to_10.1.value(), Some(&vec![20]));
 
 		// next_committed == next_prospective
 		let next_to_20 = overlay.next_child_storage_key_change(child, &[20]).unwrap();
 		assert_eq!(next_to_20.0.to_vec(), vec![30]);
-		assert_eq!(next_to_20.1.value, None);
+		assert_eq!(next_to_20.1.value(), None);
 
 		// next_committed, no next_prospective
 		let next_to_30 = overlay.next_child_storage_key_change(child, &[30]).unwrap();
 		assert_eq!(next_to_30.0.to_vec(), vec![40]);
-		assert_eq!(next_to_30.1.value, Some(vec![40]));
+		assert_eq!(next_to_30.1.value(), Some(&vec![40]));
 
 		overlay.set_child_storage(child_info, vec![50], Some(vec![50]));
 		// next_prospective, no next_committed
 		let next_to_40 = overlay.next_child_storage_key_change(child, &[40]).unwrap();
 		assert_eq!(next_to_40.0.to_vec(), vec![50]);
-		assert_eq!(next_to_40.1.value, Some(vec![50]));
+		assert_eq!(next_to_40.1.value(), Some(&vec![50]));
 	}
 }

substrate/primitives/state-machine/src/read_only.rs

@@ -170,6 +170,18 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 		unimplemented!("storage_changes_root is not supported in ReadOnlyExternalities")
 	}
 
+	fn storage_start_transaction(&mut self) {
+		unimplemented!("Transactions are not supported by ReadOnlyExternalities");
+	}
+
+	fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
+		unimplemented!("Transactions are not supported by ReadOnlyExternalities");
+	}
+
+	fn storage_commit_transaction(&mut self) -> Result<(), ()> {
+		unimplemented!("Transactions are not supported by ReadOnlyExternalities");
+	}
+
 	fn wipe(&mut self) {}
 
 	fn commit(&mut self) {}

substrate/primitives/state-machine/src/testing.rs

@@ -153,15 +153,15 @@ impl<H: Hasher, N: ChangesTrieBlockNumber> TestExternalities<H, N>
 
 	/// Return a new backend with all pending value.
 	pub fn commit_all(&self) -> InMemoryBackend<H> {
-		let top: Vec<_> = self.overlay.changes(None)
+		let top: Vec<_> = self.overlay.changes()
 			.map(|(k, v)| (k.clone(), v.value().cloned()))
 			.collect();
 		let mut transaction = vec![(None, top)];
 
-		for child_info in self.overlay.child_infos() {
+		for (child_changes, child_info) in self.overlay.children() {
 			transaction.push((
 				Some(child_info.clone()),
-				self.overlay.changes(Some(child_info))
+				child_changes
 					.map(|(k, v)| (k.clone(), v.value().cloned()))
 					.collect(),
 			))