sp_trie::Recorder: Fix recording the same key for different tries

primitives/state-machine/src/trie_backend.rs

@@ -1117,4 +1117,95 @@ pub mod tests {
 			);
 		}
 	}
+
+	/// Test to ensure that recording the same `key` for different tries works as expected.
+	///
+	/// Each trie stores a different value under the same key. The values are big enough to
+	/// be not inlined with `StateVersion::V1`, this is important to test the expected behavior. The
+	/// trie recorder is expected to differentiate key access based on the different storage roots
+	/// of the tries.
+	#[test]
+	fn recording_same_key_access_in_different_tries() {
+		recording_same_key_access_in_different_tries_inner(StateVersion::V0);
+		recording_same_key_access_in_different_tries_inner(StateVersion::V1);
+	}
+	fn recording_same_key_access_in_different_tries_inner(state_version: StateVersion) {
+		let key = b"test_key".to_vec();
+		// Use some big values to ensure that we don't keep them inline
+		let top_trie_val = vec![1; 1024];
+		let child_trie_1_val = vec![2; 1024];
+		let child_trie_2_val = vec![3; 1024];
+
+		let child_info_1 = ChildInfo::new_default(b"sub1");
+		let child_info_2 = ChildInfo::new_default(b"sub2");
+		let child_info_1 = &child_info_1;
+		let child_info_2 = &child_info_2;
+		let contents = vec![
+			(None, vec![(key.clone(), Some(top_trie_val.clone()))]),
+			(Some(child_info_1.clone()), vec![(key.clone(), Some(child_trie_1_val.clone()))]),
+			(Some(child_info_2.clone()), vec![(key.clone(), Some(child_trie_2_val.clone()))]),
+		];
+		let in_memory = new_in_mem::<BlakeTwo256, PrefixedKey<BlakeTwo256>>();
+		let in_memory = in_memory.update(contents, state_version);
+		let child_storage_keys = vec![child_info_1.to_owned(), child_info_2.to_owned()];
+		let in_memory_root = in_memory
+			.full_storage_root(
+				std::iter::empty(),
+				child_storage_keys.iter().map(|k| (k, std::iter::empty())),
+				state_version,
+			)
+			.0;
+		assert_eq!(in_memory.storage(&key).unwrap().unwrap(), top_trie_val);
+		assert_eq!(in_memory.child_storage(child_info_1, &key).unwrap().unwrap(), child_trie_1_val);
+		assert_eq!(in_memory.child_storage(child_info_2, &key).unwrap().unwrap(), child_trie_2_val);
+
+		for cache in [Some(SharedTrieCache::new(CacheSize::Unlimited)), None] {
+			// Run multiple times to have a different cache conditions.
+			for i in 0..5 {
+				eprintln!("Running with cache {}, iteration {}", cache.is_some(), i);
+
+				if let Some(cache) = &cache {
+					if i == 2 {
+						cache.reset_node_cache();
+					} else if i == 3 {
+						cache.reset_value_cache();
+					}
+				}
+
+				let trie = in_memory.as_trie_backend();
+				let trie_root = trie.storage_root(std::iter::empty(), state_version).0;
+				assert_eq!(in_memory_root, trie_root);
+
+				let proving = TrieBackendBuilder::wrap(&trie)
+					.with_recorder(Recorder::default())
+					.with_optional_cache(cache.as_ref().map(|c| c.local_cache()))
+					.build();
+				assert_eq!(proving.storage(&key).unwrap().unwrap(), top_trie_val);
+				assert_eq!(
+					proving.child_storage(child_info_1, &key).unwrap().unwrap(),
+					child_trie_1_val
+				);
+				assert_eq!(
+					proving.child_storage(child_info_2, &key).unwrap().unwrap(),
+					child_trie_2_val
+				);
+
+				let proof = proving.extract_proof().unwrap();
+
+				let proof_check =
+					create_proof_check_backend::<BlakeTwo256>(in_memory_root.into(), proof)
+						.unwrap();
+
+				assert_eq!(proof_check.storage(&key).unwrap().unwrap(), top_trie_val);
+				assert_eq!(
+					proof_check.child_storage(child_info_1, &key).unwrap().unwrap(),
+					child_trie_1_val
+				);
+				assert_eq!(
+					proof_check.child_storage(child_info_2, &key).unwrap().unwrap(),
+					child_trie_2_val
+				);
+			}
+		}
+	}
 }

primitives/state-machine/src/trie_backend_essence.rs

