Use trie-cache for validate_block

pallets/parachain-system/Cargo.toml

@@ -11,6 +11,9 @@ codec = { package = "parity-scale-codec", version = "3.0.0", default-features =
 environmental = { version = "1.1.4", default-features = false }
 impl-trait-for-tuples = "0.2.1"
 log = { version = "0.4.19", default-features = false }
+trie-db = { version = "0.27.1", default-features = false }
+hash-db = { version = "0.16.0", default-features = false }
+hashbrown = "0.14.0"
 scale-info = { version = "2.9.0", default-features = false, features = ["derive"] }
 
 # Substrate
@@ -69,6 +72,8 @@ std = [
 	"sp-std/std",
 	"sp-trie/std",
 	"xcm/std",
+	"trie-db/std",
+	"hash-db/std"
 ]
 
 runtime-benchmarks = [

pallets/parachain-system/src/validate_block/implementation.rs

@@ -16,7 +16,7 @@
 
 //! The actual implementation of the validate block functionality.
 
-use super::MemoryOptimizedValidationParams;
+use super::{trie_cache, MemoryOptimizedValidationParams};
 use cumulus_primitives_core::{
 	relay_chain::Hash as RHash, ParachainBlockData, PersistedValidationData,
 };
@@ -34,14 +34,22 @@ use sp_runtime::traits::{Block as BlockT, Extrinsic, HashFor, Header as HeaderT}
 use sp_std::prelude::*;
 use sp_trie::MemoryDB;
 
-type TrieBackend<B> = sp_state_machine::TrieBackend<MemoryDB<HashFor<B>>, HashFor<B>>;
+type TrieBackend<B> = sp_state_machine::TrieBackend<
+	MemoryDB<HashFor<B>>,
+	HashFor<B>,
+	trie_cache::CacheProvider<HashFor<B>>,
+>;
 
 type Ext<'a, B> = sp_state_machine::Ext<'a, HashFor<B>, TrieBackend<B>>;
 
 fn with_externalities<F: FnOnce(&mut dyn Externalities) -> R, R>(f: F) -> R {
 	sp_externalities::with_externalities(f).expect("Environmental externalities not set.")
 }
 
+use core::marker::PhantomData;
+use hash_db::Hasher;
+use sp_trie::NodeCodec;
+
 /// Validate the given parachain block.
 ///
 /// This function is doing roughly the following:
@@ -114,10 +122,15 @@ where
 
 	sp_std::mem::drop(storage_proof);
 
+	let cache_provider = trie_cache::CacheProvider::new();
 	// We use the storage root of the `parent_head` to ensure that it is the correct root.
 	// This is already being done above while creating the in-memory db, but let's be paranoid!!
-	let backend =
-		sp_state_machine::TrieBackendBuilder::new(db, *parent_header.state_root()).build();
+	let backend = sp_state_machine::TrieBackendBuilder::new_with_cache(
+		db,
+		*parent_header.state_root(),
+		cache_provider,
+	)
+	.build();
 
 	let _guard = (
 		// Replace storage calls with our own implementations

pallets/parachain-system/src/validate_block/mod.rs

@@ -22,6 +22,10 @@ pub mod implementation;
 #[cfg(test)]
 mod tests;
 
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+mod trie_cache;
+
 #[cfg(not(feature = "std"))]
 #[doc(hidden)]
 pub use bytes;

pallets/parachain-system/src/validate_block/trie_cache.rs

@@ -0,0 +1,110 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+extern crate alloc;
+
+use core::cell::RefCell;
+use hash_db::{HashDB, Hasher};
+use hashbrown::hash_map::{Entry, HashMap};
+use sp_state_machine::{TrieBackendStorage, TrieCacheProvider};
+use sp_std::boxed::Box;
+use sp_trie::NodeCodec;
+use trie_db::{node::NodeOwned, TrieCache, TrieError};
+
+/// Special purpose trie cache implementation that is able to cache an unlimited number
+/// of values. To be used in `validate_block` to serve values and nodes that
+/// have already been loaded and decoded from the storage proof.
+pub(crate) struct SimpleTrieCache<'a, H: Hasher> {
+	node_cache: core::cell::RefMut<'a, HashMap<H::Out, NodeOwned<H::Out>>>,
+	value_cache: core::cell::RefMut<'a, HashMap<Box<[u8]>, trie_db::CachedValue<H::Out>>>,
+}
+
+impl<'a, H: Hasher> trie_db::TrieCache<NodeCodec<H>> for SimpleTrieCache<'a, H> {
+	fn lookup_value_for_key(&mut self, key: &[u8]) -> Option<&trie_db::CachedValue<H::Out>> {
+		self.value_cache.get(key)
+	}
+
+	fn cache_value_for_key(&mut self, key: &[u8], value: trie_db::CachedValue<H::Out>) {
+		self.value_cache.insert(key.into(), value);
+	}
+
+	fn get_or_insert_node(
+		&mut self,
+		hash: <NodeCodec<H> as trie_db::NodeCodec>::HashOut,
+		fetch_node: &mut dyn FnMut() -> trie_db::Result<
+			NodeOwned<H::Out>,
+			H::Out,
+			<NodeCodec<H> as trie_db::NodeCodec>::Error,
+		>,
+	) -> trie_db::Result<&NodeOwned<H::Out>, H::Out, <NodeCodec<H> as trie_db::NodeCodec>::Error> {
+		match self.node_cache.entry(hash) {
+			Entry::Occupied(entry) => Ok(entry.into_mut()),
+			Entry::Vacant(entry) => {
+				let fetched = match fetch_node() {
+					Ok(new_node) => new_node,
+					Err(e) => return Err(e),
+				};
+				Ok(entry.insert(fetched))
+			},
+		}
+	}
+
+	fn get_node(
+		&mut self,
+		hash: &H::Out,
+	) -> Option<&NodeOwned<<NodeCodec<H> as trie_db::NodeCodec>::HashOut>> {
+		self.node_cache.get(hash)
+	}
+}
+
+/// Provider of [`SimpleTrieCache`] instances.
+pub(crate) struct CacheProvider<H: Hasher> {
+	node_cache: RefCell<HashMap<H::Out, NodeOwned<H::Out>>>,
+	value_cache: RefCell<HashMap<Box<[u8]>, trie_db::CachedValue<H::Out>>>,
+}
+
+impl<H: Hasher> CacheProvider<H> {
+	/// Constructs a new instance of CacheProvider with an uninitialized state
+	/// and empty node and value caches.
+	pub fn new() -> Self {
+		CacheProvider { node_cache: Default::default(), value_cache: Default::default() }
+	}
+}
+
+impl<H: Hasher> TrieCacheProvider<H> for CacheProvider<H> {
+	type Cache<'a> = SimpleTrieCache<'a, H> where H: 'a;
+
+	fn as_trie_db_cache(&self, _storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+		SimpleTrieCache {
+			value_cache: self.value_cache.borrow_mut(),
+			node_cache: self.node_cache.borrow_mut(),
+		}
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		SimpleTrieCache {
+			value_cache: self.value_cache.borrow_mut(),
+			node_cache: self.node_cache.borrow_mut(),
+		}
+	}
+
+	fn merge<'a>(&'a self, _other: Self::Cache<'a>, _new_root: <H as Hasher>::Out) {}
+}
+
+// This is safe here since we are single-threaded in WASM
+unsafe impl<H: Hasher> Send for CacheProvider<H> {}
+unsafe impl<H: Hasher> Sync for CacheProvider<H> {}