diff --git a/trie-db/src/lib.rs b/trie-db/src/lib.rs index 29b78bbe..c94268ca 100644 --- a/trie-db/src/lib.rs +++ b/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. /// diff --git a/trie-db/src/lookup.rs b/trie-db/src/lookup.rs index 7ca838eb..6f3d1adf 100644 --- a/trie-db/src/lookup.rs +++ b/trie-db/src/lookup.rs @@ -58,7 +58,13 @@ where query: Q, ) -> Result, CError> { 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::::default(); res.as_mut().copy_from_slice(hash); @@ -94,7 +100,13 @@ where recorder: &mut Option<&mut dyn TrieRecorder>>, ) -> Result<(Bytes, TrieHash), TrieHash, CError> { 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 }); diff --git a/trie-db/src/recorder.rs b/trie-db/src/recorder.rs index ba4c77e5..0111bf1f 100644 --- a/trie-db/src/recorder.rs +++ b/trie-db/src/recorder.rs @@ -71,6 +71,9 @@ impl TrieRecorder> for Recorder { }, 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); }, } diff --git a/trie-db/src/triedbmut.rs b/trie-db/src/triedbmut.rs index acbd0fe9..92ec30c1 100644 --- a/trie-db/src/triedbmut.rs +++ b/trie-db/src/triedbmut.rs @@ -120,7 +120,7 @@ impl From<(Bytes, Option)> for Value { fn from((v, threshold): (Bytes, Option)) -> 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) diff --git a/trie-db/test/src/triedb.rs b/trie-db/test/src/triedb.rs index 04608838..253da380 100644 --- a/trie-db/test/src/triedb.rs +++ b/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 = @@ -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() { + 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::, DBValue>::default(); + let mut root = Default::default(); + + { + let mut t = TrieDBMutBuilder::::new(&mut memdb, &mut root).build(); + for (key, value) in &key_value { + t.insert(key, value).unwrap(); + } + } + + for mut cache in [Some(TestTrieCache::::default()), None] { + for get_hash in [true, false] { + let mut recorder = Recorder::::default(); + { + let trie = TrieDBBuilder::::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), + ); + } + } +}