Changes tries build cache (#2933)

* changes tries build cache

added CT build cache test

* fix lines width

* fixed some grumbles

* clear cache when: digests disabled, top-level or skewed digest is built

* cached_changed_keys -> with_cached_changed_keys

substrate/core/client/db/src/lib.rs

@@ -36,7 +36,7 @@ mod utils;
 use std::sync::Arc;
 use std::path::PathBuf;
 use std::io;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use client::backend::NewBlockState;
 use client::blockchain::HeaderBackend;
@@ -58,7 +58,7 @@ use sr_primitives::traits::{
 };
 use state_machine::backend::Backend as StateBackend;
 use executor::RuntimeInfo;
-use state_machine::{CodeExecutor, DBValue};
+use state_machine::{CodeExecutor, DBValue, ChangesTrieTransaction, ChangesTrieCacheAction, ChangesTrieBuildCache};
 use crate::utils::{Meta, db_err, meta_keys, read_db, block_id_to_lookup_key, read_meta};
 use client::leaves::{LeafSet, FinalizationDisplaced};
 use client::children;
@@ -405,6 +405,7 @@ pub struct BlockImportOperation<Block: BlockT, H: Hasher> {
 	storage_updates: StorageCollection,
 	child_storage_updates: ChildStorageCollection,
 	changes_trie_updates: MemoryDB<H>,
+	changes_trie_cache_update: Option<ChangesTrieCacheAction<H::Out, NumberFor<Block>>>,
 	pending_block: Option<PendingBlock<Block>>,
 	aux_ops: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 	finalized_blocks: Vec<(BlockId<Block>, Option<Justification>)>,
@@ -487,8 +488,12 @@ where Block: BlockT<Hash=H256>,
 		Ok(root)
 	}
 
-	fn update_changes_trie(&mut self, update: MemoryDB<Blake2Hasher>) -> Result<(), client::error::Error> {
-		self.changes_trie_updates = update;
+	fn update_changes_trie(
+		&mut self,
+		update: ChangesTrieTransaction<Blake2Hasher, NumberFor<Block>>,
+	) -> Result<(), client::error::Error> {
+		self.changes_trie_updates = update.0;
+		self.changes_trie_cache_update = Some(update.1);
 		Ok(())
 	}
 
@@ -565,6 +570,7 @@ pub struct DbChangesTrieStorage<Block: BlockT> {
 	db: Arc<dyn KeyValueDB>,
 	meta: Arc<RwLock<Meta<NumberFor<Block>, Block::Hash>>>,
 	min_blocks_to_keep: Option<u32>,
+	cache: RwLock<ChangesTrieBuildCache<Block::Hash, NumberFor<Block>>>,
 	_phantom: ::std::marker::PhantomData<Block>,
 }
 
@@ -576,6 +582,11 @@ impl<Block: BlockT<Hash=H256>> DbChangesTrieStorage<Block> {
 		}
 	}
 
+	/// Commit changes into changes trie build cache.
+	pub fn commit_cache(&self, cache_update: ChangesTrieCacheAction<Block::Hash, NumberFor<Block>>) {
+		self.cache.write().perform(cache_update);
+	}
+
 	/// Prune obsolete changes tries.
 	pub fn prune(
 		&self,
@@ -699,6 +710,14 @@ where
 		self
 	}
 
+	fn with_cached_changed_keys(
+		&self,
+		root: &H256,
+		functor: &mut dyn FnMut(&HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>),
+	) -> bool {
+		self.cache.read().with_changed_keys(root, functor)
+	}
+
 	fn get(&self, key: &H256, _prefix: Prefix) -> Result<Option<DBValue>, String> {
 		self.db.get(columns::CHANGES_TRIE, &key[..])
 			.map_err(|err| format!("{}", err))
@@ -783,6 +802,7 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 			db,
 			meta,
 			min_blocks_to_keep: if is_archive_pruning { None } else { Some(MIN_BLOCKS_TO_KEEP_CHANGES_TRIES_FOR) },
+			cache: RwLock::new(ChangesTrieBuildCache::new()),
 			_phantom: Default::default(),
 		};
 
@@ -1098,7 +1118,6 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 			self.changes_tries_storage.commit(&mut transaction, changes_trie_updates);
 			let cache = operation.old_state.release(); // release state reference so that it can be finalized
 
-
 			if finalized {
 				// TODO: ensure best chain contains this block.
 				self.ensure_sequential_finalization(header, Some(last_finalized_hash))?;
@@ -1155,6 +1174,10 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 
 		let write_result = self.storage.db.write(transaction).map_err(db_err);
 
+		if let Some(changes_trie_cache_update) = operation.changes_trie_cache_update {
+			self.changes_tries_storage.commit_cache(changes_trie_cache_update);
+		}
+
 		if let Some((number, hash, enacted, retracted, displaced_leaf, is_best, mut cache)) = imported {
 			if let Err(e) = write_result {
 				let mut leaves = self.blockchain.leaves.write();
@@ -1288,6 +1311,7 @@ impl<Block> client::backend::Backend<Block, Blake2Hasher> for Backend<Block> whe
 			storage_updates: Default::default(),
 			child_storage_updates: Default::default(),
 			changes_trie_updates: MemoryDB::default(),
+			changes_trie_cache_update: None,
 			aux_ops: Vec::new(),
 			finalized_blocks: Vec::new(),
 			set_head: None,
@@ -1515,7 +1539,7 @@ mod tests {
 		let mut op = backend.begin_operation().unwrap();
 		backend.begin_state_operation(&mut op, block_id).unwrap();
 		op.set_block_data(header, None, None, NewBlockState::Best).unwrap();
-		op.update_changes_trie(changes_trie_update).unwrap();
+		op.update_changes_trie((changes_trie_update, ChangesTrieCacheAction::Clear)).unwrap();
 		backend.commit_operation(op).unwrap();
 
 		header_hash

substrate/core/client/src/backend.rs

@@ -24,10 +24,9 @@ use primitives::ChangesTrieConfiguration;
 use sr_primitives::{generic::BlockId, Justification, StorageOverlay, ChildrenStorageOverlay};
 use sr_primitives::traits::{Block as BlockT, NumberFor};
 use state_machine::backend::Backend as StateBackend;
-use state_machine::ChangesTrieStorage as StateChangesTrieStorage;
+use state_machine::{ChangesTrieStorage as StateChangesTrieStorage, ChangesTrieTransaction};
 use consensus::{well_known_cache_keys, BlockOrigin};
 use hash_db::Hasher;
-use trie::MemoryDB;
 use parking_lot::Mutex;
 
 /// In memory array of storage values.
@@ -116,7 +115,7 @@ pub trait BlockImportOperation<Block, H> where
 		child_update: ChildStorageCollection,
 	) -> error::Result<()>;
 	/// Inject changes trie data into the database.
-	fn update_changes_trie(&mut self, update: MemoryDB<H>) -> error::Result<()>;
+	fn update_changes_trie(&mut self, update: ChangesTrieTransaction<H, NumberFor<Block>>) -> error::Result<()>;
 	/// Insert auxiliary keys. Values are `None` if should be deleted.
 	fn insert_aux<I>(&mut self, ops: I) -> error::Result<()>
 		where I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>;

substrate/core/client/src/call_executor.rs

@@ -17,15 +17,15 @@
 use std::{sync::Arc, cmp::Ord, panic::UnwindSafe, result, cell::RefCell, rc::Rc};
 use codec::{Encode, Decode};
 use sr_primitives::{
-	generic::BlockId, traits::Block as BlockT,
+	generic::BlockId, traits::Block as BlockT, traits::NumberFor,
 };
 use state_machine::{
 	self, OverlayedChanges, Ext, CodeExecutor, ExecutionManager,
 	ExecutionStrategy, NeverOffchainExt, backend::Backend as _,
+	ChangesTrieTransaction,
 };
 use executor::{RuntimeVersion, RuntimeInfo, NativeVersion};
 use hash_db::Hasher;
-use trie::MemoryDB;
 use primitives::{offchain, H256, Blake2Hasher, NativeOrEncoded, NeverNativeValue};
 
 use crate::runtime_api::{ProofRecorder, InitializeBlock};
@@ -111,7 +111,14 @@ where
 		manager: ExecutionManager<F>,
 		native_call: Option<NC>,
 		side_effects_handler: Option<&mut O>,
-	) -> Result<(NativeOrEncoded<R>, (S::Transaction, H::Out), Option<MemoryDB<H>>), error::Error>;
+	) -> Result<
+		(
+			NativeOrEncoded<R>,
+			(S::Transaction, H::Out),
+			Option<ChangesTrieTransaction<Blake2Hasher, NumberFor<B>>>
+		),
+		error::Error,
+	>;
 
 	/// Execute a call to a contract on top of given state, gathering execution proof.
 	///
@@ -345,7 +352,7 @@ where
 	) -> error::Result<(
 		NativeOrEncoded<R>,
 		(S::Transaction, <Blake2Hasher as Hasher>::Out),
-		Option<MemoryDB<Blake2Hasher>>,
+		Option<ChangesTrieTransaction<Blake2Hasher, NumberFor<Block>>>,
 	)> {
 		state_machine::new(
 			state,

substrate/core/client/src/client.rs

@@ -43,7 +43,8 @@ use sr_primitives::{
 use state_machine::{
 	DBValue, Backend as StateBackend, CodeExecutor, ChangesTrieAnchorBlockId,
 	ExecutionStrategy, ExecutionManager, prove_read, prove_child_read,
-	ChangesTrieRootsStorage, ChangesTrieStorage, ChangesTrieConfigurationRange,
+	ChangesTrieRootsStorage, ChangesTrieStorage,
+	ChangesTrieTransaction, ChangesTrieConfigurationRange,
 	key_changes, key_changes_proof, OverlayedChanges, NeverOffchainExt,
 };
 use executor::{RuntimeVersion, RuntimeInfo};
@@ -86,7 +87,7 @@ type StorageUpdate<B, Block> = <
 		<B as backend::Backend<Block, Blake2Hasher>>::BlockImportOperation
 			as BlockImportOperation<Block, Blake2Hasher>
 	>::State as state_machine::Backend<Blake2Hasher>>::Transaction;
-type ChangesUpdate = trie::MemoryDB<Blake2Hasher>;
+type ChangesUpdate<Block> = ChangesTrieTransaction<Blake2Hasher, NumberFor<Block>>;
 
 /// Execution strategies settings.
 #[derive(Debug, Clone)]
@@ -635,6 +636,14 @@ impl<B, E, Block, RA> Client<B, E, Block, RA> where
 				self
 			}
 
+			fn with_cached_changed_keys(
+				&self,
+				root: &H256,
+				functor: &mut dyn FnMut(&HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>),
+			) -> bool {
+				self.storage.with_cached_changed_keys(root, functor)
+			}
+
 			fn get(&self, key: &H256, prefix: Prefix) -> Result<Option<DBValue>, String> {
 				self.storage.get(key, prefix)
 			}
@@ -1001,7 +1010,7 @@ impl<B, E, Block, RA> Client<B, E, Block, RA> where
 		body: Option<Vec<Block::Extrinsic>>,
 	) -> error::Result<(
 		Option<StorageUpdate<B, Block>>,
-		Option<Option<ChangesUpdate>>,
+		Option<Option<ChangesUpdate<Block>>>,
 		Option<(
 			Vec<(Vec<u8>, Option<Vec<u8>>)>,
 			Vec<(Vec<u8>, Vec<(Vec<u8>, Option<Vec<u8>>)>)>

substrate/core/client/src/in_mem.rs

@@ -16,7 +16,7 @@
 
 //! In memory client backend
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use parking_lot::{RwLock, Mutex};
 use primitives::{ChangesTrieConfiguration, storage::well_known_keys};
@@ -24,7 +24,7 @@ use sr_primitives::generic::{BlockId, DigestItem};
 use sr_primitives::traits::{Block as BlockT, Header as HeaderT, Zero, NumberFor};
 use sr_primitives::{Justification, StorageOverlay, ChildrenStorageOverlay};
 use state_machine::backend::{Backend as StateBackend, InMemory};
-use state_machine::{self, InMemoryChangesTrieStorage, ChangesTrieAnchorBlockId};
+use state_machine::{self, InMemoryChangesTrieStorage, ChangesTrieAnchorBlockId, ChangesTrieTransaction};
 use hash_db::{Hasher, Prefix};
 use trie::MemoryDB;
 use consensus::well_known_cache_keys::Id as CacheKeyId;
@@ -484,8 +484,8 @@ where
 		Ok(())
 	}
 
-	fn update_changes_trie(&mut self, update: MemoryDB<H>) -> error::Result<()> {
-		self.changes_trie_update = Some(update);
+	fn update_changes_trie(&mut self, update: ChangesTrieTransaction<H, NumberFor<Block>>) -> error::Result<()> {
+		self.changes_trie_update = Some(update.0);
 		Ok(())
 	}
 
@@ -766,6 +766,14 @@ impl<Block, H> state_machine::ChangesTrieStorage<H, NumberFor<Block>> for Change
 		self
 	}
 
+	fn with_cached_changed_keys(
+		&self,
+		_root: &H::Out,
+		_functor: &mut dyn FnMut(&HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>),
+	) -> bool {
+		false
+	}
+
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Result<Option<state_machine::DBValue>, String> {
 		self.0.get(key, prefix)
 	}

substrate/core/client/src/light/backend.rs

@@ -22,7 +22,7 @@ use std::sync::{Arc, Weak};
 use parking_lot::{RwLock, Mutex};
 
 use sr_primitives::{generic::BlockId, Justification, StorageOverlay, ChildrenStorageOverlay};
-use state_machine::{Backend as StateBackend, TrieBackend, backend::InMemory as InMemoryState};
+use state_machine::{Backend as StateBackend, TrieBackend, backend::InMemory as InMemoryState, ChangesTrieTransaction};
 use sr_primitives::traits::{Block as BlockT, NumberFor, Zero, Header};
 use crate::in_mem::{self, check_genesis_storage};
 use crate::backend::{
@@ -284,7 +284,7 @@ where
 		Ok(())
 	}
 
-	fn update_changes_trie(&mut self, _update: MemoryDB<H>) -> ClientResult<()> {
+	fn update_changes_trie(&mut self, _update: ChangesTrieTransaction<H, NumberFor<Block>>) -> ClientResult<()> {
 		// we're not storing anything locally => ignore changes
 		Ok(())
 	}

substrate/core/client/src/light/call_executor.rs

@@ -25,10 +25,10 @@ use std::{
 use codec::{Encode, Decode};
 use primitives::{offchain, H256, Blake2Hasher, convert_hash, NativeOrEncoded};
 use sr_primitives::generic::BlockId;
-use sr_primitives::traits::{One, Block as BlockT, Header as HeaderT};
+use sr_primitives::traits::{One, Block as BlockT, Header as HeaderT, NumberFor};
 use state_machine::{
 	self, Backend as StateBackend, CodeExecutor, OverlayedChanges,
-	ExecutionStrategy, create_proof_check_backend,
+	ExecutionStrategy, ChangesTrieTransaction, create_proof_check_backend,
 	execution_proof_check_on_trie_backend, ExecutionManager, NeverOffchainExt
 };
 use hash_db::Hasher;
@@ -40,7 +40,6 @@ use crate::call_executor::CallExecutor;
 use crate::error::{Error as ClientError, Result as ClientResult};
 use crate::light::fetcher::{Fetcher, RemoteCallRequest};
 use executor::{RuntimeVersion, NativeVersion};
-use trie::MemoryDB;
 
 /// Call executor that executes methods on remote node, querying execution proof
 /// and checking proof by re-executing locally.
@@ -185,7 +184,7 @@ where
 	) -> ClientResult<(
 		NativeOrEncoded<R>,
 		(S::Transaction, <Blake2Hasher as Hasher>::Out),
-		Option<MemoryDB<Blake2Hasher>>,
+		Option<ChangesTrieTransaction<Blake2Hasher, NumberFor<Block>>>,
 	)> {
 		Err(ClientError::NotAvailableOnLightClient.into())
 	}
@@ -365,7 +364,7 @@ impl<Block, B, Remote, Local> CallExecutor<Block, Blake2Hasher> for
 	) -> ClientResult<(
 		NativeOrEncoded<R>,
 		(S::Transaction, <Blake2Hasher as Hasher>::Out),
-		Option<MemoryDB<Blake2Hasher>>,
+		Option<ChangesTrieTransaction<Blake2Hasher, NumberFor<Block>>>,
 	)> {
 		// there's no actual way/need to specify native/wasm execution strategy on light node
 		// => we can safely ignore passed values

substrate/core/client/src/light/fetcher.rs

@@ -460,7 +460,7 @@ struct RootsStorage<'a, Number: SimpleArithmetic, Hash: 'a> {
 impl<'a, H, Number, Hash> ChangesTrieRootsStorage<H, Number> for RootsStorage<'a, Number, Hash>
 	where
 		H: Hasher,
-		Number: ::std::fmt::Display + Clone + SimpleArithmetic + Encode + Decode + Send + Sync + 'static,
+		Number: ::std::fmt::Display + ::std::hash::Hash + Clone + SimpleArithmetic + Encode + Decode + Send + Sync + 'static,
 		Hash: 'a + Send + Sync + Clone + AsRef<[u8]>,
 {
 	fn build_anchor(

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

@@ -33,7 +33,7 @@ use crate::changes_trie::input::ChildIndex;
 ///
 /// Returns Err if storage error has occurred OR if storage haven't returned
 /// required data.
-pub fn prepare_input<'a, B, H, Number>(
+pub(crate) fn prepare_input<'a, B, H, Number>(
 	backend: &'a B,
 	storage: &'a dyn Storage<H, Number>,
 	config: ConfigurationRange<'a, Number>,
@@ -42,6 +42,7 @@ pub fn prepare_input<'a, B, H, Number>(
 ) -> Result<(
 		impl Iterator<Item=InputPair<Number>> + 'a,
 		Vec<(ChildIndex<Number>, impl Iterator<Item=InputPair<Number>> + 'a)>,
+		Vec<Number>,
 	), String>
 	where
 		B: Backend<H>,
@@ -52,12 +53,14 @@ pub fn prepare_input<'a, B, H, Number>(
 	let (extrinsics_input, children_extrinsics_input) = prepare_extrinsics_input(
 		backend,
 		&number,
-		changes)?;
-	let (digest_input, mut children_digest_input) = prepare_digest_input::<H, Number>(
+		changes,
+	)?;
+	let (digest_input, mut children_digest_input, digest_input_blocks) = prepare_digest_input::<H, Number>(
 		parent,
 		config,
-		&number,
-		storage)?;
+		number,
+		storage,
+	)?;
 
 	let mut children_digest = Vec::with_capacity(children_extrinsics_input.len());
 	for (child_index, ext_iter) in children_extrinsics_input.into_iter() {
@@ -79,6 +82,7 @@ pub fn prepare_input<'a, B, H, Number>(
 	Ok((
 		extrinsics_input.chain(digest_input),
 		children_digest,
+		digest_input_blocks,
 	))
 }
 /// Prepare ExtrinsicIndex input pairs.
@@ -186,12 +190,13 @@ fn prepare_extrinsics_input_inner<'a, B, H, Number>(
 /// Prepare DigestIndex input pairs.
 fn prepare_digest_input<'a, H, Number>(
 	parent: &'a AnchorBlockId<H::Out, Number>,
-	config: ConfigurationRange<'a, Number>,
-	block: &Number,
+	config: ConfigurationRange<Number>,
+	block: Number,
 	storage: &'a dyn Storage<H, Number>,
 ) -> Result<(
 		impl Iterator<Item=InputPair<Number>> + 'a,
 		BTreeMap<ChildIndex<Number>, impl Iterator<Item=InputPair<Number>> + 'a>,
+		Vec<Number>,
 	), String>
 	where
 		H: Hasher,
@@ -207,16 +212,16 @@ fn prepare_digest_input<'a, H, Number>(
 		block.clone()
 	};
 
-	digest_build_iterator(config, block_for_digest)
+	let digest_input_blocks = digest_build_iterator(config, block_for_digest).collect::<Vec<_>>();
+	digest_input_blocks.clone().into_iter()
 		.try_fold(
-			(BTreeMap::new(), BTreeMap::new()),
-			move |(mut map, mut child_map), digest_build_block| {
+			(BTreeMap::new(), BTreeMap::new()), move |(mut map, mut child_map), digest_build_block| {
 			let extrinsic_prefix = ExtrinsicIndex::key_neutral_prefix(digest_build_block.clone());
 			let digest_prefix = DigestIndex::key_neutral_prefix(digest_build_block.clone());
 			let child_prefix = ChildIndex::key_neutral_prefix(digest_build_block.clone());
 			let trie_root = storage.root(parent, digest_build_block.clone())?;
 			let trie_root = trie_root.ok_or_else(|| format!("No changes trie root for block {}", digest_build_block.clone()))?;
-			
+
 			let insert_to_map = |map: &mut BTreeMap<_,_>, key: Vec<u8>| {
 				match map.entry(key.clone()) {
 					Entry::Vacant(entry) => {
@@ -239,6 +244,30 @@ fn prepare_digest_input<'a, H, Number>(
 				}
 			};
 
+			// try to get all updated keys from cache
+			let populated_from_cache = storage.with_cached_changed_keys(
+				&trie_root,
+				&mut |changed_keys| {
+					for (storage_key, changed_keys) in changed_keys {
+						let map = match storage_key {
+							Some(storage_key) => child_map
+								.entry(ChildIndex::<Number> {
+									block: block.clone(),
+									storage_key: storage_key.clone(),
+								})
+								.or_default(),
+							None => &mut map,
+						};
+						for changed_key in changed_keys.iter().cloned() {
+							insert_to_map(map, changed_key);
+						}
+					}
+				}
+			);
+			if populated_from_cache {
+				return Ok((map, child_map));
+			}
+
 			let mut children_roots = BTreeMap::<Vec<u8>, _>::new();
 			{
 				let trie_storage = TrieBackendEssence::<_, H>::new(
@@ -272,7 +301,7 @@ fn prepare_digest_input<'a, H, Number>(
 					storage_key,
 				};
 
-				let mut map = child_map.entry(child_index).or_insert_with(|| BTreeMap::<Vec<u8>, _>::new());
+				let mut map = child_map.entry(child_index).or_default();
 				let trie_storage = TrieBackendEssence::<_, H>::new(
 					crate::changes_trie::TrieBackendStorageAdapter(storage),
 					trie_root,
@@ -288,13 +317,12 @@ fn prepare_digest_input<'a, H, Number>(
 					});
 			}
 			Ok((map, child_map))
-
 		})
-
 		.map(|(pairs, child_pairs)| (
 			pairs.into_iter().map(|(_, (k, v))| InputPair::DigestIndex(k, v)),
 			child_pairs.into_iter().map(|(sk, pairs)|
 				(sk, pairs.into_iter().map(|(_, (k, v))| InputPair::DigestIndex(k, v)))).collect(),
+			digest_input_blocks,
 		))
 }
 
@@ -304,8 +332,8 @@ mod test {
 	use primitives::Blake2Hasher;
 	use primitives::storage::well_known_keys::{EXTRINSIC_INDEX};
 	use crate::backend::InMemory;
-	use crate::changes_trie::Configuration;
-	use crate::changes_trie::storage::InMemoryStorage;
+	use crate::changes_trie::{RootsStorage, Configuration, storage::InMemoryStorage};
+	use crate::changes_trie::build_cache::{IncompleteCacheAction, IncompleteCachedBuildData};
 	use crate::overlayed_changes::{OverlayedValue, OverlayedChangeSet};
 	use super::*;
 
@@ -662,4 +690,73 @@ mod test {
 		test_with_zero(16);
 		test_with_zero(17);
 	}
+
+	#[test]
+	fn cache_is_used_when_changes_trie_is_built() {
+		let (backend, mut storage, changes, _) = prepare_for_build(0);
+		let config = changes.changes_trie_config.as_ref().unwrap();
+		let parent = AnchorBlockId { hash: Default::default(), number: 15 };
+
+		// override some actual values from storage with values from the cache
+		//
+		// top-level storage:
+		// (keys 100, 101, 103, 105 are now missing from block#4 => they do not appear
+		// in l2 digest at block 16)
+		//
+		// "1" child storage:
+		// key 102 is now missing from block#4 => it doesn't appear in l2 digest at block 16
+		// (keys 103, 104) are now added to block#4 => they appear in l2 digest at block 16
+		//
+		// "2" child storage:
+		// (keys 105, 106) are now added to block#4 => they appear in l2 digest at block 16
+		let trie_root4 = storage.root(&parent, 4).unwrap().unwrap();
+		let cached_data4 = IncompleteCacheAction::CacheBuildData(IncompleteCachedBuildData::new())
+			.set_digest_input_blocks(vec![1, 2, 3])
+			.insert(None, vec![vec![100], vec![102]].into_iter().collect())
+			.insert(Some(b"1".to_vec()), vec![vec![103], vec![104]].into_iter().collect())
+			.insert(Some(b"2".to_vec()), vec![vec![105], vec![106]].into_iter().collect())
+			.complete(4, &trie_root4);
+		storage.cache_mut().perform(cached_data4);
+
+		let (root_changes_trie_nodes, child_changes_tries_nodes, _) = prepare_input(
+			&backend,
+			&storage,
+			configuration_range(&config, 0),
+			&changes,
+			&parent,
+		).unwrap();
+		assert_eq!(root_changes_trie_nodes.collect::<Vec<InputPair<u64>>>(), vec![
+			InputPair::ExtrinsicIndex(ExtrinsicIndex { block: 16, key: vec![100] }, vec![0, 2, 3]),
+			InputPair::ExtrinsicIndex(ExtrinsicIndex { block: 16, key: vec![101] }, vec![1]),
+			InputPair::ExtrinsicIndex(ExtrinsicIndex { block: 16, key: vec![103] }, vec![0, 1]),
+
+			InputPair::DigestIndex(DigestIndex { block: 16, key: vec![100] }, vec![4]),
+			InputPair::DigestIndex(DigestIndex { block: 16, key: vec![102] }, vec![4]),
+			InputPair::DigestIndex(DigestIndex { block: 16, key: vec![105] }, vec![8]),
+		]);
+
+		let child_changes_tries_nodes = child_changes_tries_nodes
+			.into_iter()
+			.map(|(k, i)| (k, i.collect::<Vec<_>>()))
+			.collect::<BTreeMap<_, _>>();
+		assert_eq!(
+			child_changes_tries_nodes.get(&ChildIndex { block: 16u64, storage_key: b"1".to_vec() }).unwrap(),
+			&vec![
+				InputPair::ExtrinsicIndex(ExtrinsicIndex { block: 16u64, key: vec![100] }, vec![0, 2, 3]),
+
+				InputPair::DigestIndex(DigestIndex { block: 16u64, key: vec![103] }, vec![4]),
+				InputPair::DigestIndex(DigestIndex { block: 16u64, key: vec![104] }, vec![4]),
+			],
+		);
+		assert_eq!(
+			child_changes_tries_nodes.get(&ChildIndex { block: 16u64, storage_key: b"2".to_vec() }).unwrap(),
+			&vec![
+				InputPair::ExtrinsicIndex(ExtrinsicIndex { block: 16u64, key: vec![100] }, vec![0, 2]),
+
+				InputPair::DigestIndex(DigestIndex { block: 16u64, key: vec![105] }, vec![4]),
+				InputPair::DigestIndex(DigestIndex { block: 16u64, key: vec![106] }, vec![4]),
+			],
+		);
+
+	}
 }

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

@@ -0,0 +1,262 @@
+// Copyright 2019 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Changes tries build cache.
+
+use std::collections::{HashMap, HashSet};
+
+/// Changes trie build cache.
+///
+/// Helps to avoid read of changes tries from the database when digest trie
+/// is built. It holds changed keys for every block (indexed by changes trie
+/// root) that could be referenced by future digest items. For digest entries
+/// it also holds keys covered by this digest. Entries for top level digests
+/// are never created, because they'll never be used to build other digests.
+///
+/// Entries are pruned from the cache once digest block that is using this entry
+/// is inserted (because digest block will includes all keys from this entry).
+/// When there's a fork, entries are pruned when first changes trie is inserted.
+pub struct BuildCache<H, N> {
+	/// Map of block (implies changes true) number => changes trie root.
+	roots_by_number: HashMap<N, H>,
+	/// Map of changes trie root => set of storage keys that are in this trie.
+	/// The `Option<Vec<u8>>` in inner `HashMap` stands for the child storage key.
+	/// If it is `None`, then the `HashSet` contains keys changed in top-level storage.
+	/// If it is `Some`, then the `HashSet` contains keys changed in child storage, identified by the key.
+	changed_keys: HashMap<H, HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>>,
+}
+
+/// The action to perform when block-with-changes-trie is imported.
+#[derive(Debug, PartialEq)]
+pub enum CacheAction<H, N> {
+	/// Cache data that has been collected when CT has been built.
+	CacheBuildData(CachedBuildData<H, N>),
+	/// Clear cache from all existing entries.
+	Clear,
+}
+
+/// The data that has been cached during changes trie building.
+#[derive(Debug, PartialEq)]
+pub struct CachedBuildData<H, N> {
+	block: N,
+	trie_root: H,
+	digest_input_blocks: Vec<N>,
+	changed_keys: HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>,
+}
+
+/// The action to perform when block-with-changes-trie is imported.
+#[derive(Debug, PartialEq)]
+pub(crate) enum IncompleteCacheAction<N> {
+	/// Cache data that has been collected when CT has been built.
+	CacheBuildData(IncompleteCachedBuildData<N>),
+	/// Clear cache from all existing entries.
+	Clear,
+}
+
+/// The data (without changes trie root) that has been cached during changes trie building.
+#[derive(Debug, PartialEq)]
+pub(crate) struct IncompleteCachedBuildData<N> {
+	digest_input_blocks: Vec<N>,
+	changed_keys: HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>,
+}
+
+impl<H, N> BuildCache<H, N>
+	where
+		N: Eq + ::std::hash::Hash,
+		H: Eq + ::std::hash::Hash + Clone,
+{
+	/// Create new changes trie build cache.
+	pub fn new() -> Self {
+		BuildCache {
+			roots_by_number: HashMap::new(),
+			changed_keys: HashMap::new(),
+		}
+	}
+
+	/// Get cached changed keys for changes trie with given root.
+	pub fn get(&self, root: &H) -> Option<&HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>> {
+		self.changed_keys.get(&root)
+	}
+
+	/// Execute given functor with cached entry for given block.
+	/// Returns true if the functor has been called and false otherwise.
+	pub fn with_changed_keys(
+		&self,
+		root: &H,
+		functor: &mut dyn FnMut(&HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>),
+	) -> bool {
+		match self.changed_keys.get(&root) {
+			Some(changed_keys) => {
+				functor(changed_keys);
+				true
+			},
+			None => false,
+		}
+	}
+
+	/// Insert data into cache.
+	pub fn perform(&mut self, action: CacheAction<H, N>) {
+		match action {
+			CacheAction::CacheBuildData(data) => {
+				self.roots_by_number.insert(data.block, data.trie_root.clone());
+				self.changed_keys.insert(data.trie_root, data.changed_keys);
+
+				for digest_input_block in data.digest_input_blocks {
+					let digest_input_block_hash = self.roots_by_number.remove(&digest_input_block);
+					if let Some(digest_input_block_hash) = digest_input_block_hash {
+						self.changed_keys.remove(&digest_input_block_hash);
+					}
+				}
+			},
+			CacheAction::Clear => {
+				self.roots_by_number.clear();
+				self.changed_keys.clear();
+			},
+		}
+	}
+}
+
+impl<N> IncompleteCacheAction<N> {
+	/// Returns true if we need to collect changed keys for this action.
+	pub fn collects_changed_keys(&self) -> bool {
+		match *self {
+			IncompleteCacheAction::CacheBuildData(_) => true,
+			IncompleteCacheAction::Clear => false,
+		}
+	}
+
+	/// Complete cache action with computed changes trie root.
+	pub(crate) fn complete<H: Clone>(self, block: N, trie_root: &H) -> CacheAction<H, N> {
+		match self {
+			IncompleteCacheAction::CacheBuildData(build_data) =>
+				CacheAction::CacheBuildData(build_data.complete(block, trie_root.clone())),
+			IncompleteCacheAction::Clear => CacheAction::Clear,
+		}
+	}
+
+	/// Set numbers of blocks that are superseded by this new entry.
+	///
+	/// If/when this build data is committed to the cache, entries for these blocks
+	/// will be removed from the cache.
+	pub(crate) fn set_digest_input_blocks(self, digest_input_blocks: Vec<N>) -> Self {
+		match self {
+			IncompleteCacheAction::CacheBuildData(build_data) =>
+				IncompleteCacheAction::CacheBuildData(build_data.set_digest_input_blocks(digest_input_blocks)),
+			IncompleteCacheAction::Clear => IncompleteCacheAction::Clear,
+		}
+	}
+
+	/// Insert changed keys of given storage into cached data.
+	pub(crate) fn insert(
+		self,
+		storage_key: Option<Vec<u8>>,
+		changed_keys: HashSet<Vec<u8>>,
+	) -> Self {
+		match self {
+			IncompleteCacheAction::CacheBuildData(build_data) =>
+				IncompleteCacheAction::CacheBuildData(build_data.insert(storage_key, changed_keys)),
+			IncompleteCacheAction::Clear => IncompleteCacheAction::Clear,
+		}
+	}
+}
+
+impl<N> IncompleteCachedBuildData<N> {
+	/// Create new cached data.
+	pub(crate) fn new() -> Self {
+		IncompleteCachedBuildData {
+			digest_input_blocks: Vec::new(),
+			changed_keys: HashMap::new(),
+		}
+	}
+
+	fn complete<H>(self, block: N, trie_root: H) -> CachedBuildData<H, N> {
+		CachedBuildData {
+			block,
+			trie_root,
+			digest_input_blocks: self.digest_input_blocks,
+			changed_keys: self.changed_keys,
+		}
+	}
+
+	fn set_digest_input_blocks(mut self, digest_input_blocks: Vec<N>) -> Self {
+		self.digest_input_blocks = digest_input_blocks;
+		self
+	}
+
+	fn insert(
+		mut self,
+		storage_key: Option<Vec<u8>>,
+		changed_keys: HashSet<Vec<u8>>,
+	) -> Self {
+		self.changed_keys.insert(storage_key, changed_keys);
+		self
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn updated_keys_are_stored_when_non_top_level_digest_is_built() {
+		let mut data = IncompleteCachedBuildData::<u32>::new();
+		data = data.insert(None, vec![vec![1]].into_iter().collect());
+		assert_eq!(data.changed_keys.len(), 1);
+
+		let mut cache = BuildCache::new();
+		cache.perform(CacheAction::CacheBuildData(data.complete(1, 1)));
+		assert_eq!(cache.changed_keys.len(), 1);
+		assert_eq!(
+			cache.get(&1).unwrap().clone(),
+			vec![(None, vec![vec![1]].into_iter().collect())].into_iter().collect(),
+		);
+	}
+
+	#[test]
+	fn obsolete_entries_are_purged_when_new_ct_is_built() {
+		let mut cache = BuildCache::<u32, u32>::new();
+		cache.perform(CacheAction::CacheBuildData(IncompleteCachedBuildData::new()
+			.insert(None, vec![vec![1]].into_iter().collect())
+			.complete(1, 1)));
+		cache.perform(CacheAction::CacheBuildData(IncompleteCachedBuildData::new()
+			.insert(None, vec![vec![2]].into_iter().collect())
+			.complete(2, 2)));
+		cache.perform(CacheAction::CacheBuildData(IncompleteCachedBuildData::new()
+			.insert(None, vec![vec![3]].into_iter().collect())
+			.complete(3, 3)));
+
+		assert_eq!(cache.changed_keys.len(), 3);
+
+		cache.perform(CacheAction::CacheBuildData(IncompleteCachedBuildData::new()
+			.set_digest_input_blocks(vec![1, 2, 3])
+			.complete(4, 4)));
+
+		assert_eq!(cache.changed_keys.len(), 1);
+
+		cache.perform(CacheAction::CacheBuildData(IncompleteCachedBuildData::new()
+			.insert(None, vec![vec![8]].into_iter().collect())
+			.complete(8, 8)));
+		cache.perform(CacheAction::CacheBuildData(IncompleteCachedBuildData::new()
+			.insert(None, vec![vec![12]].into_iter().collect())
+			.complete(12, 12)));
+
+		assert_eq!(cache.changed_keys.len(), 3);
+
+		cache.perform(CacheAction::Clear);
+
+		assert_eq!(cache.changed_keys.len(), 0);
+	}
+}

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

@@ -78,6 +78,17 @@ pub enum InputKey<Number: BlockNumber> {
 	ChildIndex(ChildIndex<Number>),
 }
 
+impl<Number: BlockNumber> InputPair<Number> {
+	/// Extract storage key that this pair corresponds to.
+	pub fn key(&self) -> Option<&[u8]> {
+		match *self {
+			InputPair::ExtrinsicIndex(ref key, _) => Some(&key.key),
+			InputPair::DigestIndex(ref key, _) => Some(&key.key),
+			InputPair::ChildIndex(_, _) => None,
+		}
+	}
+}
+
 impl<Number: BlockNumber> Into<(Vec<u8>, Vec<u8>)> for InputPair<Number> {
 	fn into(self) -> (Vec<u8>, Vec<u8>) {
 		match self {

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

@@ -49,6 +49,7 @@
 //! are propagated through its storage root on the top level storage.
 
 mod build;
+mod build_cache;
 mod build_iterator;
 mod changes_iterator;
 mod input;
@@ -56,6 +57,7 @@ mod prune;
 mod storage;
 mod surface_iterator;
 
+pub use self::build_cache::{BuildCache, CachedBuildData, CacheAction};
 pub use self::storage::InMemoryStorage;
 pub use self::changes_iterator::{
 	key_changes, key_changes_proof,
@@ -63,6 +65,7 @@ pub use self::changes_iterator::{
 };
 pub use self::prune::{prune, oldest_non_pruned_trie};
 
+use std::collections::{HashMap, HashSet};
 use std::convert::TryInto;
 use hash_db::{Hasher, Prefix};
 use crate::backend::Backend;
@@ -70,6 +73,7 @@ use num_traits::{One, Zero};
 use codec::{Decode, Encode};
 use primitives;
 use crate::changes_trie::build::prepare_input;
+use crate::changes_trie::build_cache::{IncompleteCachedBuildData, IncompleteCacheAction};
 use crate::overlayed_changes::OverlayedChanges;
 use trie::{MemoryDB, DBValue, TrieMut};
 use trie::trie_types::TrieDBMut;
@@ -84,6 +88,7 @@ pub trait BlockNumber:
 	Clone +
 	From<u32> + TryInto<u32> + One + Zero +
 	PartialEq + Ord +
+	::std::hash::Hash +
 	::std::ops::Add<Self, Output=Self> + ::std::ops::Sub<Self, Output=Self> +
 	::std::ops::Mul<Self, Output=Self> + ::std::ops::Div<Self, Output=Self> +
 	::std::ops::Rem<Self, Output=Self> +
@@ -98,6 +103,7 @@ impl<T> BlockNumber for T where T:
 	Clone +
 	From<u32> + TryInto<u32> + One + Zero +
 	PartialEq + Ord +
+	::std::hash::Hash +
 	::std::ops::Add<Self, Output=Self> + ::std::ops::Sub<Self, Output=Self> +
 	::std::ops::Mul<Self, Output=Self> + ::std::ops::Div<Self, Output=Self> +
 	::std::ops::Rem<Self, Output=Self> +
@@ -128,6 +134,13 @@ pub trait RootsStorage<H: Hasher, Number: BlockNumber>: Send + Sync {
 pub trait Storage<H: Hasher, Number: BlockNumber>: RootsStorage<H, Number> {
 	/// Casts from self reference to RootsStorage reference.
 	fn as_roots_storage(&self) -> &dyn RootsStorage<H, Number>;
+	/// Execute given functor with cached entry for given trie root.
+	/// Returns true if the functor has been called (cache entry exists) and false otherwise.
+	fn with_cached_changed_keys(
+		&self,
+		root: &H::Out,
+		functor: &mut dyn FnMut(&HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>),
+	) -> bool;
 	/// Get a trie node.
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Result<Option<DBValue>, String>;
 }
@@ -166,7 +179,7 @@ pub fn build_changes_trie<'a, B: Backend<H>, S: Storage<H, Number>, H: Hasher, N
 	storage: Option<&'a S>,
 	changes: &OverlayedChanges,
 	parent_hash: H::Out,
-) -> Result<Option<(MemoryDB<H>, H::Out)>, ()>
+) -> Result<Option<(MemoryDB<H>, H::Out, CacheAction<H::Out, Number>)>, ()>
 	where
 		H::Out: Ord + 'static,
 {
@@ -184,15 +197,22 @@ pub fn build_changes_trie<'a, B: Backend<H>, S: Storage<H, Number>, H: Hasher, N
 
 	// build_anchor error should not be considered fatal
 	let parent = storage.build_anchor(parent_hash).map_err(|_| ())?;
+	let block = parent.number.clone() + One::one();
 
 	// storage errors are considered fatal (similar to situations when runtime fetches values from storage)
-	let (input_pairs, child_input_pairs) = prepare_input::<B, H, Number>(
+	let (input_pairs, child_input_pairs, digest_input_blocks) = prepare_input::<B, H, Number>(
 		backend,
 		storage,
-		config,
+		config.clone(),
 		changes,
 		&parent,
 	).expect("changes trie: storage access is not allowed to fail within runtime");
+
+	// prepare cached data
+	let mut cache_action = prepare_cached_build_data(config, block.clone());
+	let needs_changed_keys = cache_action.collects_changed_keys();
+	cache_action = cache_action.set_digest_input_blocks(digest_input_blocks);
+
 	let mut mdb = MemoryDB::default();
 	let mut child_roots = Vec::with_capacity(child_input_pairs.len());
 	for (child_index, input_pairs) in child_input_pairs {
@@ -200,11 +220,24 @@ pub fn build_changes_trie<'a, B: Backend<H>, S: Storage<H, Number>, H: Hasher, N
 		let mut root = Default::default();
 		{
 			let mut trie = TrieDBMut::<H>::new(&mut mdb, &mut root);
-			for (key, value) in input_pairs.map(Into::into) {
+			let mut storage_changed_keys = HashSet::new();
+			for input_pair in input_pairs {
+				if needs_changed_keys {
+					if let Some(key) = input_pair.key() {
+						storage_changed_keys.insert(key.to_vec());
+					}
+				}
+
+				let (key, value) = input_pair.into();
 				not_empty = true;
 				trie.insert(&key, &value)
 					.expect("changes trie: insertion to trie is not allowed to fail within runtime");
 			}
+
+			cache_action = cache_action.insert(
+				Some(child_index.storage_key.clone()),
+				storage_changed_keys,
+			);
 		}
 		if not_empty {
 			child_roots.push(input::InputPair::ChildIndex(child_index, root.as_ref().to_vec()));
@@ -213,11 +246,91 @@ pub fn build_changes_trie<'a, B: Backend<H>, S: Storage<H, Number>, H: Hasher, N
 	let mut root = Default::default();
 	{
 		let mut trie = TrieDBMut::<H>::new(&mut mdb, &mut root);
-		for (key, value) in input_pairs.chain(child_roots.into_iter()).map(Into::into) {
+		for (key, value) in child_roots.into_iter().map(Into::into) {
 			trie.insert(&key, &value)
 				.expect("changes trie: insertion to trie is not allowed to fail within runtime");
 		}
+
+		let mut storage_changed_keys = HashSet::new();
+		for input_pair in input_pairs {
+			if needs_changed_keys {
+				if let Some(key) = input_pair.key() {
+					storage_changed_keys.insert(key.to_vec());
+				}
+			}
+
+			let (key, value) = input_pair.into();
+			trie.insert(&key, &value)
+				.expect("changes trie: insertion to trie is not allowed to fail within runtime");
+		}
+		cache_action = cache_action.insert(
+			None,
+			storage_changed_keys,
+		);
 	}
 
-	Ok(Some((mdb, root)))
+	let cache_action = cache_action.complete(block, &root);
+	Ok(Some((mdb, root, cache_action)))
+}
+
+/// Prepare empty cached build data for given block.
+fn prepare_cached_build_data<Number: BlockNumber>(
+	config: ConfigurationRange<Number>,
+	block: Number,
+) -> IncompleteCacheAction<Number> {
+	// when digests are not enabled in configuration, we do not need to cache anything
+	// because it'll never be used again for building other tries
+	// => let's clear the cache
+	if !config.config.is_digest_build_enabled() {
+		return IncompleteCacheAction::Clear;
+	}
+
+	// when this is the last block where current configuration is active
+	// => let's clear the cache
+	if config.end.as_ref() == Some(&block) {
+		return IncompleteCacheAction::Clear;
+	}
+
+	// we do not need to cache anything when top-level digest trie is created, because
+	// it'll never be used again for building other tries
+	// => let's clear the cache
+	match config.config.digest_level_at_block(config.zero.clone(), block) {
+		Some((digest_level, _, _)) if digest_level == config.config.digest_levels => IncompleteCacheAction::Clear,
+		_ => IncompleteCacheAction::CacheBuildData(IncompleteCachedBuildData::new()),
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn cache_is_cleared_when_digests_are_disabled() {
+		let config = Configuration { digest_interval: 0, digest_levels: 0 };
+		let config_range = ConfigurationRange { zero: 0, end: None, config: &config };
+		assert_eq!(prepare_cached_build_data(config_range, 8u32), IncompleteCacheAction::Clear);
+	}
+
+	#[test]
+	fn build_data_is_cached_when_digests_are_enabled() {
+		let config = Configuration { digest_interval: 8, digest_levels: 2 };
+		let config_range = ConfigurationRange { zero: 0, end: None, config: &config };
+		assert!(prepare_cached_build_data(config_range.clone(), 4u32).collects_changed_keys());
+		assert!(prepare_cached_build_data(config_range.clone(), 7u32).collects_changed_keys());
+		assert!(prepare_cached_build_data(config_range, 8u32).collects_changed_keys());
+	}
+
+	#[test]
+	fn cache_is_cleared_when_digests_are_enabled_and_top_level_digest_is_built() {
+		let config = Configuration { digest_interval: 8, digest_levels: 2 };
+		let config_range = ConfigurationRange { zero: 0, end: None, config: &config };
+		assert_eq!(prepare_cached_build_data(config_range, 64u32), IncompleteCacheAction::Clear);
+	}
+
+	#[test]
+	fn cache_is_cleared_when_end_block_of_configuration_is_built() {
+		let config = Configuration { digest_interval: 8, digest_levels: 2 };
+		let config_range = ConfigurationRange { zero: 0, end: Some(4u32), config: &config };
+		assert_eq!(prepare_cached_build_data(config_range.clone(), 4u32), IncompleteCacheAction::Clear);
+	}
 }

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

@@ -16,16 +16,14 @@
 
 //! Changes trie storage utilities.
 
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashSet, HashMap};
 use hash_db::{Hasher, Prefix, EMPTY_PREFIX};
 use trie::DBValue;
 use trie::MemoryDB;
 use parking_lot::RwLock;
-use crate::changes_trie::{RootsStorage, Storage, AnchorBlockId, BlockNumber};
+use crate::changes_trie::{BuildCache, RootsStorage, Storage, AnchorBlockId, BlockNumber};
 use crate::trie_backend_essence::TrieBackendStorage;
 
-#[cfg(test)]
-use std::collections::HashSet;
 #[cfg(test)]
 use crate::backend::insert_into_memory_db;
 #[cfg(test)]
@@ -34,6 +32,7 @@ use crate::changes_trie::input::{InputPair, ChildIndex};
 /// In-memory implementation of changes trie storage.
 pub struct InMemoryStorage<H: Hasher, Number: BlockNumber> {
 	data: RwLock<InMemoryStorageData<H, Number>>,
+	cache: BuildCache<H::Out, Number>,
 }
 
 /// Adapter for using changes trie storage as a TrieBackendEssence' storage.
@@ -55,6 +54,7 @@ impl<H: Hasher, Number: BlockNumber> InMemoryStorage<H, Number> {
 				roots: BTreeMap::new(),
 				mdb,
 			}),
+			cache: BuildCache::new(),
 		}
 	}
 
@@ -74,6 +74,11 @@ impl<H: Hasher, Number: BlockNumber> InMemoryStorage<H, Number> {
 		Self::with_db(proof_db)
 	}
 
+	/// Get mutable cache reference.
+	pub fn cache_mut(&mut self) -> &mut BuildCache<H::Out, Number> {
+		&mut self.cache
+	}
+
 	/// Create the storage with given blocks.
 	pub fn with_blocks(blocks: Vec<(Number, H::Out)>) -> Self {
 		Self {
@@ -81,6 +86,7 @@ impl<H: Hasher, Number: BlockNumber> InMemoryStorage<H, Number> {
 				roots: blocks.into_iter().collect(),
 				mdb: MemoryDB::default(),
 			}),
+			cache: BuildCache::new(),
 		}
 	}
 
@@ -122,6 +128,7 @@ impl<H: Hasher, Number: BlockNumber> InMemoryStorage<H, Number> {
 				roots,
 				mdb,
 			}),
+			cache: BuildCache::new(),
 		}
 	}
 
@@ -169,6 +176,14 @@ impl<H: Hasher, Number: BlockNumber> Storage<H, Number> for InMemoryStorage<H, N
 		self
 	}
 
+	fn with_cached_changed_keys(
+		&self,
+		root: &H::Out,
+		functor: &mut dyn FnMut(&HashMap<Option<Vec<u8>>, HashSet<Vec<u8>>>),
+	) -> bool {
+		self.cache.with_changed_keys(root, functor)
+	}
+
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Result<Option<DBValue>, String> {
 		MemoryDB::<H>::get(&self.data.read().mdb, key, prefix)
 	}

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

@@ -19,7 +19,10 @@
 use std::{error, fmt, cmp::Ord};
 use log::warn;
 use crate::backend::Backend;
-use crate::changes_trie::{Storage as ChangesTrieStorage, build_changes_trie};
+use crate::changes_trie::{
+	Storage as ChangesTrieStorage, CacheAction as ChangesTrieCacheAction,
+	build_changes_trie,
+};
 use crate::{Externalities, OverlayedChanges, ChildStorageKey};
 use hash_db::Hasher;
 use primitives::{offchain, storage::well_known_keys::is_child_storage_key, traits::BareCryptoStorePtr};
@@ -78,7 +81,7 @@ where
 	/// This differs from `storage_transaction` behavior, because the moment when
 	/// `storage_changes_root` is called matters + we need to remember additional
 	/// data at this moment (block number).
-	changes_trie_transaction: Option<(MemoryDB<H>, H::Out)>,
+	changes_trie_transaction: Option<(MemoryDB<H>, H::Out, ChangesTrieCacheAction<H::Out, N>)>,
 	/// Additional externalities for offchain workers.
 	///
 	/// If None, some methods from the trait might not be supported.
@@ -119,14 +122,14 @@ where
 	}
 
 	/// Get the transaction necessary to update the backend.
-	pub fn transaction(mut self) -> ((B::Transaction, H::Out), Option<MemoryDB<H>>) {
+	pub fn transaction(mut self) -> ((B::Transaction, H::Out), Option<crate::ChangesTrieTransaction<H, N>>) {
 		let _ = self.storage_root();
 
 		let (storage_transaction, changes_trie_transaction) = (
 			self.storage_transaction
 				.expect("storage_transaction always set after calling storage root; qed"),
 			self.changes_trie_transaction
-				.map(|(tx, _)| tx),
+				.map(|(tx, _, cache)| (tx, cache)),
 		);
 
 		(
@@ -355,7 +358,7 @@ where
 			self.overlay,
 			parent_hash,
 		)?;
-		Ok(self.changes_trie_transaction.as_ref().map(|(_, root)| root.clone()))
+		Ok(self.changes_trie_transaction.as_ref().map(|(_, root, _)| root.clone()))
 	}
 
 	fn offchain(&mut self) -> Option<&mut dyn offchain::Externalities> {

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

@@ -50,10 +50,12 @@ pub use changes_trie::{
 	Storage as ChangesTrieStorage,
 	RootsStorage as ChangesTrieRootsStorage,
 	InMemoryStorage as InMemoryChangesTrieStorage,
+	BuildCache as ChangesTrieBuildCache,
+	CacheAction as ChangesTrieCacheAction,
 	ConfigurationRange as ChangesTrieConfigurationRange,
 	key_changes, key_changes_proof, key_changes_proof_check,
 	prune as prune_changes_tries,
-	oldest_non_pruned_trie as oldest_non_pruned_changes_trie
+	oldest_non_pruned_trie as oldest_non_pruned_changes_trie,
 };
 pub use overlayed_changes::OverlayedChanges;
 pub use proving_backend::{
@@ -63,6 +65,11 @@ pub use proving_backend::{
 pub use trie_backend_essence::{TrieBackendStorage, Storage};
 pub use trie_backend::TrieBackend;
 
+/// Type of changes trie transaction.
+pub type ChangesTrieTransaction<H, N> = (
+	MemoryDB<H>,
+	ChangesTrieCacheAction<<H as Hasher>::Out, N>,
+);
 
 /// A wrapper around a child storage key.
 ///
@@ -513,7 +520,7 @@ impl<'a, H, N, B, T, O, Exec> StateMachine<'a, H, N, B, T, O, Exec> where
 	pub fn execute(
 		&mut self,
 		strategy: ExecutionStrategy,
-	) -> Result<(Vec<u8>, (B::Transaction, H::Out), Option<MemoryDB<H>>), Box<dyn Error>> {
+	) -> Result<(Vec<u8>, (B::Transaction, H::Out), Option<ChangesTrieTransaction<H, N>>), Box<dyn Error>> {
 		// We are not giving a native call and thus we are sure that the result can never be a native
 		// value.
 		self.execute_using_consensus_failure_handler::<_, NeverNativeValue, fn() -> _>(
@@ -537,7 +544,7 @@ impl<'a, H, N, B, T, O, Exec> StateMachine<'a, H, N, B, T, O, Exec> where
 		CallResult<R, Exec::Error>,
 		bool,
 		Option<(B::Transaction, H::Out)>,
-		Option<MemoryDB<H>>,
+		Option<ChangesTrieTransaction<H, N>>,
 	) where
 		R: Decode + Encode + PartialEq,
 		NC: FnOnce() -> result::Result<R, &'static str> + UnwindSafe,
@@ -571,7 +578,7 @@ impl<'a, H, N, B, T, O, Exec> StateMachine<'a, H, N, B, T, O, Exec> where
 		mut native_call: Option<NC>,
 		orig_prospective: OverlayedChangeSet,
 		on_consensus_failure: Handler,
-	) -> (CallResult<R, Exec::Error>, Option<(B::Transaction, H::Out)>, Option<MemoryDB<H>>) where
+	) -> (CallResult<R, Exec::Error>, Option<(B::Transaction, H::Out)>, Option<ChangesTrieTransaction<H, N>>) where
 		R: Decode + Encode + PartialEq,
 		NC: FnOnce() -> result::Result<R, &'static str> + UnwindSafe,
 		Handler: FnOnce(
@@ -602,7 +609,7 @@ impl<'a, H, N, B, T, O, Exec> StateMachine<'a, H, N, B, T, O, Exec> where
 		compute_tx: bool,
 		mut native_call: Option<NC>,
 		orig_prospective: OverlayedChangeSet,
-	) -> (CallResult<R, Exec::Error>, Option<(B::Transaction, H::Out)>, Option<MemoryDB<H>>) where
+	) -> (CallResult<R, Exec::Error>, Option<(B::Transaction, H::Out)>, Option<ChangesTrieTransaction<H, N>>) where
 		R: Decode + Encode + PartialEq,
 		NC: FnOnce() -> result::Result<R, &'static str> + UnwindSafe,
 	{
@@ -633,7 +640,7 @@ impl<'a, H, N, B, T, O, Exec> StateMachine<'a, H, N, B, T, O, Exec> where
 	) -> Result<(
 		NativeOrEncoded<R>,
 		Option<(B::Transaction, H::Out)>,
-		Option<MemoryDB<H>>
+		Option<ChangesTrieTransaction<H, N>>,
 	), Box<dyn Error>> where
 		R: Decode + Encode + PartialEq,
 		NC: FnOnce() -> result::Result<R, &'static str> + UnwindSafe,

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

@@ -276,7 +276,7 @@ impl<H, N> Externalities<H> for TestExternalities<H, N>
 			Some(&self.changes_trie_storage),
 			&self.overlay,
 			parent,
-		)?.map(|(_, root)| root))
+		)?.map(|(_, root, _)| root))
 	}
 
 	fn offchain(&mut self) -> Option<&mut dyn offchain::Externalities> {