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/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)
 	}