Refactor OverlayedChanges (#5989)

* Hide internal structure of OverlayChanges

* Fix tests for OverlayChanges refactor

* Do not clone pending changes

Discarding prospective changes should be equivalent as a state machine
is not to be called with peding changes.

This will be replaced by a storage transaction that is rolled back before
executing the call the second time removing this constraint.

* Doc fixes

* Remove overlong line

* Revert "Do not clone pending changes"

This reverts commit 4799491f4ac16f8517287a0fcf4a3f84ad56f46e.

* Deduplicate chield tries returned from child_infos()

* Remove redundant type annotation

* Avoid changing the storage root in tests

* Preserve extrinsic indices in trie build test

* Swap order of comitted and prospective in fn child_infos

This is only for consistency and does not impact the result.

* Rename set_pending to replace_pending for clearity

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

@@ -17,7 +17,7 @@
 
 //! Structures and functions required to build changes trie for given block.
 
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::BTreeMap;
 use std::collections::btree_map::Entry;
 use codec::{Decode, Encode};
 use hash_db::Hasher;
@@ -33,7 +33,7 @@ use crate::{
 		input::{InputKey, InputPair, DigestIndex, ExtrinsicIndex, ChildIndex},
 	},
 };
-use sp_core::storage::{ChildInfo, ChildType, PrefixedStorageKey};
+use sp_core::storage::{ChildInfo, PrefixedStorageKey};
 
 /// Prepare input pairs for building a changes trie of given block.
 ///
@@ -106,20 +106,18 @@ fn prepare_extrinsics_input<'a, B, H, Number>(
 		H: Hasher + 'a,
 		Number: BlockNumber,
 {
-
-	let mut children_info = BTreeSet::<ChildInfo>::new();
 	let mut children_result = BTreeMap::new();
-	for (_storage_key, (_map, child_info)) in changes.prospective.children_default.iter()
-		.chain(changes.committed.children_default.iter()) {
-		children_info.insert(child_info.clone());
-	}
-	for child_info in children_info {
+
+	for child_info in changes.child_infos() {
 		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))?;
+		let iter = prepare_extrinsics_input_inner(
+			backend, block, changes,
+			Some(child_info.clone())
+		)?;
 		children_result.insert(child_index, iter);
 	}
 
