feat: implement additional leaf traversal methods on MerkleStore

This commit is contained in:
Bobbin Threadbare
2023-08-16 02:54:11 -07:00
parent 85034af1df
commit 9f54c82d62
5 changed files with 163 additions and 30 deletions

View File

@@ -173,27 +173,24 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
// the path is computed from root to leaf, so it must be reversed
path.reverse();
Ok(ValuePath {
value: hash,
path: MerklePath::new(path),
})
Ok(ValuePath::new(hash, path))
}
/// Reconstructs a path from the root until a leaf or empty node and returns its depth.
// LEAF TRAVERSAL
// --------------------------------------------------------------------------------------------
/// Returns the depth of the first leaf or an empty node encountered while traversing the tree
/// from the specified root down according to the provided index.
///
/// The `tree_depth` parameter defines up to which depth the tree will be traversed, starting
/// from `root`. The maximum value the argument accepts is [u64::BITS].
///
/// The traversed path from leaf to root will start at the least significant bit of `index`,
/// and will be executed for `tree_depth` bits.
/// The `tree_depth` parameter specifies the depth of the tree rooted at `root`. The
/// maximum value the argument accepts is [u64::BITS].
///
/// # Errors
/// Will return an error if:
/// - The provided root is not found.
/// - The path from the root continues to a depth greater than `tree_depth`.
/// - The provided `tree_depth` is greater than `64.
/// - The provided `index` is not valid for a depth equivalent to `tree_depth`. For more
/// information, check [NodeIndex::new].
/// - The provided `tree_depth` is greater than 64.
/// - The provided `index` is not valid for a depth equivalent to `tree_depth`.
/// - No leaf or an empty node was found while traversing the tree down to `tree_depth`.
pub fn get_leaf_depth(
&self,
root: RpoDigest,
@@ -206,13 +203,6 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
}
NodeIndex::new(tree_depth, index)?;
// it's not illegal to have a maximum depth of `0`; we should just return the root in that
// case. this check will simplify the implementation as we could overflow bits for depth
// `0`.
if tree_depth == 0 {
return Ok(0);
}
// check if the root exists, providing the proper error report if it doesn't
let empty = EmptySubtreeRoots::empty_hashes(tree_depth);
let mut hash = root;
@@ -224,7 +214,7 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
let mut path = (index << (64 - tree_depth)).reverse_bits();
// iterate every depth and reconstruct the path from root to leaf
for depth in 0..tree_depth {
for depth in 0..=tree_depth {
// we short-circuit if an empty node has been found
if hash == empty[depth as usize] {
return Ok(depth);
@@ -241,13 +231,77 @@ impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
path >>= 1;
}
// at max depth assert it doesn't have sub-trees
if self.nodes.contains_key(&hash) {
return Err(MerkleError::DepthTooBig(tree_depth as u64 + 1));
// return an error because we exhausted the index but didn't find either a leaf or an
// empty node
Err(MerkleError::DepthTooBig(tree_depth as u64 + 1))
}
/// Returns index and value of a leaf node which is the only leaf node in a subtree defined by
/// the provided root. If the subtree contains zero or more than one leaf nodes None is
/// returned.
///
/// The `tree_depth` parameter specifies the depth of the parent tree such that `root` is
/// located in this tree at `root_index`. The maximum value the argument accepts is
/// [u64::BITS].
///
/// # Errors
/// Will return an error if:
/// - The provided root is not found.
/// - The provided `tree_depth` is greater than 64.
/// - The provided `root_index` has depth greater than `tree_depth`.
/// - A lone node at depth `tree_depth` is not a leaf node.
pub fn find_lone_leaf(
&self,
root: RpoDigest,
root_index: NodeIndex,
tree_depth: u8,
) -> Result<Option<(NodeIndex, RpoDigest)>, MerkleError> {
// we set max depth at u64::BITS as this is the largest meaningful value for a 64-bit index
const MAX_DEPTH: u8 = u64::BITS as u8;
if tree_depth > MAX_DEPTH {
return Err(MerkleError::DepthTooBig(tree_depth as u64));
}
let empty = EmptySubtreeRoots::empty_hashes(MAX_DEPTH);
let mut node = root;
if !self.nodes.contains_key(&node) {
return Err(MerkleError::RootNotInStore(node));
}
// depleted bits; return max depth
Ok(tree_depth)
let mut index = root_index;
if index.depth() > tree_depth {
return Err(MerkleError::DepthTooBig(index.depth() as u64));
}
// traverse down following the path of single non-empty nodes; this works because if a
// node has two empty children it cannot contain a lone leaf. similarly if a node has
// two non-empty children it must contain at least two leaves.
for depth in index.depth()..tree_depth {
// if the node is a leaf, return; otherwise, examine the node's children
let children = match self.nodes.get(&node) {
Some(node) => node,
None => return Ok(Some((index, node))),
};
let empty_node = empty[depth as usize + 1];
node = if children.left != empty_node && children.right == empty_node {
index = index.left_child();
children.left
} else if children.left == empty_node && children.right != empty_node {
index = index.right_child();
children.right
} else {
return Ok(None);
};
}
// if we are here, we got to `tree_depth`; thus, either the current node is a leaf node,
// and so we return it, or it is an internal node, and then we return an error
if self.nodes.contains_key(&node) {
Err(MerkleError::DepthTooBig(tree_depth as u64 + 1))
} else {
Ok(Some((index, node)))
}
}
// DATA EXTRACTORS