zec
/
incrementalmerkletree
mirror of https://github.com/zcash/incrementalmerkletree.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge 422889cd5a into 561a6dc29b

This commit is contained in:
Kris Nuttycombe 2025-02-19 15:42:40 +00:00 committed by GitHub
parent 561a6dc29b 422889cd5a
commit 587fe8a12d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 259 additions and 171 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
shardtree/src/batch.rs
View File

@ -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<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// 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<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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<H: Hashable + Clone + PartialEq>(
lroot: LocatedPrunableTree<H>,
rroot: LocatedPrunableTree<H>,
prune_below: Level,
) -> LocatedPrunableTree<H> {
) -> Result<LocatedPrunableTree<H>, 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<H: Hashable + Clone + PartialEq>(
incomplete: &mut Vec<IncompleteAt>,
contains_marked: bool,
prune_below: Level,
) -> LocatedPrunableTree<H> {
) -> Result<LocatedPrunableTree<H>, Address> {
assert_eq!(expect_left_child, root.root_addr.is_left_child());
let sibling_addr = root.root_addr.sibling();
incomplete.push(IncompleteAt {
@ -309,31 +314,32 @@ fn combine_with_empty<H: Hashable + Clone + PartialEq>(
// and in position order. Returns the resulting tree, a flag indicating whether the
// resulting tree contains a `MARKED` node, and the vector of [`IncompleteAt`] values for
// [`Node::Nil`] nodes that were introduced in the process of constructing the tree.
#[allow(clippy::type_complexity)]
fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
mut xs: Vec<(LocatedPrunableTree<H>, bool)>,
root_addr: Address,
prune_below: Level,
) -> Option<(LocatedPrunableTree<H>, bool, Vec<IncompleteAt>)> {
) -> Result<Option<(LocatedPrunableTree<H>, bool, Vec<IncompleteAt>)>, 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 +352,7 @@ fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
// 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 +360,7 @@ fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
// 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<LocatedPrunableTree<H>>, (next_tree, next_marked)| {
if let Some(mut prev_tree) = acc {
@ -368,19 +374,19 @@ fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
&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)
}
}

50
shardtree/src/legacy.rs
View File

@ -101,7 +101,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
leaf_addr: Address,
mut filled: impl Iterator<Item = H>,
split_at: Level,
) -> (Self, Option<Self>) {
) -> Result<(Self, Option<Self>), 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<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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(
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<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
None
};
(subtree, supertree)
Ok((subtree, supertree))
}
/// Insert the nodes belonging to the given incremental witness to this tree, truncating the
@ -196,10 +204,15 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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| {
let cursor_trees = witness
.cursor()
.as_ref()
.filter(|c| c.size() > 0)
.map(|c| {
Self::from_frontier_parts(
witness.tip_position(),
c.leaf()
@ -212,7 +225,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
},
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!(

32
shardtree/src/lib.rs
View File

@ -414,44 +414,44 @@ impl<
fn go<H: Hashable + Clone + PartialEq>(
root_addr: Address,
root: &PrunableTree<H>,
) -> Option<(PrunableTree<H>, Position)> {
) -> Result<Option<(PrunableTree<H>, 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,
))
)
.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,

224
shardtree/src/prunable.rs
View File

@ -258,22 +258,30 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
/// `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<Arc<H>>, left: Self, right: Self) -> Result<Self, ()> {
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<Self, MergeError> {
/// Pre-condition: `root_addr` must be the address of `t0` and `t1`.
#[allow(clippy::type_complexity)]
@ -330,12 +338,13 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
{
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<H: Hashable + Clone + PartialEq> PrunableTree<H> {
/// replacement `Nil` value).
///
/// `level` must be the level of the two nodes that are being joined.
pub(crate) fn unite(level: Level, ann: Option<Arc<H>>, left: Self, right: Self) -> Self {
pub(crate) fn unite(
level: Level,
ann: Option<Arc<H>>,
left: Self,
right: Self,
) -> Result<Self, ()> {
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<H: Hashable + Clone + PartialEq> PrunableTree<H> {
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<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
position: Position,
root_addr: Address,
root: &PrunableTree<H>,
) -> Option<PrunableTree<H>> {
) -> Result<Option<PrunableTree<H>>, Address> {
match &root.0 {
Node::Parent { ann, left, right } => {
let (l_child, r_child) = root_addr
@ -586,31 +600,49 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// 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 {
go(position, self.root_addr, &self.root)
.expect("truncation does not introduce invalid subtree roots")
.map(|root| LocatedTree {
root_addr: self.root_addr,
root,
})
@ -619,6 +651,37 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
}
}
// 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. The newly created root node will be annotated with the
// provided value.
//
// If the root level of `self` is greater than or equal to the requested level, no extension
// will be performed, but the root node will still be reannotated.
fn extend_to_level(
self,
level: Level,
ann: Option<Arc<H>>,
) -> Result<(PrunableTree<H>, Vec<Address>), 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 +707,15 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
root_addr: Address,
into: &PrunableTree<H>,
subtree: LocatedPrunableTree<H>,
contains_marked: bool,
) -> Result<(PrunableTree<H>, Vec<IncompleteAt>), InsertionError> {
) -> Result<(PrunableTree<H>, Vec<Address>), 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<Arc<H>>, mut node: LocatedPrunableTree<H>| {
// 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 +770,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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 +791,25 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
.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(
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(),
),
incomplete,
))
)
.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(
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,
),
incomplete,
))
)
.map_err(|_| InsertionError::InputMalformed(root_addr))
.map(|t| (t, incomplete))
}
}
}
@ -787,13 +825,22 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
);
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 +883,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
frontier: NonEmptyFrontier<H>,
leaf_retention: &Retention<C>,
split_at: Level,
) -> (Self, Option<Self>) {
) -> Result<(Self, Option<Self>), Address> {
let (position, leaf, ommers) = frontier.into_parts();
Self::from_frontier_parts(position, leaf, ommers.into_iter(), leaf_retention, split_at)
}
@ -850,7 +897,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
mut ommers: impl Iterator<Item = H>,
leaf_retention: &Retention<C>,
split_at: Level,
) -> (Self, Option<Self>) {
) -> Result<(Self, Option<Self>), Address> {
let mut addr = Address::from(position);
let mut subtree = Tree::leaf((leaf, leaf_retention.into()));
@ -858,12 +905,17 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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 +934,24 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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(
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 +965,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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 +988,8 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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 +1010,10 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
to_clear: &[(Position, RetentionFlags)],
root_addr: Address,
root: &PrunableTree<H>,
) -> PrunableTree<H> {
) -> Result<PrunableTree<H>, 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 +1023,18 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
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 +1044,13 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// 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 +1058,8 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
let to_clear = to_clear.into_iter().collect::<Vec<_>>();
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"),
}
}