trie-db: Fetch the closest merkle value

trie-db/src/fatdb.rs

@@ -72,6 +72,13 @@ where
 		self.raw.get_with(L::Hash::hash(key).as_ref(), query)
 	}
 
+	fn get_closest_merkle_value(
+		&self,
+		_key: &[u8],
+	) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
+		Ok(None)
+	}
+
 	fn iter<'a>(
 		&'a self,
 	) -> Result<

trie-db/src/lib.rs

@@ -277,6 +277,12 @@ pub trait Trie<L: TrieLayout> {
 		query: Q,
 	) -> Result<Option<Q::Item>, TrieHash<L>, CError<L>>;
 
+	/// Returns the merkle value of the closest descendant node of the key.
+	fn get_closest_merkle_value(
+		&self,
+		key: &[u8],
+	) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>>;
+
 	/// Returns a depth-first iterator over the elements of trie.
 	fn iter<'a>(
 		&'a self,
@@ -404,6 +410,13 @@ impl<'db, 'cache, L: TrieLayout> Trie<L> for TrieKinds<'db, 'cache, L> {
 		wrapper!(self, get_with, key, query)
 	}
 
+	fn get_closest_merkle_value(
+		&self,
+		key: &[u8],
+	) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
+		wrapper!(self, get_closest_merkle_value, key)
+	}
+
 	fn iter<'a>(
 		&'a self,
 	) -> Result<

trie-db/src/lookup.rs

@@ -134,6 +134,19 @@ where
 		}
 	}
 
+	/// Look up the closest merkle value.
+	///
+	/// When the provided key leads to a node, then the merkle value of that node
+	/// is returned. However, if the key does not lead to a node, then the closest
+	/// merkle value is returned.
+	pub fn look_up_merkle(
+		self,
+		full_key: &[u8],
+		nibble_key: NibbleSlice,
+	) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
+		self.lookup_first_descendent_without_cache(nibble_key, full_key)
+	}
+
 	/// Look up the given `nibble_key`.
 	///
 	/// If the value is found, it will be passed to the given function to decode or copy.
@@ -655,4 +668,139 @@ where
 		}
 		Ok(None)
 	}
+
+	/// Look up the merkle value (hash) of the node that is the closest descendant for the provided
+	/// key.
+	///
+	/// When the provided key leads to a node, then the merkle value of that node
+	/// is returned. However, if the key does not lead to a node, then the merkle value
+	/// of the closest descendant is returned. `None` if no such descendant exists.
+	fn lookup_first_descendent_without_cache(
+		mut self,
+		nibble_key: NibbleSlice,
+		full_key: &[u8],
+	) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
+		let mut partial = nibble_key;
+		let mut hash = self.hash;
+		let mut key_nibbles = 0;
+
+		// this loop iterates through non-inline nodes.
+		for depth in 0.. {
+			let node_data = match self.db.get(&hash, nibble_key.mid(key_nibbles).left()) {
+				Some(value) => value,
+				None =>
+					return Err(Box::new(match depth {
+						0 => TrieError::InvalidStateRoot(hash),
+						_ => TrieError::IncompleteDatabase(hash),
+					})),
+			};
+
+			self.record(|| TrieAccess::EncodedNode {
+				hash,
+				encoded_node: node_data.as_slice().into(),
+			});
+
+			// this loop iterates through all inline children (usually max 1)
+			// without incrementing the depth.
+			let mut node_data = &node_data[..];
+			loop {
+				let decoded = match L::Codec::decode(node_data) {
+					Ok(node) => node,
+					Err(e) => return Err(Box::new(TrieError::DecoderError(hash, e))),
+				};
+
+				let next_node = match decoded {
+					Node::Leaf(slice, _) => {
+						if !partial.starts_with(&slice) {
+							self.record(|| TrieAccess::NonExisting { full_key });
+							return Ok(None)
+						}
+
+						if partial.len() != slice.len() {
+							self.record(|| TrieAccess::NonExisting { full_key });
+						}
+
+						return Ok(Some(hash))
+					},
+					Node::Extension(slice, item) =>
+						if partial.starts_with(&slice) {
+							partial = partial.mid(slice.len());
+							key_nibbles += slice.len();
+							item
+						} else {
+							self.record(|| TrieAccess::NonExisting { full_key });
+
+							return Ok(None)
+						},
+					Node::Branch(children, value) =>
+						if partial.is_empty() {
+							if value.is_none() {
+								self.record(|| TrieAccess::NonExisting { full_key });
+							}
+
+							return Ok(Some(hash))
+						} else {
+							match children[partial.at(0) as usize] {
+								Some(x) => {
+									partial = partial.mid(1);
+									key_nibbles += 1;
+									x
+								},
+								None => {
+									self.record(|| TrieAccess::NonExisting { full_key });
+
+									return Ok(Some(hash))
+								},
+							}
+						},
+					Node::NibbledBranch(slice, children, value) => {
+						if !partial.starts_with(&slice) {
+							self.record(|| TrieAccess::NonExisting { full_key });
+
+							return Ok(None)
+						}
+
+						if partial.len() == slice.len() {
+							if value.is_none() {
+								self.record(|| TrieAccess::NonExisting { full_key });
+							}
+
+							return Ok(Some(hash))
+						} else {
+							match children[partial.at(slice.len()) as usize] {
+								Some(x) => {
+									partial = partial.mid(slice.len() + 1);
+									key_nibbles += slice.len() + 1;
+									x
+								},
+								None => {
+									self.record(|| TrieAccess::NonExisting { full_key });
+
+									return Ok(Some(hash))
+								},
+							}
+						}
+					},
+					Node::Empty => {
+						self.record(|| TrieAccess::NonExisting { full_key });
+
+						return Ok(Some(hash))
+					},
+				};
+
+				// check if new node data is inline or hash.
+				match next_node {
+					NodeHandle::Hash(data) => {
+						hash = decode_hash::<L::Hash>(data)
+							.ok_or_else(|| Box::new(TrieError::InvalidHash(hash, data.to_vec())))?;
+						break
+					},
+					NodeHandle::Inline(data) => {
+						node_data = data;
+					},
+				}
+			}
+		}
+		Ok(None)
+	}
 }

