Skip to content

Commit

Permalink
shardtree: Ensure LocatedTree is correct by construction
Browse files Browse the repository at this point in the history 
This ensures that the various existing `unwrap`s are correctly
justified.
  • Loading branch information
@str4d
str4d committed Nov 23, 2024
1 parent 1e81f3c commit 166872e
Show file tree
Hide file tree
Showing 4 changed files with 84 additions and 19 deletions.
4 changes: 4 additions & 0 deletions shardtree/CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@ and this project adheres to Rust's notion of

## Unreleased

### Changed
- `shardtree::LocatedTree::from_parts` now returns `Option<Self>` (returning
`None` if the provided `Address` and `Tree` are inconsistent).

## [0.5.0] - 2024-10-04

This release includes a significant refactoring and rework of several methods
Expand Down
10 changes: 8 additions & 2 deletions shardtree/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -406,13 +406,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.0 {
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 @@ -765,7 +768,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
42 changes: 31 additions & 11 deletions shardtree/src/prunable.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -358,6 +358,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
/// Note that no actual leaf value may exist at this position, as it may have previously been
/// pruned.
pub fn max_position(&self) -> Option<Position> {
/// Pre-condition: `addr` must be the address of `root`.
fn go<H>(
addr: Address,
root: &Tree<Option<Arc<H>>, (H, RetentionFlags)>,
Expand All @@ -369,7 +370,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
if ann.is_some() {
Some(addr.max_position())
} else {
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");
go(r_addr, right.as_ref()).or_else(|| go(l_addr, left.as_ref()))
}
}
Expand Down Expand Up @@ -406,14 +409,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 @@ -440,8 +446,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 @@ -450,7 +458,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 @@ -525,14 +535,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 @@ -586,8 +599,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 @@ -694,7 +709,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, contains_marked)?;
Expand Down Expand Up @@ -892,6 +909,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 @@ -903,7 +921,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
} else {
match &root.0 {
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 Expand Up @@ -1228,7 +1248,7 @@ mod tests {
root in arb_prunable_tree(arb_char_str(), 8, 2^6)
) {
let root_addr = Address::from_parts(Level::from(7), 0);
let tree = LocatedTree::from_parts(root_addr, root);
let tree = LocatedTree::from_parts(root_addr, root).unwrap();

let (to_clear, to_retain) = tree.flag_positions().into_iter().enumerate().fold(
(BTreeMap::new(), BTreeMap::new()),
Expand Down
47 changes: 41 additions & 6 deletions shardtree/src/tree.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -197,9 +197,35 @@ pub struct LocatedTree<A, V> {
}

impl<A, V> LocatedTree<A, V> {
/// Constructs a new LocatedTree from its constituent parts
pub fn from_parts(root_addr: Address, root: Tree<A, V>) -> Self {
LocatedTree { root_addr, root }
/// Constructs a new LocatedTree from its constituent parts.
///
/// Returns `None` if `root_addr` is inconsistent with `root` (in particular, if the
/// level of `root_addr` is too small to contain `tree`).
pub fn from_parts(root_addr: Address, root: Tree<A, V>) -> Option<Self> {
// In order to meet various pre-conditions throughout the crate, we require that
// no `Node::Parent` in `root` has a level of 0 relative to `root_addr`.
fn is_consistent<A, V>(addr: Address, root: &Tree<A, V>) -> bool {
match (&root.0, addr.children()) {
// Found an inconsistency!
(Node::Parent { .. }, None) => false,
// Check consistency of children recursively.
(Node::Parent { left, right, .. }, Some((l_addr, r_addr))) => {
is_consistent(l_addr, left) && is_consistent(r_addr, right)
}

// Leaves are technically allowed to occur at any level, so we do not
// require `addr` to have no children.
(Node::Leaf { .. }, _) => true,

// Nil nodes have no information, so we cannot verify that the data it
// represents is consistent with `root_addr`. Instead we rely on methods
// that mutate `LocatedTree` to verify that the insertion address is not
// inconsistent with `root_addr`.
(Node::Nil, _) => true,
}
}

is_consistent(root_addr, &root).then_some(LocatedTree { root_addr, root })
}

/// Returns the root address of this tree.
Expand Down Expand Up @@ -234,10 +260,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 @@ -306,6 +335,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 @@ -319,7 +349,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 @@ -343,6 +375,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 @@ -353,7 +386,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 166872e

Please sign in to comment.