Skip to content

Commit

Permalink
Merge pull request #121 from zcash/shardtree-justify-unwraps-v0.3.x-b…
Browse files Browse the repository at this point in the history
…ackport

Backport #119 to shardtree v0.3.x
  • Loading branch information
nuttycom authored Dec 9, 2024
2 parents 7c862c1 + 766a00e commit c8ce5ef
Show file tree
Hide file tree
Showing 5 changed files with 86 additions and 26 deletions.
6 changes: 5 additions & 1 deletion shardtree/src/legacy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,11 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// construct the subtree and cap based on the frontier containing the
// witnessed position
let (past_subtree, past_supertree) = self.insert_frontier_nodes::<C>(
witness.tree().to_frontier().take().unwrap(),
witness
.tree()
.to_frontier()
.take()
.expect("IncrementalWitness must not be created from the empty tree."),
&Retention::Marked,
)?;

Expand Down
13 changes: 10 additions & 3 deletions shardtree/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -392,13 +392,16 @@ impl<

/// Adds a checkpoint at the rightmost leaf state of the tree.
pub fn checkpoint(&mut self, checkpoint_id: C) -> Result<bool, ShardTreeError<S::Error>> {
/// Pre-condition: `root_addr` must be the address of `root`.
fn go<H: Hashable + Clone + PartialEq>(
root_addr: Address,
root: &PrunableTree<H>,
) -> Option<(PrunableTree<H>, Position)> {
match root {
Tree(Node::Parent { ann, left, right }) => {
let (l_addr, r_addr) = root_addr.children().unwrap();
let (l_addr, r_addr) = root_addr
.children()
.expect("has children because we checked `root` is a parent");
go(r_addr, right).map_or_else(
|| {
go(l_addr, left).map(|(new_left, pos)| {
Expand Down Expand Up @@ -750,7 +753,10 @@ impl<
// Compute the roots of the left and right children and hash them together.
// We skip computation in any subtrees that will not have data included in
// the final result.
let (l_addr, r_addr) = cap.root_addr.children().unwrap();
let (l_addr, r_addr) = cap
.root_addr
.children()
.expect("has children because we checked `cap.root` is a parent");
let l_result = if r_addr.contains(&target_addr) {
None
} else {
Expand Down Expand Up @@ -1103,7 +1109,8 @@ impl<
cur_addr = cur_addr.parent();
}

Ok(MerklePath::from_parts(witness, position).unwrap())
Ok(MerklePath::from_parts(witness, position)
.expect("witness has length DEPTH because we extended it to the root"))
}

fn witness_internal(
Expand Down
51 changes: 39 additions & 12 deletions shardtree/src/prunable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,9 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
|| {
// Compute the roots of the left and right children and hash them
// together.
let (l_addr, r_addr) = root_addr.children().unwrap();
let (l_addr, r_addr) = root_addr
.children()
.expect("The root address of a parent node must have children.");
accumulate_result_with(
left.root_hash(l_addr, truncate_at),
right.root_hash(r_addr, truncate_at),
Expand Down Expand Up @@ -207,6 +209,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
/// would cause information loss or if a conflict between root hashes occurs at a node. The
/// returned error contains the address of the node where such a conflict occurred.
pub fn merge_checked(self, root_addr: Address, other: Self) -> Result<Self, Address> {
/// Pre-condition: `root_addr` must be the address of `t0` and `t1`.
#[allow(clippy::type_complexity)]
fn go<H: Hashable + Clone + PartialEq>(
addr: Address,
Expand Down Expand Up @@ -261,7 +264,9 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
}),
) = (lparent, rparent)
{
let (l_addr, r_addr) = addr.children().unwrap();
let (l_addr, r_addr) = addr
.children()
.expect("The root address of a parent node must have children.");
Ok(Tree::unite(
addr.level() - 1,
lann.or(rann),
Expand Down Expand Up @@ -357,14 +362,17 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {

/// Returns the positions of marked leaves in the tree.
pub fn marked_positions(&self) -> BTreeSet<Position> {
/// Pre-condition: `root_addr` must be the address of `root`.
fn go<H: Hashable + Clone + PartialEq>(
root_addr: Address,
root: &PrunableTree<H>,
acc: &mut BTreeSet<Position>,
) {
match &root.0 {
Node::Parent { left, right, .. } => {
let (l_addr, r_addr) = root_addr.children().unwrap();
let (l_addr, r_addr) = root_addr
.children()
.expect("has children because we checked `root` is a parent");
go(l_addr, left.as_ref(), acc);
go(r_addr, right.as_ref(), acc);
}
Expand All @@ -391,8 +399,10 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
/// Returns either the witness for the leaf at the specified position, or an error that
/// describes the causes of failure.
pub fn witness(&self, position: Position, truncate_at: Position) -> Result<Vec<H>, QueryError> {
// traverse down to the desired leaf position, and then construct
// the authentication path on the way back up.
/// Traverse down to the desired leaf position, and then construct
/// the authentication path on the way back up.
//
/// Pre-condition: `root_addr` must be the address of `root`.
fn go<H: Hashable + Clone + PartialEq>(
root: &PrunableTree<H>,
root_addr: Address,
Expand All @@ -401,7 +411,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
) -> Result<Vec<H>, Vec<Address>> {
match &root.0 {
Node::Parent { left, right, .. } => {
let (l_addr, r_addr) = root_addr.children().unwrap();
let (l_addr, r_addr) = root_addr
.children()
.expect("has children because we checked `root` is a parent");
if root_addr.level() > 1.into() {
let r_start = r_addr.position_range_start();
if position < r_start {
Expand Down Expand Up @@ -476,14 +488,17 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
/// subtree root with the specified position as its maximum position exists, or `None`
/// otherwise.
pub fn truncate_to_position(&self, position: Position) -> Option<Self> {
/// Pre-condition: `root_addr` must be the address of `root`.
fn go<H: Hashable + Clone + PartialEq>(
position: Position,
root_addr: Address,
root: &PrunableTree<H>,
) -> Option<PrunableTree<H>> {
match &root.0 {
Node::Parent { ann, left, right } => {
let (l_child, r_child) = root_addr.children().unwrap();
let (l_child, r_child) = root_addr
.children()
.expect("has children because we checked `root` is a parent");
if position < r_child.position_range_start() {
// we are truncating within the range of the left node, so recurse
// to the left to truncate the left child and then reconstruct the
Expand Down Expand Up @@ -537,8 +552,10 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
subtree: Self,
contains_marked: bool,
) -> Result<(Self, Vec<IncompleteAt>), InsertionError> {
// A function to recursively dig into the tree, creating a path downward and introducing
// empty nodes as necessary until we can insert the provided subtree.
/// A function to recursively dig into the tree, creating a path downward and introducing
/// empty nodes as necessary until we can insert the provided subtree.
///
/// Pre-condition: `root_addr` must be the address of `into`.
#[allow(clippy::type_complexity)]
fn go<H: Hashable + Clone + PartialEq>(
root_addr: Address,
Expand Down Expand Up @@ -621,7 +638,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
Tree(Node::Parent { ann, left, right }) => {
// In this case, we have an existing parent but we need to dig down farther
// before we can insert the subtree that we're carrying for insertion.
let (l_addr, r_addr) = root_addr.children().unwrap();
let (l_addr, r_addr) = root_addr
.children()
.expect("has children because we checked `into` is a parent");
if l_addr.contains(&subtree.root_addr) {
let (new_left, incomplete) =
go(l_addr, left.as_ref(), subtree, is_complete, contains_marked)?;
Expand Down Expand Up @@ -696,7 +715,12 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
if r.remainder.next().is_some() {
Err(InsertionError::TreeFull)
} else {
Ok((r.subtree, r.max_insert_position.unwrap(), checkpoint_id))
Ok((
r.subtree,
r.max_insert_position
.expect("Batch insertion result position is never initialized to None"),
checkpoint_id,
))
}
})
}
Expand Down Expand Up @@ -820,6 +844,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
/// Clears the specified retention flags at all positions specified, pruning any branches
/// that no longer need to be retained.
pub fn clear_flags(&self, to_clear: BTreeMap<Position, RetentionFlags>) -> Self {
/// Pre-condition: `root_addr` must be the address of `root`.
fn go<H: Hashable + Clone + PartialEq>(
to_clear: &[(Position, RetentionFlags)],
root_addr: Address,
Expand All @@ -831,7 +856,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
} else {
match root {
Tree(Node::Parent { ann, left, right }) => {
let (l_addr, r_addr) = root_addr.children().unwrap();
let (l_addr, r_addr) = root_addr
.children()
.expect("has children because we checked `root` is a parent");

let p = to_clear.partition_point(|(p, _)| p < &l_addr.position_range_end());
trace!(
Expand Down
27 changes: 20 additions & 7 deletions shardtree/src/store/caching.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,9 +46,11 @@ where
let _ = cache.put_cap(backend.get_cap()?);

backend.with_checkpoints(backend.checkpoint_count()?, |checkpoint_id, checkpoint| {
// TODO: Once MSRV is at least 1.82, replace this (and similar `expect()`s below) with:
// `let Ok(_) = cache.add_checkpoint(checkpoint_id.clone(), checkpoint.clone());`
cache
.add_checkpoint(checkpoint_id.clone(), checkpoint.clone())
.unwrap();
.expect("error type is Infallible");
Ok(())
})?;

Expand All @@ -74,26 +76,37 @@ where
}
self.deferred_actions.clear();

for shard_root in self.cache.get_shard_roots().unwrap() {
for shard_root in self
.cache
.get_shard_roots()
.expect("error type is Infallible")
{
self.backend.put_shard(
self.cache
.get_shard(shard_root)
.unwrap()
.expect("error type is Infallible")
.expect("known address"),
)?;
}
self.backend.put_cap(self.cache.get_cap().unwrap())?;
self.backend
.put_cap(self.cache.get_cap().expect("error type is Infallible"))?;

let mut checkpoints = Vec::with_capacity(self.cache.checkpoint_count().unwrap());
let mut checkpoints = Vec::with_capacity(
self.cache
.checkpoint_count()
.expect("error type is Infallible"),
);
self.cache
.with_checkpoints(
self.cache.checkpoint_count().unwrap(),
self.cache
.checkpoint_count()
.expect("error type is Infallible"),
|checkpoint_id, checkpoint| {
checkpoints.push((checkpoint_id.clone(), checkpoint.clone()));
Ok(())
},
)
.unwrap();
.expect("error type is Infallible");
for (checkpoint_id, checkpoint) in checkpoints {
self.backend.add_checkpoint(checkpoint_id, checkpoint)?;
}
Expand Down
15 changes: 12 additions & 3 deletions shardtree/src/tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -218,10 +218,13 @@ impl<A, V> LocatedTree<A, V> {

/// Returns the value at the specified position, if any.
pub fn value_at_position(&self, position: Position) -> Option<&V> {
/// Pre-condition: `addr` must be the address of `root`.
fn go<A, V>(pos: Position, addr: Address, root: &Tree<A, V>) -> Option<&V> {
match &root.0 {
Node::Parent { left, right, .. } => {
let (l_addr, r_addr) = addr.children().unwrap();
let (l_addr, r_addr) = addr
.children()
.expect("has children because we checked `root` is a parent");
if l_addr.position_range().contains(&pos) {
go(pos, l_addr, left)
} else {
Expand Down Expand Up @@ -265,6 +268,7 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
/// if the tree is terminated by a [`Node::Nil`] or leaf node before the specified address can
/// be reached.
pub fn subtree(&self, addr: Address) -> Option<Self> {
/// Pre-condition: `root_addr` must be the address of `root`.
fn go<A: Clone, V: Clone>(
root_addr: Address,
root: &Tree<A, V>,
Expand All @@ -278,7 +282,9 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
} else {
match &root.0 {
Node::Parent { left, right, .. } => {
let (l_addr, r_addr) = root_addr.children().unwrap();
let (l_addr, r_addr) = root_addr
.children()
.expect("has children because we checked `root` is a parent");
if l_addr.contains(&addr) {
go(l_addr, left.as_ref(), addr)
} else {
Expand All @@ -302,6 +308,7 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
/// If this root address of this tree is lower down in the tree than the level specified,
/// the entire tree is returned as the sole element of the result vector.
pub fn decompose_to_level(self, level: Level) -> Vec<Self> {
/// Pre-condition: `root_addr` must be the address of `root`.
fn go<A: Clone, V: Clone>(
level: Level,
root_addr: Address,
Expand All @@ -312,7 +319,9 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
} else {
match root.0 {
Node::Parent { left, right, .. } => {
let (l_addr, r_addr) = root_addr.children().unwrap();
let (l_addr, r_addr) = root_addr
.children()
.expect("has children because we checked `root` is a parent");
let mut l_decomposed = go(
level,
l_addr,
Expand Down

0 comments on commit c8ce5ef

Please sign in to comment.