trie-db/src/sectriedb.rs

@@ -75,6 +75,13 @@ where
 		self.raw.get_with(L::Hash::hash(key).as_ref(), query)
 	}
 
+	fn get_closest_merkle_value(
+		&self,
+		_key: &[u8],
+	) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
+		Ok(None)
+	}
+
 	fn iter<'a>(
 		&'a self,
 	) -> Result<

trie-db/src/triedb.rs

@@ -252,6 +252,23 @@ where
 		.look_up(key, NibbleSlice::new(key))
 	}
 
+	fn get_closest_merkle_value(
+		&self,
+		key: &[u8],
+	) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
+		let mut cache = self.cache.as_ref().map(|c| c.borrow_mut());
+		let mut recorder = self.recorder.as_ref().map(|r| r.borrow_mut());
+
+		Lookup::<L, _> {
+			db: self.db,
+			query: |_: &[u8]| (),
+			hash: *self.root,
+			cache: cache.as_mut().map(|c| &mut ***c as &mut dyn TrieCache<L::Codec>),
+			recorder: recorder.as_mut().map(|r| &mut ***r as &mut dyn TrieRecorder<TrieHash<L>>),
+		}
+		.look_up_merkle(key, NibbleSlice::new(key))
+	}
+
 	fn iter<'a>(
 		&'a self,
 	) -> Result<

trie-db/test/src/triedb.rs

@@ -644,6 +644,58 @@ fn test_recorder_with_cache_get_hash_internal<T: TrieLayout>() {
 	}
 }
 
+test_layouts!(test_merkle_value, test_merkle_value_internal);
+fn test_merkle_value_internal<T: TrieLayout>() {
+	let mut memdb = MemoryDB::<T::Hash, HashKey<_>, DBValue>::default();
+	let mut root = Default::default();
+
+	// Data set.
+	let key_value = vec![
+		(b"A".to_vec(), vec![1; 64]),
+		(b"AA".to_vec(), vec![2; 64]),
+		(b"AAA".to_vec(), vec![3; 64]),
+		(b"AAB".to_vec(), vec![4; 64]),
+		(b"AB".to_vec(), vec![5; 4]),
+		(b"B".to_vec(), vec![6; 64]),
+	];
+	{
+		let mut t = TrieDBMutBuilder::<T>::new(&mut memdb, &mut root).build();
+		for (key, value) in &key_value {
+			t.insert(key, value).unwrap();
+		}
+	}
+
+	// Ensure we can fetch the merkle values for all present keys.
+	let trie = TrieDBBuilder::<T>::new(&memdb, &root).build();
+	for (key, _) in &key_value {
+		trie.get_closest_merkle_value(key).unwrap().unwrap();
+	}
+
+	// Key is not present in the trie, but the closest descendant is AAA.
+	let hash = trie.get_closest_merkle_value(b"AAAA").unwrap().unwrap();
+	let expected = trie.get_closest_merkle_value(b"AAA").unwrap().unwrap();
+	assert_eq!(hash, expected);
+
+	// Key is not present in the trie, but the closest descendant is AB.
+	let hash = trie.get_closest_merkle_value(b"ABA").unwrap().unwrap();
+	let expected = trie.get_closest_merkle_value(b"AB").unwrap().unwrap();
+	assert_eq!(hash, expected);
+
+	// Key is not present in the tire.
+	// At the nibble level (4 bits), b"A" is represented by number 65 in binary: 0100 0001.
+	// The key 0000 01000 shares no common prefix in the trie.
+	let not_present = trie.get_closest_merkle_value(&[0b_0000_0100]).unwrap();
+	assert!(not_present.is_none());
+
+	// b"D" should exist since it shares a prefix with b"A" and it returns the branch node
+	// when the layout is `ExtensionLayout` or a nibble branch otherwise.
+	let d_key = trie.get_closest_merkle_value(&[0b_0100_0100]).unwrap().unwrap();
+	// The same is true for b"C".
+	let c_key = trie.get_closest_merkle_value(&[0b_0100_0011]).unwrap().unwrap();
+	// Check the same branch is returned.
+	assert_eq!(d_key, c_key);
+}
+
 test_layouts!(iterator_seek_with_recorder, iterator_seek_with_recorder_internal);
 fn iterator_seek_with_recorder_internal<T: TrieLayout>() {
 	let d = vec![b"A".to_vec(), b"AA".to_vec(), b"AB".to_vec(), b"B".to_vec()];