Skip to content

Commit

Permalink
shardtree: Justify unwraps due to upheld pre-conditions
Browse files Browse the repository at this point in the history
  • Loading branch information
str4d authored and nuttycom committed Dec 2, 2024
1 parent 150f554 commit 6298065
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 16 deletions.
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
42 changes: 32 additions & 10 deletions shardtree/src/prunable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -357,14 +357,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 +394,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 +406,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 +483,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 +547,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 +633,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 +710,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 +839,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 +851,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
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 6298065

Please sign in to comment.