From 1d7cd427f052f2dda94e3d2871f8f1e2532cbf11 Mon Sep 17 00:00:00 2001 From: Kris Nuttycombe Date: Mon, 17 Feb 2025 13:44:22 -0700 Subject: [PATCH] shardtree: make construction of fully-empty Parent nodes an error Ideally the errors here would all be panics, as construction of such a node represents a programming error; however, it is preferable to return extra contextual information about the circumstance that led to this error where possible, so we model this as an error where possible without altering the public API. --- shardtree/src/batch.rs | 47 +++---- shardtree/src/legacy.rs | 80 +++++++----- shardtree/src/lib.rs | 46 +++---- shardtree/src/prunable.rs | 252 +++++++++++++++++++++++--------------- 4 files changed, 254 insertions(+), 171 deletions(-) diff --git a/shardtree/src/batch.rs b/shardtree/src/batch.rs index 59c8714..e5eca9e 100644 --- a/shardtree/src/batch.rs +++ b/shardtree/src/batch.rs @@ -8,7 +8,7 @@ use tracing::trace; use crate::{ error::{InsertionError, ShardTreeError}, store::{Checkpoint, ShardStore}, - IncompleteAt, LocatedPrunableTree, LocatedTree, RetentionFlags, ShardTree, Tree, + IncompleteAt, LocatedPrunableTree, LocatedTree, PrunableTree, RetentionFlags, ShardTree, Tree, }; impl< @@ -215,7 +215,8 @@ impl LocatedPrunableTree { // fragments up the stack until we get the largest possible subtree while let Some((potential_sibling, marked)) = fragments.pop() { if potential_sibling.root_addr.parent() == subtree.root_addr.parent() { - subtree = unite(potential_sibling, subtree, prune_below); + subtree = unite(potential_sibling, subtree, prune_below) + .expect("subtree is non-Nil"); } else { // this is not a sibling node, so we push it back on to the stack // and are done @@ -238,7 +239,9 @@ impl LocatedPrunableTree { let minimal_tree_addr = Address::from(position_range.start).common_ancestor(&last_position.into()); trace!("Building minimal tree at {:?}", minimal_tree_addr); - build_minimal_tree(fragments, minimal_tree_addr, prune_below).map( + build_minimal_tree(fragments, minimal_tree_addr, prune_below) + .expect("construction of minimal tree does not result in creation of invalid parent nodes") + .map( |(to_insert, contains_marked, incomplete)| BatchInsertionResult { subtree: to_insert, contains_marked, @@ -263,16 +266,18 @@ fn unite( lroot: LocatedPrunableTree, rroot: LocatedPrunableTree, prune_below: Level, -) -> LocatedPrunableTree { +) -> Result, Address> { assert_eq!(lroot.root_addr.parent(), rroot.root_addr.parent()); - LocatedTree { + Ok(LocatedTree { root_addr: lroot.root_addr.parent(), root: if lroot.root_addr.level() < prune_below { - Tree::unite(lroot.root_addr.level(), None, lroot.root, rroot.root) + PrunableTree::unite(lroot.root_addr.level(), None, lroot.root, rroot.root) + .map_err(|_| lroot.root_addr)? } else { - Tree::parent(None, lroot.root, rroot.root) + PrunableTree::parent_checked(None, lroot.root, rroot.root) + .map_err(|_| lroot.root_addr)? }, - } + }) } /// Combines the given subtree with an empty sibling node to obtain the next level @@ -286,7 +291,7 @@ fn combine_with_empty( incomplete: &mut Vec, contains_marked: bool, prune_below: Level, -) -> LocatedPrunableTree { +) -> Result, Address> { assert_eq!(expect_left_child, root.root_addr.is_left_child()); let sibling_addr = root.root_addr.sibling(); incomplete.push(IncompleteAt { @@ -313,27 +318,27 @@ fn build_minimal_tree( mut xs: Vec<(LocatedPrunableTree, bool)>, root_addr: Address, prune_below: Level, -) -> Option<(LocatedPrunableTree, bool, Vec)> { +) -> Result, bool, Vec)>, Address> { // First, consume the stack from the right, building up a single tree // until we can't combine any more. if let Some((mut cur, mut contains_marked)) = xs.pop() { let mut incomplete = vec![]; while let Some((top, top_marked)) = xs.pop() { while cur.root_addr.level() < top.root_addr.level() { - cur = combine_with_empty(cur, true, &mut incomplete, top_marked, prune_below); + cur = combine_with_empty(cur, true, &mut incomplete, top_marked, prune_below)?; } if cur.root_addr.level() == top.root_addr.level() { contains_marked = contains_marked || top_marked; if cur.root_addr.is_right_child() { // We have a left child and a right child, so unite them. - cur = unite(top, cur, prune_below); + cur = unite(top, cur, prune_below)?; } else { // This is a left child, so we build it up one more level and then // we've merged as much as we can from the right and need to work from // the left xs.push((top, top_marked)); - cur = combine_with_empty(cur, true, &mut incomplete, top_marked, prune_below); + cur = combine_with_empty(cur, true, &mut incomplete, top_marked, prune_below)?; break; } } else { @@ -346,7 +351,7 @@ fn build_minimal_tree( // Ensure we can work from the left in a single pass by making this right-most subtree while cur.root_addr.level() + 1 < root_addr.level() { - cur = combine_with_empty(cur, true, &mut incomplete, contains_marked, prune_below); + cur = combine_with_empty(cur, true, &mut incomplete, contains_marked, prune_below)?; } // push our accumulated max-height right hand node back on to the stack. @@ -354,7 +359,7 @@ fn build_minimal_tree( // From the stack of subtrees, construct a single sparse tree that can be // inserted/merged into the existing tree - let res_tree = xs.into_iter().fold( + let res_tree = xs.into_iter().try_fold( None, |acc: Option>, (next_tree, next_marked)| { if let Some(mut prev_tree) = acc { @@ -368,19 +373,19 @@ fn build_minimal_tree( &mut incomplete, next_marked, prune_below, - ); + )?; } - Some(unite(prev_tree, next_tree, prune_below)) + Ok::<_, Address>(Some(unite(prev_tree, next_tree, prune_below)?)) } else { - Some(next_tree) + Ok::<_, Address>(Some(next_tree)) } }, - ); + )?; - res_tree.map(|t| (t, contains_marked, incomplete)) + Ok(res_tree.map(|t| (t, contains_marked, incomplete))) } else { - None + Ok(None) } } diff --git a/shardtree/src/legacy.rs b/shardtree/src/legacy.rs index 2abf1be..7e81bda 100644 --- a/shardtree/src/legacy.rs +++ b/shardtree/src/legacy.rs @@ -101,7 +101,7 @@ impl LocatedPrunableTree { leaf_addr: Address, mut filled: impl Iterator, split_at: Level, - ) -> (Self, Option) { + ) -> Result<(Self, Option), Address> { // add filled nodes to the subtree; here, we do not need to worry about // whether or not these nodes can be invalidated by a rewind let mut addr = leaf_addr; @@ -113,17 +113,19 @@ impl LocatedPrunableTree { if let Some(right) = filled.next() { // once we have a right-hand node, add a parent with the current tree // as the left-hand sibling - subtree = Tree::parent( + subtree = PrunableTree::parent_checked( None, subtree, Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)), - ); + ) + .map_err(|_| addr)?; } else { break; } } else { // the current address is for a right child, so add an empty left sibling - subtree = Tree::parent(None, Tree::empty(), subtree); + subtree = + PrunableTree::parent_checked(None, Tree::empty(), subtree).map_err(|_| addr)?; } addr = addr.parent(); @@ -140,16 +142,22 @@ impl LocatedPrunableTree { for right in filled { // build up the right-biased tree until we get a left-hand node while addr.is_right_child() { - supertree = supertree.map(|t| Tree::parent(None, Tree::empty(), t)); + supertree = supertree + .map(|t| PrunableTree::parent_checked(None, Tree::empty(), t)) + .transpose() + .map_err(|_| addr)?; addr = addr.parent(); } // once we have a left-hand root, add a parent with the current ommer as the right-hand sibling - supertree = Some(Tree::parent( - None, - supertree.unwrap_or_else(PrunableTree::empty), - Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)), - )); + supertree = Some( + PrunableTree::parent_checked( + None, + supertree.unwrap_or_else(PrunableTree::empty), + Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)), + ) + .map_err(|_| addr)?, + ); addr = addr.parent(); } @@ -161,7 +169,7 @@ impl LocatedPrunableTree { None }; - (subtree, supertree) + Ok((subtree, supertree)) } /// Insert the nodes belonging to the given incremental witness to this tree, truncating the @@ -196,23 +204,30 @@ impl LocatedPrunableTree { Address::from(witness.witnessed_position()), witness.filled().iter().cloned(), self.root_addr.level(), - ); + ) + .map_err(InsertionError::InputMalformed)?; // construct subtrees from the `cursor` part of the witness - let cursor_trees = witness.cursor().as_ref().filter(|c| c.size() > 0).map(|c| { - Self::from_frontier_parts( - witness.tip_position(), - c.leaf() - .cloned() - .expect("Cannot have an empty leaf for a non-empty tree"), - c.ommers_iter().cloned(), - &Retention::Checkpoint { - id: checkpoint_id, - marking: Marking::None, - }, - self.root_addr.level(), - ) - }); + let cursor_trees = witness + .cursor() + .as_ref() + .filter(|c| c.size() > 0) + .map(|c| { + Self::from_frontier_parts( + witness.tip_position(), + c.leaf() + .cloned() + .expect("Cannot have an empty leaf for a non-empty tree"), + c.ommers_iter().cloned(), + &Retention::Checkpoint { + id: checkpoint_id, + marking: Marking::None, + }, + self.root_addr.level(), + ) + }) + .transpose() + .map_err(InsertionError::InputMalformed)?; let (subtree, _) = past_subtree.insert_subtree(future_subtree, true)?; @@ -253,7 +268,7 @@ mod tests { frontier::CommitmentTree, witness::IncrementalWitness, Address, Level, Position, }; - use crate::{LocatedPrunableTree, RetentionFlags, Tree}; + use crate::{LocatedPrunableTree, PrunableTree, RetentionFlags, Tree}; #[test] fn insert_witness_nodes() { @@ -280,19 +295,22 @@ mod tests { assert_eq!( c.root, - Tree::parent( + PrunableTree::parent_checked( None, - Tree::parent( + PrunableTree::parent_checked( None, Tree::empty(), Tree::leaf(("ijklmnop".to_string(), RetentionFlags::EPHEMERAL)), - ), - Tree::parent( + ) + .unwrap(), + PrunableTree::parent_checked( None, Tree::leaf(("qrstuvwx".to_string(), RetentionFlags::EPHEMERAL)), Tree::empty() ) + .unwrap() ) + .unwrap() ); assert_eq!( diff --git a/shardtree/src/lib.rs b/shardtree/src/lib.rs index a7dbd83..9bdaa6b 100644 --- a/shardtree/src/lib.rs +++ b/shardtree/src/lib.rs @@ -414,44 +414,44 @@ impl< fn go( root_addr: Address, root: &PrunableTree, - ) -> Option<(PrunableTree, Position)> { + ) -> Result, Position)>, Address> { match &root.0 { Node::Parent { ann, left, right } => { 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(r_addr, right)?.map_or_else( || { - go(l_addr, left).map(|(new_left, pos)| { - ( - Tree::unite( + go(l_addr, left)? + .map(|(new_left, pos)| { + PrunableTree::unite( l_addr.level(), ann.clone(), new_left, Tree::empty(), - ), - pos, - ) - }) + ) + .map_err(|_| l_addr) + .map(|t| (t, pos)) + }) + .transpose() }, |(new_right, pos)| { - Some(( - Tree::unite( - l_addr.level(), - ann.clone(), - left.as_ref().clone(), - new_right, - ), - pos, - )) + Tree::unite( + l_addr.level(), + ann.clone(), + left.as_ref().clone(), + new_right, + ) + .map_err(|_| l_addr) + .map(|t| Some((t, pos))) }, ) } - Node::Leaf { value: (h, r) } => Some(( + Node::Leaf { value: (h, r) } => Ok(Some(( Tree::leaf((h.clone(), *r | RetentionFlags::CHECKPOINT)), root_addr.max_position(), - )), - Node::Nil => None, + ))), + Node::Nil => Ok(None), } } @@ -469,7 +469,9 @@ impl< // Update the rightmost subtree to add the `CHECKPOINT` flag to the right-most leaf (which // need not be a level-0 leaf; it's fine to rewind to a pruned state). if let Some(subtree) = self.store.last_shard().map_err(ShardTreeError::Storage)? { - if let Some((replacement, pos)) = go(subtree.root_addr, &subtree.root) { + if let Some((replacement, pos)) = go(subtree.root_addr, &subtree.root) + .map_err(|addr| ShardTreeError::Insert(InsertionError::InputMalformed(addr)))? + { self.store .put_shard(LocatedTree { root_addr: subtree.root_addr, diff --git a/shardtree/src/prunable.rs b/shardtree/src/prunable.rs index 5eb85c4..85eb766 100644 --- a/shardtree/src/prunable.rs +++ b/shardtree/src/prunable.rs @@ -258,22 +258,30 @@ impl PrunableTree { /// `level` must be the level of the root of the node being pruned. pub fn prune(self, level: Level) -> Self { match self { - Tree(Node::Parent { ann, left, right }) => Tree::unite( + Tree(Node::Parent { ann, left, right }) => PrunableTree::unite( level, ann, left.as_ref().clone().prune(level - 1), right.as_ref().clone().prune(level - 1), - ), + ) + .expect("pruning does not construct empty parent nodes"), other => other, } } + pub(crate) fn parent_checked(ann: Option>, left: Self, right: Self) -> Result { + if left.is_empty() && right.is_empty() && ann.is_none() { + Err(()) + } else { + Ok(Tree::parent(ann, left, right)) + } + } + /// Merge two subtrees having the same root address. /// /// The merge operation is checked to be strictly additive and returns an error if merging /// 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. - #[tracing::instrument()] pub fn merge_checked(self, root_addr: Address, other: Self) -> Result { /// Pre-condition: `root_addr` must be the address of `t0` and `t1`. #[allow(clippy::type_complexity)] @@ -330,12 +338,13 @@ impl PrunableTree { { let (l_addr, r_addr) = addr.children().ok_or(MergeError::TreeMalformed(addr))?; - Ok(Tree::unite( + PrunableTree::unite( addr.level() - 1, lann.or(rann), go(l_addr, ll.as_ref().clone(), rl.as_ref().clone())?, go(r_addr, lr.as_ref().clone(), rr.as_ref().clone())?, - )) + ) + .map_err(|_| MergeError::TreeMalformed(addr)) } else { unreachable!() } @@ -355,9 +364,14 @@ impl PrunableTree { /// replacement `Nil` value). /// /// `level` must be the level of the two nodes that are being joined. - pub(crate) fn unite(level: Level, ann: Option>, left: Self, right: Self) -> Self { + pub(crate) fn unite( + level: Level, + ann: Option>, + left: Self, + right: Self, + ) -> Result { match (left, right) { - (Tree(Node::Nil), Tree(Node::Nil)) if ann.is_none() => Tree::empty(), + (Tree(Node::Nil), Tree(Node::Nil)) if ann.is_none() => Ok(Tree::empty()), (Tree(Node::Leaf { value: lv }), Tree(Node::Leaf { value: rv })) // we can prune right-hand leaves that are not marked or reference leaves; if a // leaf is a checkpoint then that information will be propagated to the replacement @@ -365,9 +379,9 @@ impl PrunableTree { if lv.1 == RetentionFlags::EPHEMERAL && (rv.1 & (RetentionFlags::MARKED | RetentionFlags::REFERENCE)) == RetentionFlags::EPHEMERAL => { - Tree::leaf((H::combine(level, &lv.0, &rv.0), rv.1)) + Ok(Tree::leaf((H::combine(level, &lv.0, &rv.0), rv.1))) } - (left, right) => Tree::parent(ann, left, right), + (left, right) => PrunableTree::parent_checked(ann, left, right), } } } @@ -576,7 +590,7 @@ impl LocatedPrunableTree { position: Position, root_addr: Address, root: &PrunableTree, - ) -> Option> { + ) -> Result>, Address> { match &root.0 { Node::Parent { ann, left, right } => { let (l_child, r_child) = root_addr @@ -586,39 +600,84 @@ impl LocatedPrunableTree { // we are truncating within the range of the left node, so recurse // to the left to truncate the left child and then reconstruct the // node with `Nil` as the right sibling - go(position, l_child, left.as_ref()).map(|left| { - Tree::unite(l_child.level(), ann.clone(), left, Tree::empty()) - }) + go(position, l_child, left.as_ref())? + .map(|left| { + PrunableTree::unite( + l_child.level(), + ann.clone(), + left, + Tree::empty(), + ) + .map_err(|_| root_addr) + }) + .transpose() } else { // we are truncating within the range of the right node, so recurse // to the right to truncate the right child and then reconstruct the // node with the left sibling unchanged - go(position, r_child, right.as_ref()).map(|right| { - Tree::unite(r_child.level(), ann.clone(), left.as_ref().clone(), right) - }) + go(position, r_child, right.as_ref())? + .map(|right| { + PrunableTree::unite( + r_child.level(), + ann.clone(), + left.as_ref().clone(), + right, + ) + .map_err(|_| root_addr) + }) + .transpose() } } Node::Leaf { .. } => { if root_addr.max_position() <= position { - Some(root.clone()) + Ok(Some(root.clone())) } else { - None + Ok(None) } } - Node::Nil => None, + Node::Nil => Ok(None), } } if self.root_addr.position_range().contains(&position) { - go(position, self.root_addr, &self.root).map(|root| LocatedTree { - root_addr: self.root_addr, - root, - }) + go(position, self.root_addr, &self.root) + .expect("truncation does not introduce invalid subtree roots") + .map(|root| LocatedTree { + root_addr: self.root_addr, + root, + }) } else { None } } + // In the case that we are replacing a node entirely, we need to extend the + // subtree up to the level of the node being replaced, adding Nil siblings + // and recording the presence of those incomplete nodes when necessary + fn extend_to_level( + self, + level: Level, + ann: Option>, + ) -> Result<(PrunableTree, Vec
), Address> { + // construct the replacement node bottom-up + let mut node = self; + let mut incomplete = vec![]; + while node.root_addr.level() < level { + incomplete.push(node.root_addr.sibling()); + let full = node.root; + node = LocatedTree { + root_addr: node.root_addr.parent(), + root: if node.root_addr.is_right_child() { + PrunableTree::parent_checked(None, Tree::empty(), full) + } else { + PrunableTree::parent_checked(None, full, Tree::empty()) + } + .map_err(|_| node.root_addr.parent())?, + }; + } + Ok((node.root.reannotate_root(ann), incomplete)) + } + /// Inserts a descendant subtree into this subtree, creating empty sibling nodes as necessary /// to fill out the tree. /// @@ -644,38 +703,15 @@ impl LocatedPrunableTree { root_addr: Address, into: &PrunableTree, subtree: LocatedPrunableTree, - contains_marked: bool, - ) -> Result<(PrunableTree, Vec), InsertionError> { + ) -> Result<(PrunableTree, Vec
), InsertionError> { // An empty tree cannot replace any other type of tree. if subtree.root().is_nil() { Ok((into.clone(), vec![])) } else { - // In the case that we are replacing a node entirely, we need to extend the - // subtree up to the level of the node being replaced, adding Nil siblings - // and recording the presence of those incomplete nodes when necessary - let replacement = |ann: Option>, mut node: LocatedPrunableTree| { - // construct the replacement node bottom-up - let mut incomplete = vec![]; - while node.root_addr.level() < root_addr.level() { - incomplete.push(IncompleteAt { - address: node.root_addr.sibling(), - required_for_witness: contains_marked, - }); - let full = node.root; - node = LocatedTree { - root_addr: node.root_addr.parent(), - root: if node.root_addr.is_right_child() { - Tree::parent(None, Tree::empty(), full) - } else { - Tree::parent(None, full, Tree::empty()) - }, - }; - } - (node.root.reannotate_root(ann), incomplete) - }; - match into { - Tree(Node::Nil) => Ok(replacement(None, subtree)), + Tree(Node::Nil) => subtree + .extend_to_level(root_addr.level(), None) + .map_err(InsertionError::InputMalformed), Tree(Node::Leaf { value: (value, retention), }) => { @@ -730,7 +766,9 @@ impl LocatedPrunableTree { Err(InsertionError::Conflict(root_addr)) } } else { - Ok(replacement(Some(Arc::new(value.clone())), subtree)) + subtree + .extend_to_level(root_addr.level(), Some(Arc::new(value.clone()))) + .map_err(InsertionError::InputMalformed) } } parent if root_addr == subtree.root_addr => { @@ -749,29 +787,25 @@ impl LocatedPrunableTree { .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)?; - Ok(( - Tree::unite( - root_addr.level() - 1, - ann.clone(), - new_left, - right.as_ref().clone(), - ), - incomplete, - )) + let (new_left, incomplete) = go(l_addr, left.as_ref(), subtree)?; + PrunableTree::unite( + root_addr.level() - 1, + ann.clone(), + new_left, + right.as_ref().clone(), + ) + .map_err(|_| InsertionError::InputMalformed(root_addr)) + .map(|t| (t, incomplete)) } else { - let (new_right, incomplete) = - go(r_addr, right.as_ref(), subtree, contains_marked)?; - Ok(( - Tree::unite( - root_addr.level() - 1, - ann.clone(), - left.as_ref().clone(), - new_right, - ), - incomplete, - )) + let (new_right, incomplete) = go(r_addr, right.as_ref(), subtree)?; + PrunableTree::unite( + root_addr.level() - 1, + ann.clone(), + left.as_ref().clone(), + new_right, + ) + .map_err(|_| InsertionError::InputMalformed(root_addr)) + .map(|t| (t, incomplete)) } } } @@ -787,13 +821,22 @@ impl LocatedPrunableTree { ); let LocatedTree { root_addr, root } = self; if root_addr.contains(&subtree.root_addr) { - go(*root_addr, root, subtree, contains_marked).map(|(root, incomplete)| { + go(*root_addr, root, subtree).map(|(root, incomplete)| { let new_tree = LocatedTree { root_addr: *root_addr, root, }; assert!(new_tree.max_position() >= max_position); - (new_tree, incomplete) + ( + new_tree, + incomplete + .into_iter() + .map(|address| IncompleteAt { + address, + required_for_witness: contains_marked, + }) + .collect(), + ) }) } else { Err(InsertionError::NotContained(subtree.root_addr)) @@ -836,7 +879,7 @@ impl LocatedPrunableTree { frontier: NonEmptyFrontier, leaf_retention: &Retention, split_at: Level, - ) -> (Self, Option) { + ) -> Result<(Self, Option), Address> { let (position, leaf, ommers) = frontier.into_parts(); Self::from_frontier_parts(position, leaf, ommers.into_iter(), leaf_retention, split_at) } @@ -850,7 +893,7 @@ impl LocatedPrunableTree { mut ommers: impl Iterator, leaf_retention: &Retention, split_at: Level, - ) -> (Self, Option) { + ) -> Result<(Self, Option), Address> { let mut addr = Address::from(position); let mut subtree = Tree::leaf((leaf, leaf_retention.into())); @@ -858,12 +901,17 @@ impl LocatedPrunableTree { if addr.is_left_child() { // the current address is a left child, so create a parent with // an empty right-hand tree - subtree = Tree::parent(None, subtree, Tree::empty()); + subtree = + PrunableTree::parent_checked(None, subtree, Tree::empty()).map_err(|_| addr)?; } else if let Some(left) = ommers.next() { // the current address corresponds to a right child, so create a parent that // takes the left sibling to that child from the ommers - subtree = - Tree::parent(None, Tree::leaf((left, RetentionFlags::EPHEMERAL)), subtree); + subtree = PrunableTree::parent_checked( + None, + Tree::leaf((left, RetentionFlags::EPHEMERAL)), + subtree, + ) + .map_err(|_| addr)?; } else { break; } @@ -882,17 +930,24 @@ impl LocatedPrunableTree { for left in ommers { // build up the left-biased tree until we get a right-hand node while addr.is_left_child() { - supertree = supertree.map(|t| Tree::parent(None, t, Tree::empty())); + supertree = supertree + .map(|t| { + PrunableTree::parent_checked(None, t, Tree::empty()).map_err(|_| addr) + }) + .transpose()?; addr = addr.parent(); } // once we have a right-hand root, add a parent with the current ommer as the // left-hand sibling - supertree = Some(Tree::parent( - None, - Tree::leaf((left, RetentionFlags::EPHEMERAL)), - supertree.unwrap_or_else(Tree::empty), - )); + supertree = Some( + PrunableTree::parent_checked( + None, + Tree::leaf((left, RetentionFlags::EPHEMERAL)), + supertree.unwrap_or_else(Tree::empty), + ) + .map_err(|_| addr)?, + ); addr = addr.parent(); } @@ -906,7 +961,7 @@ impl LocatedPrunableTree { None }; - (located_subtree, located_supertree) + Ok((located_subtree, located_supertree)) } /// Inserts leaves and subtree roots from the provided frontier into this tree, up to the level @@ -929,7 +984,8 @@ impl LocatedPrunableTree { if subtree_range.contains(&frontier.position()) { let leaf_is_marked = leaf_retention.is_marked(); let (subtree, supertree) = - Self::from_frontier(frontier, leaf_retention, self.root_addr.level()); + Self::from_frontier(frontier, leaf_retention, self.root_addr.level()) + .map_err(InsertionError::InputMalformed)?; let subtree = self.insert_subtree(subtree, leaf_is_marked)?.0; @@ -950,10 +1006,10 @@ impl LocatedPrunableTree { to_clear: &[(Position, RetentionFlags)], root_addr: Address, root: &PrunableTree, - ) -> PrunableTree { + ) -> Result, Address> { if to_clear.is_empty() { // nothing to do, so we just return the root - root.clone() + Ok(root.clone()) } else { match &root.0 { Node::Parent { ann, left, right } => { @@ -963,17 +1019,18 @@ impl LocatedPrunableTree { let p = to_clear.partition_point(|(p, _)| p < &l_addr.position_range_end()); trace!( - "Tree::unite at {:?}, partitioned: {:?} {:?}", + "PrunableTree::unite at {:?}, partitioned: {:?} {:?}", root_addr, &to_clear[0..p], &to_clear[p..], ); - Tree::unite( + PrunableTree::unite( l_addr.level(), ann.clone(), - go(&to_clear[0..p], l_addr, left), - go(&to_clear[p..], r_addr, right), + go(&to_clear[0..p], l_addr, left)?, + go(&to_clear[p..], r_addr, right)?, ) + .map_err(|_| root_addr) } Node::Leaf { value: (h, r) } => { trace!("In {:?}, clearing {:?}", root_addr, to_clear); @@ -983,13 +1040,13 @@ impl LocatedPrunableTree { // a partially-pruned branch, and if it's a marked node then it will // be a level-0 leaf. match to_clear { - [(_, flags)] => Tree::leaf((h.clone(), *r & !*flags)), + [(_, flags)] => Ok(Tree::leaf((h.clone(), *r & !*flags))), _ => { panic!("Tree state inconsistent with checkpoints."); } } } - Node::Nil => Tree::empty(), + Node::Nil => Ok(Tree::empty()), } } } @@ -997,7 +1054,8 @@ impl LocatedPrunableTree { let to_clear = to_clear.into_iter().collect::>(); Self { root_addr: self.root_addr, - root: go(&to_clear, self.root_addr, &self.root), + root: go(&to_clear, self.root_addr, &self.root) + .expect("clearing flags cannot result in invalid tree state"), } }