@@ -139,19 +137,8 @@ fn prepare_extrinsics_input_inner<'a, B, H, Number>(
 		H: Hasher,
 		Number: BlockNumber,
 {
-	let (committed, prospective) = if let Some(child_info) = child_info.as_ref() {
-		match child_info.child_type() {
-			ChildType::ParentKeyId => (
-				changes.committed.children_default.get(child_info.storage_key()).map(|c| &c.0),
-				changes.prospective.children_default.get(child_info.storage_key()).map(|c| &c.0),
-			),
-		}
-	} else {
-		(Some(&changes.committed.top), Some(&changes.prospective.top))
-	};
-	committed.iter().flat_map(|c| c.iter())
-		.chain(prospective.iter().flat_map(|c| c.iter()))
-		.filter(|( _, v)| v.extrinsics.is_some())
+	changes.changes(child_info.as_ref())
+		.filter(|( _, v)| v.extrinsics().is_some())
 		.try_fold(BTreeMap::new(), |mut map: BTreeMap<&[u8], (ExtrinsicIndex<Number>, Vec<u32>)>, (k, v)| {
 			match map.entry(k) {
 				Entry::Vacant(entry) => {
@@ -172,9 +159,10 @@ fn prepare_extrinsics_input_inner<'a, B, H, Number>(
 						}
 					};
 
-					let extrinsics = v.extrinsics.as_ref()
+					let extrinsics = v.extrinsics()
 						.expect("filtered by filter() call above; qed")
-						.iter().cloned().collect();
+						.cloned()
+						.collect();
 					entry.insert((ExtrinsicIndex {
 						block: block.clone(),
 						key: k.to_vec(),
@@ -185,9 +173,8 @@ 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.as_ref()
+						v.extrinsics()
 							.expect("filtered by filter() call above; qed")
-							.iter()
 							.cloned()
 					);
 					extrinsics.sort_unstable();
@@ -341,13 +328,10 @@ fn prepare_digest_input<'a, H, Number>(
 
 #[cfg(test)]
 mod test {
-	use codec::Encode;
 	use sp_core::Blake2Hasher;
-	use sp_core::storage::well_known_keys::EXTRINSIC_INDEX;
 	use crate::InMemoryBackend;
 	use crate::changes_trie::{RootsStorage, Configuration, storage::InMemoryStorage};
 	use crate::changes_trie::build_cache::{IncompleteCacheAction, IncompleteCachedBuildData};
-	use crate::overlayed_changes::{OverlayedValue, OverlayedChangeSet};
 	use super::*;
 
 	fn prepare_for_build(zero: u64) -> (
@@ -367,8 +351,6 @@ mod test {
 			(vec![105], vec![255]),
 		].into_iter().collect::<std::collections::BTreeMap<_, _>>().into();
 		let prefixed_child_trie_key1 = child_info_1.prefixed_storage_key();
-		let child_trie_key1 = child_info_1.storage_key().to_vec();
-		let child_trie_key2 = child_info_2.storage_key().to_vec();
 		let storage = InMemoryStorage::with_inputs(vec![
 			(zero + 1, vec![
 				InputPair::ExtrinsicIndex(ExtrinsicIndex { block: zero + 1, key: vec![100] }, vec![1, 3]),
@@ -418,58 +400,41 @@ mod test {
 				]),
 			]),
 		]);
-		let changes = OverlayedChanges {
-			prospective: OverlayedChangeSet { top: vec![
-				(vec![100], OverlayedValue {
-					value: Some(vec![200]),
-					extrinsics: Some(vec![0, 2].into_iter().collect())
-				}),
-				(vec![103], OverlayedValue {
-					value: None,
-					extrinsics: Some(vec![0, 1].into_iter().collect())
-				}),
-			].into_iter().collect(),
-				children_default: vec![
-					(child_trie_key1.clone(), (vec![
-						(vec![100], OverlayedValue {
-							value: Some(vec![200]),
-							extrinsics: Some(vec![0, 2].into_iter().collect())
-						})
-					].into_iter().collect(), child_info_1.to_owned())),
-					(child_trie_key2, (vec![
-						(vec![100], OverlayedValue {
-							value: Some(vec![200]),
-							extrinsics: Some(vec![0, 2].into_iter().collect())
-						})
-					].into_iter().collect(), child_info_2.to_owned())),
-				].into_iter().collect()
-			},
-			committed: OverlayedChangeSet { top: vec![
-				(EXTRINSIC_INDEX.to_vec(), OverlayedValue {
-					value: Some(3u32.encode()),
-					extrinsics: None,
-				}),
-				(vec![100], OverlayedValue {
-					value: Some(vec![202]),
-					extrinsics: Some(vec![3].into_iter().collect())
-				}),
-				(vec![101], OverlayedValue {
-					value: Some(vec![203]),
-					extrinsics: Some(vec![1].into_iter().collect())
-				}),
-			].into_iter().collect(),
-				children_default: vec![
-					(child_trie_key1, (vec![
-						(vec![100], OverlayedValue {
-							value: Some(vec![202]),
-							extrinsics: Some(vec![3].into_iter().collect())
-						})
-					].into_iter().collect(), child_info_1.to_owned())),
-				].into_iter().collect(),
-			},
-			collect_extrinsics: true,
-			stats: Default::default(),
-		};
+
+		let mut changes = OverlayedChanges::default();
+		changes.set_collect_extrinsics(true);
+
+		changes.set_extrinsic_index(1);
+		changes.set_storage(vec![101], Some(vec![203]));
+
+		changes.set_extrinsic_index(3);
+		changes.set_storage(vec![100], Some(vec![202]));
+		changes.set_child_storage(&child_info_1, vec![100], Some(vec![202]));
+
+		changes.commit_prospective();
+
+		changes.set_extrinsic_index(0);
+		changes.set_storage(vec![100], Some(vec![0]));
+		changes.set_extrinsic_index(2);
+		changes.set_storage(vec![100], Some(vec![200]));
+
+		changes.set_extrinsic_index(0);
+		changes.set_storage(vec![103], Some(vec![0]));
+		changes.set_extrinsic_index(1);
+		changes.set_storage(vec![103], None);
+
+		changes.set_extrinsic_index(0);
+		changes.set_child_storage(&child_info_1, vec![100], Some(vec![0]));
+		changes.set_extrinsic_index(2);
+		changes.set_child_storage(&child_info_1, vec![100], Some(vec![200]));
+
+		changes.set_extrinsic_index(0);
+		changes.set_child_storage(&child_info_2, vec![100], Some(vec![0]));
+		changes.set_extrinsic_index(2);
+		changes.set_child_storage(&child_info_2, vec![100], Some(vec![200]));
+
+		changes.set_extrinsic_index(1);
+
 		let config = Configuration { digest_interval: 4, digest_levels: 2 };
 
 		(backend, storage, changes, config)
@@ -667,10 +632,7 @@ mod test {
 			let (backend, storage, mut changes, config) = prepare_for_build(zero);
 
 			// 110: missing from backend, set to None in overlay
-			changes.prospective.top.insert(vec![110], OverlayedValue {
-				value: None,
-				extrinsics: Some(vec![1].into_iter().collect())
-			});
+			changes.set_storage(vec![110], None);
 
 			let parent = AnchorBlockId { hash: Default::default(), number: zero + 3 };
 			let changes_trie_nodes = prepare_input(

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

@@ -147,8 +147,7 @@ where
 
 		self.backend.pairs().iter()
 			.map(|&(ref k, ref v)| (k.to_vec(), Some(v.to_vec())))
-			.chain(self.overlay.committed.top.clone().into_iter().map(|(k, v)| (k, v.value)))
-			.chain(self.overlay.prospective.top.clone().into_iter().map(|(k, v)| (k, v.value)))
+			.chain(self.overlay.changes(None).map(|(k, v)| (k.clone(), v.value().cloned())))
 			.collect::<HashMap<_, _>>()
 			.into_iter()
 			.filter_map(|(k, maybe_val)| maybe_val.map(|val| (k, val)))
@@ -293,7 +292,7 @@ where
 		match (next_backend_key, next_overlay_key_change) {
 			(Some(backend_key), Some(overlay_key)) if &backend_key[..] < overlay_key.0 => Some(backend_key),
 			(backend_key, None) => backend_key,
-			(_, Some(overlay_key)) => if overlay_key.1.value.is_some() {
+			(_, Some(overlay_key)) => if overlay_key.1.value().is_some() {
 				Some(overlay_key.0.to_vec())
 			} else {
 				self.next_storage_key(&overlay_key.0[..])
@@ -317,7 +316,7 @@ where
 		match (next_backend_key, next_overlay_key_change) {
 			(Some(backend_key), Some(overlay_key)) if &backend_key[..] < overlay_key.0 => Some(backend_key),
 			(backend_key, None) => backend_key,
-			(_, Some(overlay_key)) => if overlay_key.1.value.is_some() {
+			(_, Some(overlay_key)) => if overlay_key.1.value().is_some() {
 				Some(overlay_key.0.to_vec())
 			} else {
 				self.next_child_storage_key(
@@ -479,18 +478,12 @@ where
 			root.encode()
 		} else {
 
-			if let Some(child_info) = self.overlay.default_child_info(storage_key).cloned() {
+			if let Some(child_info) = self.overlay.default_child_info(storage_key) {
 				let (root, is_empty, _) = {
-					let delta = self.overlay.committed.children_default.get(storage_key)
-						.into_iter()
-						.flat_map(|(map, _)| map.clone().into_iter().map(|(k, v)| (k, v.value)))
-						.chain(
-							self.overlay.prospective.children_default.get(storage_key)
-								.into_iter()
-								.flat_map(|(map, _)| map.clone().into_iter().map(|(k, v)| (k, v.value)))
-						);
+					let delta = self.overlay.changes(Some(child_info))
+						.map(|(k, v)| (k.clone(), v.value().cloned()));
 
-					self.backend.child_storage_root(&child_info, delta)
+					self.backend.child_storage_root(child_info, delta)
 				};
 
 				let root = root.encode();
@@ -677,28 +670,19 @@ mod tests {
 		changes_trie::{
 			Configuration as ChangesTrieConfiguration,
 			InMemoryStorage as TestChangesTrieStorage,
-		}, InMemoryBackend, overlayed_changes::OverlayedValue,
+		}, InMemoryBackend,
 	};
 
 	type TestBackend = InMemoryBackend<Blake2Hasher>;
 	type TestExt<'a> = Ext<'a, Blake2Hasher, u64, TestBackend>;
 
 	fn prepare_overlay_with_changes() -> OverlayedChanges {
-		OverlayedChanges {
-			prospective: vec![
-				(EXTRINSIC_INDEX.to_vec(), OverlayedValue {
-					value: Some(3u32.encode()),
-					extrinsics: Some(vec![1].into_iter().collect())
-				}),
-				(vec![1], OverlayedValue {
-					value: Some(vec![100].into_iter().collect()),
-					extrinsics: Some(vec![1].into_iter().collect())
-				}),
-			].into_iter().collect(),
-			committed: Default::default(),
-			collect_extrinsics: true,
-			stats: Default::default(),
-		}
+		let mut changes = OverlayedChanges::default();
+		changes.set_collect_extrinsics(true);
+		changes.set_extrinsic_index(1);
+		changes.set_storage(vec![1], Some(vec![100]));
+		changes.set_storage(EXTRINSIC_INDEX.to_vec(), Some(3u32.encode()));
+		changes
 	}
 
 	fn prepare_offchain_overlay_with_changes() -> OffchainOverlayedChanges {
@@ -755,7 +739,8 @@ mod tests {
 		let mut overlay = prepare_overlay_with_changes();
 		let mut offchain_overlay = prepare_offchain_overlay_with_changes();
 		let mut cache = StorageTransactionCache::default();
-		overlay.prospective.top.get_mut(&vec![1]).unwrap().value = None;
+		overlay.set_collect_extrinsics(false);
+		overlay.set_storage(vec![1], None);
 		let storage = TestChangesTrieStorage::with_blocks(vec![(99, Default::default())]);
 		let state = Some(ChangesTrieState::new(changes_trie_config(), Zero::zero(), &storage));
 		let backend = TestBackend::default();

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

@@ -28,7 +28,6 @@ use sp_core::{
 	storage::ChildInfo, NativeOrEncoded, NeverNativeValue, hexdisplay::HexDisplay,
 	traits::{CodeExecutor, CallInWasmExt, RuntimeCode},
 };
-use overlayed_changes::OverlayedChangeSet;
 use sp_externalities::Extensions;
 
 pub mod backend;
@@ -336,7 +335,6 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 	fn execute_call_with_both_strategy<Handler, R, NC>(
 		&mut self,
 		mut native_call: Option<NC>,
-		orig_prospective: OverlayedChangeSet,
 		on_consensus_failure: Handler,
 	) -> CallResult<R, Exec::Error>
 		where
@@ -347,10 +345,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();
 		let (result, was_native) = self.execute_aux(true, native_call.take());
 
 		if was_native {
-			self.overlay.prospective = orig_prospective.clone();
+			self.overlay.replace_pending(pending_changes);
 			let (wasm_result, _) = self.execute_aux(
 				false,
 				native_call,
@@ -372,12 +371,12 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 	fn execute_call_with_native_else_wasm_strategy<R, NC>(
 		&mut self,
 		mut native_call: Option<NC>,
-		orig_prospective: OverlayedChangeSet,
 	) -> CallResult<R, Exec::Error>
 		where
 			R: Decode + Encode + PartialEq,
 			NC: FnOnce() -> result::Result<R, String> + UnwindSafe,
 	{
+		let pending_changes = self.overlay.clone_pending();
 		let (result, was_native) = self.execute_aux(
 			true,
 			native_call.take(),
@@ -386,7 +385,7 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 		if !was_native || result.is_ok() {
 			result
 		} else {
-			self.overlay.prospective = orig_prospective.clone();
+			self.overlay.replace_pending(pending_changes);
 			let (wasm_result, _) = self.execute_aux(
 				false,
 				native_call,
@@ -421,20 +420,16 @@ impl<'a, B, H, N, Exec> StateMachine<'a, B, H, N, Exec> where
 		self.overlay.set_collect_extrinsics(changes_tries_enabled);
 
 		let result = {
-			let orig_prospective = self.overlay.prospective.clone();
-
 			match manager {
 				ExecutionManager::Both(on_consensus_failure) => {
 					self.execute_call_with_both_strategy(
 						native_call.take(),
-						orig_prospective,
 						on_consensus_failure,
 					)
 				},
 				ExecutionManager::NativeElseWasm => {
 					self.execute_call_with_native_else_wasm_strategy(
 						native_call.take(),
-						orig_prospective,
 					)
 				},
 				ExecutionManager::AlwaysWasm(trust_level) => {
@@ -754,7 +749,6 @@ where
 mod tests {
 	use std::collections::BTreeMap;
 	use codec::Encode;
-	use overlayed_changes::OverlayedValue;
 	use super::*;
 	use super::ext::Ext;
 	use super::changes_trie::Configuration as ChangesTrieConfig;
@@ -977,21 +971,16 @@ mod tests {
 		];
 		let mut state = InMemoryBackend::<BlakeTwo256>::from(initial);
 		let backend = state.as_trie_backend().unwrap();
-		let mut overlay = OverlayedChanges {
-			committed: map![
-				b"aba".to_vec() => OverlayedValue::from(Some(b"1312".to_vec())),
-				b"bab".to_vec() => OverlayedValue::from(Some(b"228".to_vec()))
-			],
-			prospective: map![
-				b"abd".to_vec() => OverlayedValue::from(Some(b"69".to_vec())),
-				b"bbd".to_vec() => OverlayedValue::from(Some(b"42".to_vec()))
-			],
-			..Default::default()
-		};
-		let mut offchain_overlay = Default::default();
 
+		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.set_storage(b"abd".to_vec(), Some(b"69".to_vec()));
+		overlay.set_storage(b"bbd".to_vec(), Some(b"42".to_vec()));
 
 		{
+			let mut offchain_overlay = Default::default();
 			let mut cache = StorageTransactionCache::default();
 			let mut ext = Ext::new(
 				&mut overlay,
@@ -1006,7 +995,8 @@ mod tests {
 		overlay.commit_prospective();
 
 		assert_eq!(
-			overlay.committed,
+			overlay.changes(None).map(|(k, v)| (k.clone(), v.value().cloned()))
+				.collect::<HashMap<_, _>>(),
 			map![
 				b"abc".to_vec() => None.into(),
 				b"abb".to_vec() => None.into(),

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

@@ -30,7 +30,7 @@ use crate::{
 use std::iter::FromIterator;
 use std::collections::{HashMap, BTreeMap, BTreeSet};
 use codec::{Decode, Encode};
-use sp_core::storage::{well_known_keys::EXTRINSIC_INDEX, ChildInfo};
+use sp_core::storage::{well_known_keys::EXTRINSIC_INDEX, ChildInfo, ChildType};
 use sp_core::offchain::storage::OffchainOverlayedChanges;
 use std::{mem, ops};
 
@@ -55,13 +55,13 @@ pub type ChildStorageCollection = Vec<(StorageKey, StorageCollection)>;
 #[derive(Debug, Default, Clone)]
 pub struct OverlayedChanges {
 	/// Changes that are not yet committed.
-	pub(crate) prospective: OverlayedChangeSet,
+	prospective: OverlayedChangeSet,
 	/// Committed changes.
-	pub(crate) committed: OverlayedChangeSet,
+	committed: OverlayedChangeSet,
 	/// True if extrinsics stats must be collected.
-	pub(crate) collect_extrinsics: bool,
+	collect_extrinsics: bool,
 	/// Collect statistic on this execution.
-	pub(crate) stats: StateMachineStats,
+	stats: StateMachineStats,
 }
 
 /// The storage value, used inside OverlayedChanges.
@@ -69,10 +69,10 @@ pub struct OverlayedChanges {
 #[cfg_attr(test, derive(PartialEq))]
 pub struct OverlayedValue {
 	/// Current value. None if value has been deleted.
-	pub value: Option<StorageValue>,
+	value: Option<StorageValue>,
 	/// The set of extrinsic indices where the values has been changed.
 	/// Is filled only if runtime has announced changes trie support.
-	pub extrinsics: Option<BTreeSet<u32>>,
+	extrinsics: Option<BTreeSet<u32>>,
 }
 
 /// Prospective or committed overlayed change set.
@@ -80,9 +80,9 @@ pub struct OverlayedValue {
 #[cfg_attr(test, derive(PartialEq))]
 pub struct OverlayedChangeSet {
 	/// Top level storage changes.
-	pub top: BTreeMap<StorageKey, OverlayedValue>,
+	top: BTreeMap<StorageKey, OverlayedValue>,
 	/// Child storage changes. The map key is the child storage key without the common prefix.
-	pub children_default: HashMap<StorageKey, (BTreeMap<StorageKey, OverlayedValue>, ChildInfo)>,
+	children_default: HashMap<StorageKey, (BTreeMap<StorageKey, OverlayedValue>, ChildInfo)>,
 }
 
 /// A storage changes structure that can be generated by the data collected in [`OverlayedChanges`].
@@ -187,6 +187,18 @@ impl FromIterator<(StorageKey, OverlayedValue)> for OverlayedChangeSet {
 	}
 }
 
+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 {
@@ -499,6 +511,44 @@ impl OverlayedChanges {
 		)
 	}
 
+	/// 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 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())
+	}
+
+	/// 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;
+	}
+
 	/// Convert this instance with all changes into a [`StorageChanges`] instance.
 	pub fn into_storage_changes<
 		B: Backend<H>, H: Hasher, N: BlockNumber

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

@@ -136,21 +136,19 @@ 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.committed.top.clone().into_iter()
-			.chain(self.overlay.prospective.top.clone().into_iter())
-			.map(|(k, v)| (k, v.value)).collect();
+		let top: Vec<_> = self.overlay.changes(None)
+			.map(|(k, v)| (k.clone(), v.value().cloned()))
+			.collect();
 		let mut transaction = vec![(None, top)];
 
-		self.overlay.committed.children_default.clone().into_iter()
-			.chain(self.overlay.prospective.children_default.clone().into_iter())
-			.for_each(|(_storage_key, (map, child_info))| {
-				transaction.push((
-					Some(child_info),
-					map.into_iter()
-						.map(|(k, v)| (k, v.value))
-						.collect::<Vec<_>>(),
-				))
-			});
+		for child_info in self.overlay.child_infos() {
+			transaction.push((
+				Some(child_info.clone()),
+				self.overlay.changes(Some(child_info))
+					.map(|(k, v)| (k.clone(), v.value().cloned()))
+					.collect(),
+			))
+		}
 
 		self.backend.update(transaction)
 	}