trie-db: Fetch the closest merkle value

trie-db/src/fatdb.rs

@@ -18,7 +18,7 @@ use super::{
 };
 use hash_db::{HashDBRef, Hasher};
 
-use crate::{rstd::boxed::Box, TrieDBBuilder};
+use crate::{rstd::boxed::Box, MerkleValue, TrieDBBuilder};
 
 /// A `Trie` implementation which hashes keys and uses a generic `HashDB` backing database.
 /// Additionaly it stores inserted hash-key mappings for later retrieval.
@@ -72,6 +72,13 @@ where
 		self.raw.get_with(L::Hash::hash(key).as_ref(), query)
 	}
 
+	fn lookup_first_descendant(
+		&self,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<TrieHash<L>>>, TrieHash<L>, CError<L>> {
+		self.raw.lookup_first_descendant(key)
+	}
+
 	fn iter<'a>(
 		&'a self,
 	) -> Result<

trie-db/src/lib.rs

@@ -277,6 +277,17 @@ pub trait Trie<L: TrieLayout> {
 		query: Q,
 	) -> Result<Option<Q::Item>, TrieHash<L>, CError<L>>;
 
+	/// Look up the merkle value 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_descendant(
+		&self,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<TrieHash<L>>>, TrieHash<L>, CError<L>>;
+
 	/// Returns a depth-first iterator over the elements of trie.
 	fn iter<'a>(
 		&'a self,
@@ -404,6 +415,13 @@ impl<'db, 'cache, L: TrieLayout> Trie<L> for TrieKinds<'db, 'cache, L> {
 		wrapper!(self, get_with, key, query)
 	}
 
+	fn lookup_first_descendant(
+		&self,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<TrieHash<L>>>, TrieHash<L>, CError<L>> {
+		wrapper!(self, lookup_first_descendant, key)
+	}
+
 	fn iter<'a>(
 		&'a self,
 	) -> Result<
@@ -738,3 +756,15 @@ impl From<Bytes> for BytesWeak {
 		Self(rstd::sync::Arc::downgrade(&bytes.0))
 	}
 }
+
+/// A value returned by [`Trie::lookup_first_descendant`].
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum MerkleValue<H> {
+	/// The merkle value is the node data itself when the
+	/// node data is smaller than `MAX_INLINE_VALUE`.
+	///
+	/// Note: The case of inline nodes.
+	Node(Vec<u8>),
+	/// The merkle value is the hash of the node.
+	Hash(H),
+}

trie-db/src/lookup.rs

@@ -18,9 +18,9 @@ use crate::{
 	nibble::NibbleSlice,
 	node::{decode_hash, Node, NodeHandle, NodeHandleOwned, NodeOwned, Value, ValueOwned},
 	node_codec::NodeCodec,
-	rstd::boxed::Box,
-	Bytes, CError, CachedValue, DBValue, Query, RecordedForKey, Result, TrieAccess, TrieCache,
-	TrieError, TrieHash, TrieLayout, TrieRecorder,
+	rstd::{boxed::Box, vec::Vec},
+	Bytes, CError, CachedValue, DBValue, MerkleValue, Query, RecordedForKey, Result, TrieAccess,
+	TrieCache, TrieError, TrieHash, TrieLayout, TrieRecorder,
 };
 use hash_db::{HashDBRef, Hasher, Prefix};
 
@@ -133,6 +133,227 @@ where
 			recorder.record(get_access());
 		}
 	}
