validate-block: Fix TrieCache implementation

cumulus/pallets/parachain-system/Cargo.toml

@@ -45,6 +45,7 @@ assert_matches = "1.5"
 hex-literal = "0.4.1"
 lazy_static = "1.4"
 rand = "0.8.5"
+futures = "0.3.28"
 
 # Substrate
 sc-client-api = { path = "../../../substrate/client/api" }

cumulus/pallets/parachain-system/src/validate_block/tests.rs

@@ -21,12 +21,13 @@ use cumulus_test_client::{
 	runtime::{
 		self as test_runtime, Block, Hash, Header, TestPalletCall, UncheckedExtrinsic, WASM_BINARY,
 	},
-	transfer, BlockData, BuildParachainBlockData, Client, DefaultTestClientBuilderExt, HeadData,
-	InitBlockBuilder, TestClientBuilder, TestClientBuilderExt, ValidationParams,
+	transfer, BlockData, BlockOrigin, BuildParachainBlockData, Client, ClientBlockImportExt,
+	DefaultTestClientBuilderExt, HeadData, InitBlockBuilder, TestClientBuilder,
+	TestClientBuilderExt, ValidationParams,
 };
 use cumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
 use sp_keyring::AccountKeyring::*;
-use sp_runtime::traits::Header as HeaderT;
+use sp_runtime::traits::{Block as BlockT, Header as HeaderT};
 use std::{env, process::Command};
 
 use crate::validate_block::MemoryOptimizedValidationParams;
@@ -100,7 +101,7 @@ fn build_block_with_witness(
 }
 
 #[test]
-fn validate_block_no_extra_extrinsics() {
+fn validate_block_works() {
 	sp_tracing::try_init_simple();
 
 	let (client, parent_head) = create_test_client();
@@ -290,3 +291,42 @@ fn validation_params_and_memory_optimized_validation_params_encode_and_decode()
 	let decoded = ValidationParams::decode_all(&mut &encoded[..]).unwrap();
 	assert_eq!(decoded, validation_params);
 }
+
+/// Test for ensuring that we are differentiating in the `validation::trie_cache` between different
+/// child tries.
+///
+/// This is achieved by first building a block using `read_and_write_child_tries` that should set
+/// the values in the child tries. In the second step we are building a second block with the same
+/// extrinsic that reads the values from the child tries and it asserts that we read the correct
+/// data from the state.
+#[test]
+fn validate_block_works_with_child_tries() {
+	sp_tracing::try_init_simple();
+
+	let (mut client, parent_head) = create_test_client();
+	let TestBlockData { block, .. } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic(&client, Charlie, TestPalletCall::read_and_write_child_tries {})],
+		parent_head.clone(),
+		Default::default(),
+	);
+	let block = block.into_block();
+
+	futures::executor::block_on(client.import(BlockOrigin::Own, block.clone())).unwrap();
+
+	let parent_head = block.header().clone();
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic(&client, Alice, TestPalletCall::read_and_write_child_tries {})],
+		parent_head.clone(),
+		Default::default(),
+	);
+
+	let header = block.header().clone();
+
+	let res_header =
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}

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