@@ -177,7 +177,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H>> TrieBackendEss
 		) -> R,
 	) -> R {
 		let storage_root = storage_root.unwrap_or_else(|| self.root);
-		let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder());
+		let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder(storage_root));
 		let recorder = match recorder.as_mut() {
 			Some(recorder) => Some(recorder as &mut dyn TrieRecorder<H::Out>),
 			None => None,
@@ -209,16 +209,19 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H>> TrieBackendEss
 	/// This function must only be used when the operation in `callback` is
 	/// calculating a `storage_root`. It is expected that `callback` returns
 	/// the new storage root. This is required to register the changes in the cache
-	/// for the correct storage root.
+	/// for the correct storage root. The given `storage_root` corresponds to the root of the "old"
+	/// trie. If the value is not given, `self.root` is used.
 	#[cfg(feature = "std")]
 	fn with_recorder_and_cache_for_storage_root<R>(
 		&self,
+		storage_root: Option<H::Out>,
 		callback: impl FnOnce(
 			Option<&mut dyn TrieRecorder<H::Out>>,
 			Option<&mut dyn TrieCache<NodeCodec<H>>>,
 		) -> (Option<H::Out>, R),
 	) -> R {
-		let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder());
+		let storage_root = storage_root.unwrap_or_else(|| self.root);
+		let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder(storage_root));
 		let recorder = match recorder.as_mut() {
 			Some(recorder) => Some(recorder as &mut dyn TrieRecorder<H::Out>),
 			None => None,
@@ -244,6 +247,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H>> TrieBackendEss
 	#[cfg(not(feature = "std"))]
 	fn with_recorder_and_cache_for_storage_root<R>(
 		&self,
+		_: Option<H::Out>,
 		callback: impl FnOnce(
 			Option<&mut dyn TrieRecorder<H::Out>>,
 			Option<&mut dyn TrieCache<NodeCodec<H>>>,
@@ -675,7 +679,7 @@ where
 	) -> (H::Out, S::Overlay) {
 		let mut write_overlay = S::Overlay::default();
 
-		let root = self.with_recorder_and_cache_for_storage_root(|recorder, cache| {
+		let root = self.with_recorder_and_cache_for_storage_root(None, |recorder, cache| {
 			let mut eph = Ephemeral::new(self.backend_storage(), &mut write_overlay);
 			let res = match state_version {
 				StateVersion::V0 => delta_trie_root::<sp_trie::LayoutV0<H>, _, _, _, _, _>(
@@ -719,35 +723,36 @@ where
 			},
 		};
 
-		let new_child_root = self.with_recorder_and_cache_for_storage_root(|recorder, cache| {
-			let mut eph = Ephemeral::new(self.backend_storage(), &mut write_overlay);
-			match match state_version {
-				StateVersion::V0 =>
-					child_delta_trie_root::<sp_trie::LayoutV0<H>, _, _, _, _, _, _>(
-						child_info.keyspace(),
-						&mut eph,
-						child_root,
-						delta,
-						recorder,
-						cache,
-					),
-				StateVersion::V1 =>
-					child_delta_trie_root::<sp_trie::LayoutV1<H>, _, _, _, _, _, _>(
-						child_info.keyspace(),
-						&mut eph,
-						child_root,
-						delta,
-						recorder,
-						cache,
-					),
-			} {
-				Ok(ret) => (Some(ret), ret),
-				Err(e) => {
-					warn!(target: "trie", "Failed to write to trie: {}", e);
-					(None, child_root)
-				},
-			}
-		});
+		let new_child_root =
+			self.with_recorder_and_cache_for_storage_root(Some(child_root), |recorder, cache| {
+				let mut eph = Ephemeral::new(self.backend_storage(), &mut write_overlay);
+				match match state_version {
+					StateVersion::V0 =>
+						child_delta_trie_root::<sp_trie::LayoutV0<H>, _, _, _, _, _, _>(
+							child_info.keyspace(),
+							&mut eph,
+							child_root,
+							delta,
+							recorder,
+							cache,
+						),
+					StateVersion::V1 =>
+						child_delta_trie_root::<sp_trie::LayoutV1<H>, _, _, _, _, _, _>(
+							child_info.keyspace(),
+							&mut eph,
+							child_root,
+							delta,
+							recorder,
+							cache,
+						),
+				} {
+					Ok(ret) => (Some(ret), ret),
+					Err(e) => {
+						warn!(target: "trie", "Failed to write to trie: {}", e);
+						(None, child_root)
+					},
+				}
+			});
 
 		let is_default = new_child_root == default_root;
 

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)

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();