+	/// 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 (hash) 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.
+	pub fn lookup_first_descendant(
+		mut self,
+		full_key: &[u8],
+		nibble_key: NibbleSlice,
+	) -> Result<Option<MerkleValue<TrieHash<L>>>, TrieHash<L>, CError<L>> {
+		let mut partial = nibble_key;
+		let mut hash = self.hash;
+		let mut key_nibbles = 0;
+
+		let mut cache = self.cache.take();
+
+		// this loop iterates through non-inline nodes.
+		for depth in 0.. {
+			// Ensure the owned node reference lives long enough.
+			// Value is never read, but the reference is.
+			let mut _owned_node = NodeOwned::Empty;
+
+			// The binary encoded data of the node fetched from the database.
+			//
+			// Populated by `get_owned_node` to avoid one extra allocation by not
+			// calling `NodeOwned::to_encoded` when computing the hash of inlined nodes.
+			let mut node_data = Vec::new();
+
+			// Get the owned node representation from the database.
+			let mut get_owned_node = |depth: i32| {
+				let 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),
+						})),
+				};
+
+				let decoded = match L::Codec::decode(&data[..]) {
+					Ok(node) => node,
+					Err(e) => return Err(Box::new(TrieError::DecoderError(hash, e))),
+				};
+
+				let owned = decoded.to_owned_node::<L>()?;
+				node_data = data;
+				Ok(owned)
+			};
+
+			let mut node = if let Some(cache) = &mut cache {
+				let node = cache.get_or_insert_node(hash, &mut || get_owned_node(depth))?;
+
+				self.record(|| TrieAccess::NodeOwned { hash, node_owned: node });
+
+				node
+			} else {
+				_owned_node = get_owned_node(depth)?;
+
+				self.record(|| TrieAccess::EncodedNode {
+					hash,
+					encoded_node: node_data.as_slice().into(),
+				});
+
+				&_owned_node
+			};
+
+			// this loop iterates through all inline children (usually max 1)
+			// without incrementing the depth.
+			let mut is_inline = false;
+			loop {
+				let next_node = match node {
+					NodeOwned::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_slice(&partial) {
+							self.record(|| TrieAccess::NonExisting { full_key });
+							return Ok(None)
+						}
+
+						if partial.len() != slice.len() {
+							self.record(|| TrieAccess::NonExisting { full_key });
+						}
+
+						let res = is_inline
+							.then(|| MerkleValue::Node(node_data))
+							.unwrap_or_else(|| MerkleValue::Hash(hash));
+						return Ok(Some(res))
+					},
+					NodeOwned::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_slice(&partial) {
+								let res = is_inline
+									.then(|| MerkleValue::Node(node_data))
+									.unwrap_or_else(|| MerkleValue::Hash(hash));
+								Ok(Some(res))
+							} 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_vec(&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)
+						}
+					},
+					NodeOwned::Branch(children, value) =>
+						if partial.is_empty() {
+							if value.is_none() {
+								self.record(|| TrieAccess::NonExisting { full_key });
+							}
+							let res = is_inline
+								.then(|| MerkleValue::Node(node_data))
+								.unwrap_or_else(|| MerkleValue::Hash(hash));
+							return Ok(Some(res))
+						} 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)
+								},
+							}
+						},
+					NodeOwned::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_slice(&partial) {
+								let res = is_inline
+									.then(|| MerkleValue::Node(node_data))
+									.unwrap_or_else(|| MerkleValue::Hash(hash));
+								Ok(Some(res))
+							} 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_vec(&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 });
+							}
+
+							let res = is_inline
+								.then(|| MerkleValue::Node(node_data))
+								.unwrap_or_else(|| MerkleValue::Hash(hash));
+							return Ok(Some(res))
+						} 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)
+								},
+							}
+						}
+					},
+					NodeOwned::Empty => {
+						self.record(|| TrieAccess::NonExisting { full_key });
+
+						return Ok(None)
+					},
+					NodeOwned::Value(_, _) => {
+						unreachable!(
+							"`NodeOwned::Value` can not be reached by using the hash of a node. \
+							 `NodeOwned::Value` is only constructed when loading a value into memory, \
+							 which needs to have a different hash than any node; qed",
+						)
+					},
+				};
+
+				// check if new node data is inline or hash.
+				match next_node {
+					NodeHandleOwned::Hash(new_hash) => {
+						hash = *new_hash;
+						break
+					},
+					NodeHandleOwned::Inline(inline_node) => {
+						node = &inline_node;
+						is_inline = true;
+					},
+				}
+			}
+		}
+		Ok(None)
+	}
 
 	/// Look up the given `nibble_key`.
 	///

trie-db/src/nibble/nibblevec.rs

@@ -235,6 +235,25 @@ impl NibbleVec {
 		true
 	}
 
+	/// Same as [`Self::starts_with`] but using [`NibbleSlice`].
+	pub fn starts_with_slice(&self, other: &NibbleSlice) -> bool {
+		if self.len() < other.len() {
+			return false
+		}
+
+		match self.as_nibbleslice() {
+			Some(slice) => slice.starts_with(&other),
+			None => {
+				for i in 0..other.len() {
+					if self.at(i) != other.at(i) {
+						return false
+					}
+				}
+				true
+			},
+		}
+	}
+
 	/// Return an iterator over `Partial` bytes representation.
 	pub fn right_iter<'a>(&'a self) -> impl Iterator<Item = u8> + 'a {
 		let require_padding = self.len % nibble_ops::NIBBLE_PER_BYTE != 0;

trie-db/src/sectriedb.rs

@@ -13,8 +13,8 @@
 // limitations under the License.
 
 use crate::{
-	rstd::boxed::Box, triedb::TrieDB, CError, DBValue, Query, Result, Trie, TrieDBBuilder,
-	TrieHash, TrieItem, TrieIterator, TrieKeyItem, TrieLayout,
+	rstd::boxed::Box, triedb::TrieDB, CError, DBValue, MerkleValue, Query, Result, Trie,
+	TrieDBBuilder, TrieHash, TrieItem, TrieIterator, TrieKeyItem, TrieLayout,
 };
 use hash_db::{HashDBRef, Hasher};
 
@@ -75,6 +75,13 @@ where
 		self.raw.get_with(L::Hash::hash(key).as_ref(), query)
 	}
 
+	fn lookup_first_descendant(
+		&self,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<TrieHash<L>>>, TrieHash<L>, CError<L>> {
+		self.raw.lookup_first_descendant(key)
+	}
+
 	fn iter<'a>(
 		&'a self,
 	) -> Result<

trie-db/src/triedb.rs

@@ -18,8 +18,8 @@ use crate::{
 	nibble::NibbleSlice,
 	node::{decode_hash, NodeHandle, OwnedNode},
 	rstd::boxed::Box,
-	CError, DBValue, Query, Result, Trie, TrieAccess, TrieCache, TrieError, TrieHash, TrieItem,
-	TrieIterator, TrieKeyItem, TrieLayout, TrieRecorder,
+	CError, DBValue, MerkleValue, Query, Result, Trie, TrieAccess, TrieCache, TrieError, TrieHash,
+	TrieItem, TrieIterator, TrieKeyItem, TrieLayout, TrieRecorder,
 };
 #[cfg(feature = "std")]
 use crate::{
@@ -252,6 +252,23 @@ where
 		.look_up(key, NibbleSlice::new(key))
 	}
 
+	fn lookup_first_descendant(
+		&self,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<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>>),
+		}
+		.lookup_first_descendant(key, NibbleSlice::new(key))
+	}
+
 	fn iter<'a>(
 		&'a self,
 	) -> Result<