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,174 @@ 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, _) => {
+						// The leaf slice can be longer than remainder of the provided key
+						// (descendent), but not the other way around.
+						if !slice.starts_with(&partial) {
+							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.len() < slice.len() {
+							self.record(|| TrieAccess::NonExisting { full_key });
+
+							// Extension slice can be longer than remainder of the provided key
+							// (descendent), ensure the extension slice starts with the remainder
+							// of the provided key.
+							return if slice.starts_with(&partial) {
+								Ok(Some(hash))
+							} else {
+								Ok(None)
+							}
+						}
+
+						// Remainder of the provided key is longer than the extension slice,
+						// must advance the node iteration if and only if keys share
+						// a common prefix.
+						if partial.starts_with(&slice) {
+							// Empties the partial key if the extension slice is longer.
+							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(None)
+								},
+							}
+						},
+					Node::NibbledBranch(slice, children, value) => {
+						// Not enough remainder key to continue the search.
+						if partial.len() < slice.len() {
+							self.record(|| TrieAccess::NonExisting { full_key });
+
+							// Branch slice starts with the remainder key, there's nothing to
+							// advance.
+							return if slice.starts_with(&partial) {
+								Ok(Some(hash))
+							} else {
+								Ok(None)
+							}
+						}
+
+						// Partial key is longer or equal than the branch slice.
+						// Ensure partial key starts with the branch slice.
+						if !partial.starts_with(&slice) {
+							self.record(|| TrieAccess::NonExisting { full_key });
+							return Ok(None)
+						}
+
+						// Partial key starts with the branch slice.
+						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(None)
+								},
+							}
+						}
+					},
+					Node::Empty => {
+						self.record(|| TrieAccess::NonExisting { full_key });
+
+						return Ok(None)
+					},
+				};
+
+				// 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<