@@ -67,7 +67,13 @@ impl<'a, H: Hasher> trie_db::TrieCache<NodeCodec<H>> for TrieCache<'a, H> {
 /// Provider of [`TrieCache`] instances.
 pub(crate) struct CacheProvider<H: Hasher> {
 	node_cache: RefCell<BTreeMap<H::Out, NodeOwned<H::Out>>>,
-	value_cache: RefCell<BTreeMap<Box<[u8]>, trie_db::CachedValue<H::Out>>>,
+	/// Cache: `storage_root` => `storage_key` => `value`.
+	///
+	/// One `block` can for example use multiple tries (child tries) and we need to distinguish the
+	/// cached (`storage_key`, `value`) between them. For this we are using the `storage_root` to
+	/// distinguish them (even if the storage root is the same for two child tries, it just means
+	/// that both are exactly the same trie and there would happen no collision).
+	value_cache: RefCell<BTreeMap<H::Out, BTreeMap<Box<[u8]>, trie_db::CachedValue<H::Out>>>>,
 }
 
 impl<H: Hasher> CacheProvider<H> {
@@ -81,22 +87,26 @@ impl<H: Hasher> CacheProvider<H> {
 impl<H: Hasher> TrieCacheProvider<H> for CacheProvider<H> {
 	type Cache<'a> = TrieCache<'a, H> where H: 'a;
 
-	fn as_trie_db_cache(&self, _storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+	fn as_trie_db_cache(&self, storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
 		TrieCache {
-			value_cache: Some(self.value_cache.borrow_mut()),
+			value_cache: Some(RefMut::map(self.value_cache.borrow_mut(), |c| {
+				c.entry(storage_root).or_default()
+			})),
 			node_cache: self.node_cache.borrow_mut(),
 		}
 	}
 
 	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
 		// This method is called when we calculate the storage root.
-		// Since we are using a simplified cache architecture,
-		// we do not have separate key spaces for different storage roots.
-		// The value cache is therefore disabled here.
+		// We are not interested in caching new values (as we would throw them away directly after a
+		// block is validated) and thus, we don't pass any `value_cache`.
 		TrieCache { value_cache: None, node_cache: self.node_cache.borrow_mut() }
 	}
 
-	fn merge<'a>(&'a self, _other: Self::Cache<'a>, _new_root: <H as Hasher>::Out) {}
+	fn merge<'a>(&'a self, _other: Self::Cache<'a>, _new_root: <H as Hasher>::Out) {
+		// This is called to merge the `value_cache` from `as_trie_db_mut_cache`, which is not
+		// activated, so we don't need to do anything here.
+	}
 }
 
 // This is safe here since we are single-threaded in WASM

cumulus/test/runtime/src/test_pallet.rs

@@ -49,6 +49,30 @@ pub mod pallet {
 			);
 			Ok(())
 		}
+
+		/// A dispatchable that first reads two values from two different child tries, asserts they
+		/// are the expected values (if the values exist in the state) and then writes two different
+		/// values to these child tries.
+		#[pallet::weight(0)]
+		pub fn read_and_write_child_tries(_: OriginFor<T>) -> DispatchResult {
+			let key = &b"hello"[..];
+			let first_trie = &b"first"[..];
+			let second_trie = &b"second"[..];
+			let first_value = "world1".encode();
+			let second_value = "world2".encode();
+
+			if let Some(res) = sp_io::default_child_storage::get(first_trie, key) {
+				assert_eq!(first_value, res);
+			}
+			if let Some(res) = sp_io::default_child_storage::get(second_trie, key) {
+				assert_eq!(second_value, res);
+			}
+
+			sp_io::default_child_storage::set(first_trie, key, &first_value);
+			sp_io::default_child_storage::set(second_trie, key, &second_value);
+
+			Ok(())
+		}
 	}
 
 	#[derive(frame_support::DefaultNoBound)]

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

@@ -52,7 +52,10 @@ pub trait TrieCacheProvider<H: Hasher> {
 
 	/// Return a [`trie_db::TrieDB`] compatible cache.
 	///
-	/// The `storage_root` parameter should be the storage root of the used trie.
+	/// The `storage_root` parameter *must* be the storage root of the trie this cache is used for.
+	///
+	/// NOTE: Implementors should use the `storage_root` to differentiate between storage keys that
+	/// may belong to different tries.
 	fn as_trie_db_cache(&self, storage_root: H::Out) -> Self::Cache<'_>;
 
 	/// Returns a cache that can be used with a [`trie_db::TrieDBMut`].

substrate/test-utils/client/src/client_ext.rs

@@ -20,9 +20,11 @@
 use sc_client_api::{backend::Finalizer, client::BlockBackend};
 use sc_consensus::{BlockImport, BlockImportParams, ForkChoiceStrategy};
 use sc_service::client::Client;
-use sp_consensus::{BlockOrigin, Error as ConsensusError};
+use sp_consensus::Error as ConsensusError;
 use sp_runtime::{traits::Block as BlockT, Justification, Justifications};
 
+pub use sp_consensus::BlockOrigin;
+
 /// Extension trait for a test client.
 pub trait ClientExt<Block: BlockT>: Sized {
 	/// Finalize a block.

substrate/test-utils/client/src/lib.rs

@@ -21,7 +21,7 @@
 
 pub mod client_ext;
 
-pub use self::client_ext::{ClientBlockImportExt, ClientExt};
+pub use self::client_ext::{BlockOrigin, ClientBlockImportExt, ClientExt};
 pub use sc_client_api::{execution_extensions::ExecutionExtensions, BadBlocks, ForkBlocks};
 pub use sc_client_db::{self, Backend, BlocksPruning};
 pub use sc_executor::{self, NativeElseWasmExecutor, WasmExecutionMethod, WasmExecutor};