sp_trie::Recorder: Fix recording the same key for different tries (#12636)

With `StateVersion::V1` values over a certain size are not inlined and being put into the backend
with their own hash. When accessing a value in the trie with a recorder, we check if the value is maybe already
recorded and thus, we can check the cache. To check if a value is already recorded, we use the key
of the value to differentiate them. The problem is when there are multiple tries, like multiple
child tries that all have different values under the same key. Before this pull request we didn't
have differentiated for which trie we already had recorded a (key, value) pair. This is now done by also taking
the storage root into account in the recorder to differentiate the different (key, value) pair in
the tries.

substrate/primitives/trie/src/cache/mod.rs

@@ -490,7 +490,7 @@ mod tests {
 
 			{
 				let mut cache = local_cache.as_trie_db_cache(root);
-				let mut recorder = recorder.as_trie_recorder();
+				let mut recorder = recorder.as_trie_recorder(root);
 				let trie = TrieDBBuilder::<Layout>::new(&db, &root)
 					.with_cache(&mut cache)
 					.with_recorder(&mut recorder)
@@ -538,7 +538,7 @@ mod tests {
 			{
 				let mut db = db.clone();
 				let mut cache = local_cache.as_trie_db_cache(root);
-				let mut recorder = recorder.as_trie_recorder();
+				let mut recorder = recorder.as_trie_recorder(root);
 				let mut trie = TrieDBMutBuilder::<Layout>::from_existing(&mut db, &mut new_root)
 					.with_cache(&mut cache)
 					.with_recorder(&mut recorder)

substrate/primitives/trie/src/recorder.rs

@@ -41,7 +41,7 @@ const LOG_TARGET: &str = "trie-recorder";
 /// The internals of [`Recorder`].
 struct RecorderInner<H> {
 	/// The keys for that we have recorded the trie nodes and if we have recorded up to the value.
-	recorded_keys: HashMap<Vec<u8>, RecordedForKey>,
+	recorded_keys: HashMap<H, HashMap<Vec<u8>, RecordedForKey>>,
 	/// The encoded nodes we accessed while recording.
 	accessed_nodes: HashMap<H, Vec<u8>>,
 }
@@ -80,9 +80,16 @@ impl<H: Hasher> Clone for Recorder<H> {
 
 impl<H: Hasher> Recorder<H> {
 	/// Returns the recorder as [`TrieRecorder`](trie_db::TrieRecorder) compatible type.
-	pub fn as_trie_recorder(&self) -> impl trie_db::TrieRecorder<H::Out> + '_ {
+	///
+	/// - `storage_root`: The storage root of the trie for which accesses are recorded. This is
+	///   important when recording access to different tries at once (like top and child tries).
+	pub fn as_trie_recorder(
+		&self,
+		storage_root: H::Out,
+	) -> impl trie_db::TrieRecorder<H::Out> + '_ {
 		TrieRecorder::<H, _> {
 			inner: self.inner.lock(),
+			storage_root,
 			encoded_size_estimation: self.encoded_size_estimation.clone(),
 			_phantom: PhantomData,
 		}
@@ -132,6 +139,7 @@ impl<H: Hasher> Recorder<H> {
 /// The [`TrieRecorder`](trie_db::TrieRecorder) implementation.
 struct TrieRecorder<H: Hasher, I> {
 	inner: I,
+	storage_root: H::Out,
 	encoded_size_estimation: Arc<AtomicUsize>,
 	_phantom: PhantomData<H>,
 }
@@ -191,6 +199,8 @@ impl<H: Hasher, I: DerefMut<Target = RecorderInner<H::Out>>> trie_db::TrieRecord
 
 				self.inner
 					.recorded_keys
+					.entry(self.storage_root)
+					.or_default()
 					.entry(full_key.to_vec())
 					.and_modify(|e| *e = RecordedForKey::Value)
 					.or_insert(RecordedForKey::Value);
@@ -206,6 +216,8 @@ impl<H: Hasher, I: DerefMut<Target = RecorderInner<H::Out>>> trie_db::TrieRecord
 				// accounted for by the recorded node that holds the hash.
 				self.inner
 					.recorded_keys
+					.entry(self.storage_root)
+					.or_default()
 					.entry(full_key.to_vec())
 					.or_insert(RecordedForKey::Hash);
 			},
@@ -221,6 +233,8 @@ impl<H: Hasher, I: DerefMut<Target = RecorderInner<H::Out>>> trie_db::TrieRecord
 				// that the value doesn't exist in the trie.
 				self.inner
 					.recorded_keys
+					.entry(self.storage_root)
+					.or_default()
 					.entry(full_key.to_vec())
 					.and_modify(|e| *e = RecordedForKey::Value)
 					.or_insert(RecordedForKey::Value);
@@ -231,7 +245,11 @@ impl<H: Hasher, I: DerefMut<Target = RecorderInner<H::Out>>> trie_db::TrieRecord
 	}
 
 	fn trie_nodes_recorded_for_key(&self, key: &[u8]) -> RecordedForKey {
-		self.inner.recorded_keys.get(key).copied().unwrap_or(RecordedForKey::None)
+		self.inner
+			.recorded_keys
+			.get(&self.storage_root)
+			.and_then(|k| k.get(key).copied())
+			.unwrap_or(RecordedForKey::None)
 	}
 }
 
@@ -267,7 +285,7 @@ mod tests {
 		let recorder = Recorder::default();
 
 		{
-			let mut trie_recorder = recorder.as_trie_recorder();
+			let mut trie_recorder = recorder.as_trie_recorder(root);
 			let trie = TrieDBBuilder::<Layout>::new(&db, &root)
 				.with_recorder(&mut trie_recorder)
 				.build();