Make trie_nodes_recorded_for_key work for inline values

trie-db/src/lib.rs

@@ -173,6 +173,13 @@ pub enum TrieAccess<'a, H> {
 	///
 	/// Should map to [`RecordedForKey::Value`] when checking the recorder.
 	Value { hash: H, value: rstd::borrow::Cow<'a, [u8]>, full_key: &'a [u8] },
+	/// A value was accessed that is stored inline a node.
+	///
+	/// As the value is stored inline there is no need to separately record the value as it is part
+	/// of a node. The given `full_key` is the key to access this value in the trie.
+	///
+	/// Should map to [`RecordedForKey::Value`] when checking the recorder.
+	InlineValue { full_key: &'a [u8] },
 	/// The hash of the value for the given `full_key` was accessed.
 	///
 	/// Should map to [`RecordedForKey::Hash`] when checking the recorder.
@@ -184,7 +191,7 @@ pub enum TrieAccess<'a, H> {
 }
 
 /// Result of [`TrieRecorder::trie_nodes_recorded_for_key`].
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum RecordedForKey {
 	/// We recorded all trie nodes up to the value for a storage key.
 	///

trie-db/src/lookup.rs

@@ -58,7 +58,13 @@ where
 		query: Q,
 	) -> Result<Q::Item, TrieHash<L>, CError<L>> {
 		match v {
-			Value::Inline(value) => Ok(query.decode(&value)),
+			Value::Inline(value) => {
+				if let Some(recorder) = recorder {
+					recorder.record(TrieAccess::InlineValue { full_key });
+				}
+
+				Ok(query.decode(&value))
+			},
 			Value::Node(hash) => {
 				let mut res = TrieHash::<L>::default();
 				res.as_mut().copy_from_slice(hash);
@@ -94,7 +100,13 @@ where
 		recorder: &mut Option<&mut dyn TrieRecorder<TrieHash<L>>>,
 	) -> Result<(Bytes, TrieHash<L>), TrieHash<L>, CError<L>> {
 		match v {
-			ValueOwned::Inline(value, hash) => Ok((value.clone(), hash)),
+			ValueOwned::Inline(value, hash) => {
+				if let Some(recorder) = recorder {
+					recorder.record(TrieAccess::InlineValue { full_key });
+				}
+
+				Ok((value.clone(), hash))
+			},
 			ValueOwned::Node(hash) => {
 				let node = cache.get_or_insert_node(hash, &mut || {
 					let value = db
@@ -388,7 +400,16 @@ where
 				full_key,
 				|v, _, full_key, _, recorder, _| {
 					Ok(match v {
-						Value::Inline(v) => L::Hash::hash(&v),
+						Value::Inline(v) => {
+							if let Some(recoder) = recorder.as_mut() {
+								// We can record this as `InlineValue`, even we are just returning
+								// the `hash`. This is done to prevent requiring to re-record this
+								// key.
+								recoder.record(TrieAccess::InlineValue { full_key });
+							}
+
+							L::Hash::hash(&v)
+						},
 						Value::Node(hash_bytes) => {
 							if let Some(recoder) = recorder.as_mut() {
 								recoder.record(TrieAccess::Hash { full_key });
@@ -432,7 +453,16 @@ where
 				full_key,
 				cache,
 				|value, _, full_key, _, _, recorder| match value {
-					ValueOwned::Inline(value, hash) => Ok((hash, Some(value.clone()))),
+					ValueOwned::Inline(value, hash) => {
+						if let Some(recoder) = recorder.as_mut() {
+							// We can record this as `InlineValue`, even we are just returning
+							// the `hash`. This is done to prevent requiring to re-record this
+							// key.
+							recoder.record(TrieAccess::InlineValue { full_key });
+						}
+
+						Ok((hash, Some(value.clone())))
+					},
 					ValueOwned::Node(hash) => {
 						if let Some(recoder) = recorder.as_mut() {
 							recoder.record(TrieAccess::Hash { full_key });

trie-db/src/recorder.rs

@@ -71,6 +71,9 @@ impl<L: TrieLayout> TrieRecorder<TrieHash<L>> for Recorder<L> {
 			},
 			TrieAccess::NonExisting { full_key } => {
 				// We handle the non existing value/hash like having recorded the value.
+				self.recorded_keys.entry(full_key.to_vec()).insert(RecordedForKey::None);
+			},
+			TrieAccess::InlineValue { full_key } => {
 				self.recorded_keys.entry(full_key.to_vec()).insert(RecordedForKey::Value);
 			},
 		}

trie-db/src/triedbmut.rs

@@ -120,7 +120,7 @@ impl<L: TrieLayout> From<(Bytes, Option<u32>)> for Value<L> {
 	fn from((v, threshold): (Bytes, Option<u32>)) -> Self {
 		match v {
 			value =>
-				if threshold.map(|threshold| value.len() >= threshold as usize).unwrap_or(false) {
+				if threshold.map_or(false, |threshold| value.len() >= threshold as usize) {
 					Value::NewNode(None, value)
 				} else {
 					Value::Inline(value)

trie-db/test/src/triedb.rs

@@ -21,8 +21,8 @@ use reference_trie::{
 	test_layouts, test_layouts_substrate, HashedValueNoExtThreshold, TestTrieCache,
 };
 use trie_db::{
-	encode_compact, CachedValue, DBValue, Lookup, NibbleSlice, Recorder, Trie, TrieCache,
-	TrieDBBuilder, TrieDBMutBuilder, TrieLayout, TrieMut,
+	encode_compact, CachedValue, DBValue, Lookup, NibbleSlice, RecordedForKey, Recorder, Trie,
+	TrieCache, TrieDBBuilder, TrieDBMutBuilder, TrieLayout, TrieMut, TrieRecorder,
 };
 
 type PrefixedMemoryDB<T> =
@@ -1101,3 +1101,75 @@ fn test_record_value() {
 	// leaf with value hash only.
 	assert_eq!(compact_proof[1].len(), 33);
 }
+
+test_layouts!(test_trie_nodes_recorded, test_trie_nodes_recorded_internal);
+fn test_trie_nodes_recorded_internal<T: TrieLayout>() {
+	let key_value = vec![
+		(b"A".to_vec(), vec![1; 64]),
+		(b"AA".to_vec(), vec![2; 64]),
+		(b"AB".to_vec(), vec![3; 4]),
+		(b"B".to_vec(), vec![4; 64]),
+		(b"BC".to_vec(), vec![4; 64]),
+	];
+	const NON_EXISTENT_KEY: &[u8] = &*b"NOT";
+
+	let mut memdb = MemoryDB::<T::Hash, HashKey<_>, DBValue>::default();
+	let mut root = Default::default();
+
+	{
+		let mut t = TrieDBMutBuilder::<T>::new(&mut memdb, &mut root).build();
+		for (key, value) in &key_value {
+			t.insert(key, value).unwrap();
+		}
+	}
+
+	for mut cache in [Some(TestTrieCache::<T>::default()), None] {
+		for get_hash in [true, false] {
+			let mut recorder = Recorder::<T>::default();
+			{
+				let trie = TrieDBBuilder::<T>::new(&memdb, &root)
+					.with_recorder(&mut recorder)
+					.with_optional_cache(cache.as_mut().map(|c| c as &mut _))
+					.build();
+				for (key, _) in &key_value {
+					if get_hash {
+						assert!(trie.get_hash(key).unwrap().is_some());
+					} else {
+						assert!(trie.get(key).unwrap().is_some());
+					}
+				}
+
+				if get_hash {
+					assert!(trie.get_hash(&NON_EXISTENT_KEY).unwrap().is_none());
+				} else {
+					assert!(trie.get(&NON_EXISTENT_KEY).unwrap().is_none());
+				}
+			}
+
+			for (key, value) in &key_value {
+				let recorded = recorder.trie_nodes_recorded_for_key(&key);
+
+				let is_inline = T::MAX_INLINE_VALUE.map_or(true, |m| value.len() < m as usize);
+
+				let expected = if get_hash && !is_inline {
+					RecordedForKey::Hash
+				} else {
+					RecordedForKey::Value
+				};
+
+				assert_eq!(
+					expected,
+					recorded,
+					"{:?} max_inline: {:?} get_hash: {get_hash}",
+					String::from_utf8(key.to_vec()),
+					T::MAX_INLINE_VALUE
+				);
+			}
+
+			assert_eq!(
+				RecordedForKey::None,
+				recorder.trie_nodes_recorded_for_key(&NON_EXISTENT_KEY),
+			);
+		}
+	}
+}