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::{ use crate::{
error::{InsertionError, ShardTreeError}, error::{InsertionError, ShardTreeError},
store::{Checkpoint, ShardStore}, store::{Checkpoint, ShardStore},
IncompleteAt, LocatedPrunableTree, LocatedTree, RetentionFlags, ShardTree, Tree, IncompleteAt, LocatedPrunableTree, LocatedTree, PrunableTree, RetentionFlags, ShardTree, Tree,
}; };
impl< impl<
@ -215,7 +215,8 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// fragments up the stack until we get the largest possible subtree // fragments up the stack until we get the largest possible subtree
while let Some((potential_sibling, marked)) = fragments.pop() { while let Some((potential_sibling, marked)) = fragments.pop() {
if potential_sibling.root_addr.parent() == subtree.root_addr.parent() { 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 { } else {
// this is not a sibling node, so we push it back on to the stack // this is not a sibling node, so we push it back on to the stack
// and are done // and are done
@ -238,7 +239,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
let minimal_tree_addr = let minimal_tree_addr =
Address::from(position_range.start).common_ancestor(&last_position.into()); Address::from(position_range.start).common_ancestor(&last_position.into());
trace!("Building minimal tree at {:?}", minimal_tree_addr); 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 { |(to_insert, contains_marked, incomplete)| BatchInsertionResult {
subtree: to_insert, subtree: to_insert,
contains_marked, contains_marked,
@ -263,16 +266,18 @@ fn unite<H: Hashable + Clone + PartialEq>(
lroot: LocatedPrunableTree<H>, lroot: LocatedPrunableTree<H>,
rroot: LocatedPrunableTree<H>, rroot: LocatedPrunableTree<H>,
prune_below: Level, prune_below: Level,
) -> LocatedPrunableTree<H> { ) -> Result<LocatedPrunableTree<H>, Address> {
assert_eq!(lroot.root_addr.parent(), rroot.root_addr.parent()); assert_eq!(lroot.root_addr.parent(), rroot.root_addr.parent());
LocatedTree { Ok(LocatedTree {
root_addr: lroot.root_addr.parent(), root_addr: lroot.root_addr.parent(),
root: if lroot.root_addr.level() < prune_below { 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 { } 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 /// 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>, incomplete: &mut Vec<IncompleteAt>,
contains_marked: bool, contains_marked: bool,
prune_below: Level, prune_below: Level,
) -> LocatedPrunableTree<H> { ) -> Result<LocatedPrunableTree<H>, Address> {
assert_eq!(expect_left_child, root.root_addr.is_left_child()); assert_eq!(expect_left_child, root.root_addr.is_left_child());
let sibling_addr = root.root_addr.sibling(); let sibling_addr = root.root_addr.sibling();
incomplete.push(IncompleteAt { 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 // 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 // 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. // [`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>( fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
mut xs: Vec<(LocatedPrunableTree<H>, bool)>, mut xs: Vec<(LocatedPrunableTree<H>, bool)>,
root_addr: Address, root_addr: Address,
prune_below: Level, 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 // First, consume the stack from the right, building up a single tree
// until we can't combine any more. // until we can't combine any more.
if let Some((mut cur, mut contains_marked)) = xs.pop() { if let Some((mut cur, mut contains_marked)) = xs.pop() {
let mut incomplete = vec![]; let mut incomplete = vec![];
while let Some((top, top_marked)) = xs.pop() { while let Some((top, top_marked)) = xs.pop() {
while cur.root_addr.level() < top.root_addr.level() { 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() { if cur.root_addr.level() == top.root_addr.level() {
contains_marked = contains_marked || top_marked; contains_marked = contains_marked || top_marked;
if cur.root_addr.is_right_child() { if cur.root_addr.is_right_child() {
// We have a left child and a right child, so unite them. // 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 { } else {
// This is a left child, so we build it up one more level and then // 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 // we've merged as much as we can from the right and need to work from
// the left // the left
xs.push((top, top_marked)); 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; break;
} }
} else { } 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 // 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() { 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. // 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 // From the stack of subtrees, construct a single sparse tree that can be
// inserted/merged into the existing tree // inserted/merged into the existing tree
let res_tree = xs.into_iter().fold( let res_tree = xs.into_iter().try_fold(
None, None,
|acc: Option<LocatedPrunableTree<H>>, (next_tree, next_marked)| { |acc: Option<LocatedPrunableTree<H>>, (next_tree, next_marked)| {
if let Some(mut prev_tree) = acc { if let Some(mut prev_tree) = acc {
@ -368,19 +374,19 @@ fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
&mut incomplete, &mut incomplete,
next_marked, next_marked,
prune_below, prune_below,
); )?;
} }
Some(unite(prev_tree, next_tree, prune_below)) Ok::<_, Address>(Some(unite(prev_tree, next_tree, prune_below)?))
} else { } 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 { } else {
None Ok(None)
} }
} }

80
shardtree/src/legacy.rs
View File

@ -101,7 +101,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
leaf_addr: Address, leaf_addr: Address,
mut filled: impl Iterator<Item = H>, mut filled: impl Iterator<Item = H>,
split_at: Level, 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 // 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 // whether or not these nodes can be invalidated by a rewind
let mut addr = leaf_addr; let mut addr = leaf_addr;
@ -113,17 +113,19 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
if let Some(right) = filled.next() { if let Some(right) = filled.next() {
// once we have a right-hand node, add a parent with the current tree // once we have a right-hand node, add a parent with the current tree
// as the left-hand sibling // as the left-hand sibling
subtree = Tree::parent( subtree = PrunableTree::parent_checked(
None, None,
subtree, subtree,
Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)), Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)),
); )
.map_err(|_| addr)?;
} else { } else {
break; break;
} }
} else { } else {
// the current address is for a right child, so add an empty left sibling // 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(); addr = addr.parent();
@ -140,16 +142,22 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
for right in filled { for right in filled {
// build up the right-biased tree until we get a left-hand node // build up the right-biased tree until we get a left-hand node
while addr.is_right_child() { 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(); addr = addr.parent();
} }
// once we have a left-hand root, add a parent with the current ommer as the right-hand sibling // 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(
None, PrunableTree::parent_checked(
supertree.unwrap_or_else(PrunableTree::empty), None,
Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)), supertree.unwrap_or_else(PrunableTree::empty),
)); Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)),
)
.map_err(|_| addr)?,
);
addr = addr.parent(); addr = addr.parent();
} }
@ -161,7 +169,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
None None
}; };
(subtree, supertree) Ok((subtree, supertree))
} }
/// Insert the nodes belonging to the given incremental witness to this tree, truncating the /// Insert the nodes belonging to the given incremental witness to this tree, truncating the
@ -196,23 +204,30 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
Address::from(witness.witnessed_position()), Address::from(witness.witnessed_position()),
witness.filled().iter().cloned(), witness.filled().iter().cloned(),
self.root_addr.level(), self.root_addr.level(),
); )
.map_err(InsertionError::InputMalformed)?;
// construct subtrees from the `cursor` part of the witness // 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
Self::from_frontier_parts( .cursor()
witness.tip_position(), .as_ref()
c.leaf() .filter(|c| c.size() > 0)
.cloned() .map(|c| {
.expect("Cannot have an empty leaf for a non-empty tree"), Self::from_frontier_parts(
c.ommers_iter().cloned(), witness.tip_position(),
&Retention::Checkpoint { c.leaf()
id: checkpoint_id, .cloned()
marking: Marking::None, .expect("Cannot have an empty leaf for a non-empty tree"),
}, c.ommers_iter().cloned(),
self.root_addr.level(), &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)?; let (subtree, _) = past_subtree.insert_subtree(future_subtree, true)?;
@ -253,7 +268,7 @@ mod tests {
frontier::CommitmentTree, witness::IncrementalWitness, Address, Level, Position, frontier::CommitmentTree, witness::IncrementalWitness, Address, Level, Position,
}; };
use crate::{LocatedPrunableTree, RetentionFlags, Tree}; use crate::{LocatedPrunableTree, PrunableTree, RetentionFlags, Tree};
#[test] #[test]
fn insert_witness_nodes() { fn insert_witness_nodes() {
@ -280,19 +295,22 @@ mod tests {
assert_eq!( assert_eq!(
c.root, c.root,
Tree::parent( PrunableTree::parent_checked(
None, None,
Tree::parent( PrunableTree::parent_checked(
None, None,
Tree::empty(), Tree::empty(),
Tree::leaf(("ijklmnop".to_string(), RetentionFlags::EPHEMERAL)), Tree::leaf(("ijklmnop".to_string(), RetentionFlags::EPHEMERAL)),
), )
Tree::parent( .unwrap(),
PrunableTree::parent_checked(
None, None,
Tree::leaf(("qrstuvwx".to_string(), RetentionFlags::EPHEMERAL)), Tree::leaf(("qrstuvwx".to_string(), RetentionFlags::EPHEMERAL)),
Tree::empty() Tree::empty()
) )
.unwrap()
) )
.unwrap()
); );
assert_eq!( assert_eq!(

46
shardtree/src/lib.rs
View File

@ -414,44 +414,44 @@ impl<
fn go<H: Hashable + Clone + PartialEq>( fn go<H: Hashable + Clone + PartialEq>(
root_addr: Address, root_addr: Address,
root: &PrunableTree<H>, root: &PrunableTree<H>,
) -> Option<(PrunableTree<H>, Position)> { ) -> Result<Option<(PrunableTree<H>, Position)>, Address> {
match &root.0 { match &root.0 {
Node::Parent { ann, left, right } => { Node::Parent { ann, left, right } => {
let (l_addr, r_addr) = root_addr let (l_addr, r_addr) = root_addr
.children() .children()
.expect("has children because we checked `root` is a parent"); .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)| { go(l_addr, left)?
( .map(|(new_left, pos)| {
Tree::unite( PrunableTree::unite(
l_addr.level(), l_addr.level(),
ann.clone(), ann.clone(),
new_left, new_left,
Tree::empty(), Tree::empty(),
), )
pos, .map_err(|_| l_addr)
) .map(|t| (t, pos))
}) })
.transpose()
}, },
|(new_right, pos)| { |(new_right, pos)| {
Some(( Tree::unite(
Tree::unite( l_addr.level(),
l_addr.level(), ann.clone(),
ann.clone(), left.as_ref().clone(),
left.as_ref().clone(), new_right,
new_right, )
), .map_err(|_| l_addr)
pos, .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)), Tree::leaf((h.clone(), *r | RetentionFlags::CHECKPOINT)),
root_addr.max_position(), 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 // 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). // 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(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 self.store
.put_shard(LocatedTree { .put_shard(LocatedTree {
root_addr: subtree.root_addr, root_addr: subtree.root_addr,

256
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. /// `level` must be the level of the root of the node being pruned.
pub fn prune(self, level: Level) -> Self { pub fn prune(self, level: Level) -> Self {
match self { match self {
Tree(Node::Parent { ann, left, right }) => Tree::unite( Tree(Node::Parent { ann, left, right }) => PrunableTree::unite(
level, level,
ann, ann,
left.as_ref().clone().prune(level - 1), left.as_ref().clone().prune(level - 1),
right.as_ref().clone().prune(level - 1), right.as_ref().clone().prune(level - 1),
), )
.expect("pruning does not construct empty parent nodes"),
other => other, 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. /// Merge two subtrees having the same root address.
/// ///
/// The merge operation is checked to be strictly additive and returns an error if merging /// 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 /// 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. /// 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> { 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`. /// Pre-condition: `root_addr` must be the address of `t0` and `t1`.
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
@ -330,12 +338,13 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
{ {
let (l_addr, r_addr) = let (l_addr, r_addr) =
addr.children().ok_or(MergeError::TreeMalformed(addr))?; addr.children().ok_or(MergeError::TreeMalformed(addr))?;
Ok(Tree::unite( PrunableTree::unite(
addr.level() - 1, addr.level() - 1,
lann.or(rann), lann.or(rann),
go(l_addr, ll.as_ref().clone(), rl.as_ref().clone())?, go(l_addr, ll.as_ref().clone(), rl.as_ref().clone())?,
go(r_addr, lr.as_ref().clone(), rr.as_ref().clone())?, go(r_addr, lr.as_ref().clone(), rr.as_ref().clone())?,
)) )
.map_err(|_| MergeError::TreeMalformed(addr))
} else { } else {
unreachable!() unreachable!()
} }
@ -355,9 +364,14 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
/// replacement `Nil` value). /// replacement `Nil` value).
/// ///
/// `level` must be the level of the two nodes that are being joined. /// `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) { 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 })) (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 // 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 // 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 && if lv.1 == RetentionFlags::EPHEMERAL &&
(rv.1 & (RetentionFlags::MARKED | RetentionFlags::REFERENCE)) == 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, position: Position,
root_addr: Address, root_addr: Address,
root: &PrunableTree<H>, root: &PrunableTree<H>,
) -> Option<PrunableTree<H>> { ) -> Result<Option<PrunableTree<H>>, Address> {
match &root.0 { match &root.0 {
Node::Parent { ann, left, right } => { Node::Parent { ann, left, right } => {
let (l_child, r_child) = root_addr let (l_child, r_child) = root_addr
@ -586,39 +600,88 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// we are truncating within the range of the left node, so recurse // we are truncating within the range of the left node, so recurse
// to the left to truncate the left child and then reconstruct the // to the left to truncate the left child and then reconstruct the
// node with `Nil` as the right sibling // node with `Nil` as the right sibling
go(position, l_child, left.as_ref()).map(|left| { go(position, l_child, left.as_ref())?
Tree::unite(l_child.level(), ann.clone(), left, Tree::empty()) .map(|left| {
}) PrunableTree::unite(
l_child.level(),
ann.clone(),
left,
Tree::empty(),
)
.map_err(|_| root_addr)
})
.transpose()
} else { } else {
// we are truncating within the range of the right node, so recurse // we are truncating within the range of the right node, so recurse
// to the right to truncate the right child and then reconstruct the // to the right to truncate the right child and then reconstruct the
// node with the left sibling unchanged // node with the left sibling unchanged
go(position, r_child, right.as_ref()).map(|right| { go(position, r_child, right.as_ref())?
Tree::unite(r_child.level(), ann.clone(), left.as_ref().clone(), right) .map(|right| {
}) PrunableTree::unite(
r_child.level(),
ann.clone(),
left.as_ref().clone(),
right,
)
.map_err(|_| root_addr)
})
.transpose()
} }
} }
Node::Leaf { .. } => { Node::Leaf { .. } => {
if root_addr.max_position() <= position { if root_addr.max_position() <= position {
Some(root.clone()) Ok(Some(root.clone()))
} else { } else {
None Ok(None)
} }
} }
Node::Nil => None, Node::Nil => Ok(None),
} }
} }
if self.root_addr.position_range().contains(&position) { 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)
root_addr: self.root_addr, .expect("truncation does not introduce invalid subtree roots")
root, .map(|root| LocatedTree {
}) root_addr: self.root_addr,
root,
})
} else { } else {
None 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. 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 /// Inserts a descendant subtree into this subtree, creating empty sibling nodes as necessary
/// to fill out the tree. /// to fill out the tree.
/// ///
@ -644,38 +707,15 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
root_addr: Address, root_addr: Address,
into: &PrunableTree<H>, into: &PrunableTree<H>,
subtree: LocatedPrunableTree<H>, subtree: LocatedPrunableTree<H>,
contains_marked: bool, ) -> Result<(PrunableTree<H>, Vec<Address>), InsertionError> {
) -> Result<(PrunableTree<H>, Vec<IncompleteAt>), InsertionError> {
// An empty tree cannot replace any other type of tree. // An empty tree cannot replace any other type of tree.
if subtree.root().is_nil() { if subtree.root().is_nil() {
Ok((into.clone(), vec![])) Ok((into.clone(), vec![]))
} else { } 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 { 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 { Tree(Node::Leaf {
value: (value, retention), value: (value, retention),
}) => { }) => {
@ -730,7 +770,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
Err(InsertionError::Conflict(root_addr)) Err(InsertionError::Conflict(root_addr))
} }
} else { } 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 => { parent if root_addr == subtree.root_addr => {
@ -749,29 +791,25 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
.children() .children()
.expect("has children because we checked `into` is a parent"); .expect("has children because we checked `into` is a parent");
if l_addr.contains(&subtree.root_addr) { if l_addr.contains(&subtree.root_addr) {
let (new_left, incomplete) = let (new_left, incomplete) = go(l_addr, left.as_ref(), subtree)?;
go(l_addr, left.as_ref(), subtree, contains_marked)?; PrunableTree::unite(
Ok(( root_addr.level() - 1,
Tree::unite( ann.clone(),
root_addr.level() - 1, new_left,
ann.clone(), right.as_ref().clone(),
new_left, )
right.as_ref().clone(), .map_err(|_| InsertionError::InputMalformed(root_addr))
), .map(|t| (t, incomplete))
incomplete,
))
} else { } else {
let (new_right, incomplete) = let (new_right, incomplete) = go(r_addr, right.as_ref(), subtree)?;
go(r_addr, right.as_ref(), subtree, contains_marked)?; PrunableTree::unite(
Ok(( root_addr.level() - 1,
Tree::unite( ann.clone(),
root_addr.level() - 1, left.as_ref().clone(),
ann.clone(), new_right,
left.as_ref().clone(), )
new_right, .map_err(|_| InsertionError::InputMalformed(root_addr))
), .map(|t| (t, incomplete))
incomplete,
))
} }
} }
} }
@ -787,13 +825,22 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
); );
let LocatedTree { root_addr, root } = self; let LocatedTree { root_addr, root } = self;
if root_addr.contains(&subtree.root_addr) { 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 { let new_tree = LocatedTree {
root_addr: *root_addr, root_addr: *root_addr,
root, root,
}; };
assert!(new_tree.max_position() >= max_position); 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 { } else {
Err(InsertionError::NotContained(subtree.root_addr)) Err(InsertionError::NotContained(subtree.root_addr))
@ -836,7 +883,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
frontier: NonEmptyFrontier<H>, frontier: NonEmptyFrontier<H>,
leaf_retention: &Retention<C>, leaf_retention: &Retention<C>,
split_at: Level, split_at: Level,
) -> (Self, Option<Self>) { ) -> Result<(Self, Option<Self>), Address> {
let (position, leaf, ommers) = frontier.into_parts(); let (position, leaf, ommers) = frontier.into_parts();
Self::from_frontier_parts(position, leaf, ommers.into_iter(), leaf_retention, split_at) 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>, mut ommers: impl Iterator<Item = H>,
leaf_retention: &Retention<C>, leaf_retention: &Retention<C>,
split_at: Level, split_at: Level,
) -> (Self, Option<Self>) { ) -> Result<(Self, Option<Self>), Address> {
let mut addr = Address::from(position); let mut addr = Address::from(position);
let mut subtree = Tree::leaf((leaf, leaf_retention.into())); 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() { if addr.is_left_child() {
// the current address is a left child, so create a parent with // the current address is a left child, so create a parent with
// an empty right-hand tree // 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() { } else if let Some(left) = ommers.next() {
// the current address corresponds to a right child, so create a parent that // the current address corresponds to a right child, so create a parent that
// takes the left sibling to that child from the ommers // takes the left sibling to that child from the ommers
subtree = subtree = PrunableTree::parent_checked(
Tree::parent(None, Tree::leaf((left, RetentionFlags::EPHEMERAL)), subtree); None,
Tree::leaf((left, RetentionFlags::EPHEMERAL)),
subtree,
)
.map_err(|_| addr)?;
} else { } else {
break; break;
} }
@ -882,17 +934,24 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
for left in ommers { for left in ommers {
// build up the left-biased tree until we get a right-hand node // build up the left-biased tree until we get a right-hand node
while addr.is_left_child() { 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(); addr = addr.parent();
} }
// once we have a right-hand root, add a parent with the current ommer as the // once we have a right-hand root, add a parent with the current ommer as the
// left-hand sibling // left-hand sibling
supertree = Some(Tree::parent( supertree = Some(
None, PrunableTree::parent_checked(
Tree::leaf((left, RetentionFlags::EPHEMERAL)), None,
supertree.unwrap_or_else(Tree::empty), Tree::leaf((left, RetentionFlags::EPHEMERAL)),
)); supertree.unwrap_or_else(Tree::empty),
)
.map_err(|_| addr)?,
);
addr = addr.parent(); addr = addr.parent();
} }
@ -906,7 +965,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
None 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 /// 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()) { if subtree_range.contains(&frontier.position()) {
let leaf_is_marked = leaf_retention.is_marked(); let leaf_is_marked = leaf_retention.is_marked();
let (subtree, supertree) = 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; 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)], to_clear: &[(Position, RetentionFlags)],
root_addr: Address, root_addr: Address,
root: &PrunableTree<H>, root: &PrunableTree<H>,
) -> PrunableTree<H> { ) -> Result<PrunableTree<H>, Address> {
if to_clear.is_empty() { if to_clear.is_empty() {
// nothing to do, so we just return the root // nothing to do, so we just return the root
root.clone() Ok(root.clone())
} else { } else {
match &root.0 { match &root.0 {
Node::Parent { ann, left, right } => { 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()); let p = to_clear.partition_point(|(p, _)| p < &l_addr.position_range_end());
trace!( trace!(
"Tree::unite at {:?}, partitioned: {:?} {:?}", "PrunableTree::unite at {:?}, partitioned: {:?} {:?}",
root_addr, root_addr,
&to_clear[0..p], &to_clear[0..p],
&to_clear[p..], &to_clear[p..],
); );
Tree::unite( PrunableTree::unite(
l_addr.level(), l_addr.level(),
ann.clone(), ann.clone(),
go(&to_clear[0..p], l_addr, left), go(&to_clear[0..p], l_addr, left)?,
go(&to_clear[p..], r_addr, right), go(&to_clear[p..], r_addr, right)?,
) )
.map_err(|_| root_addr)
} }
Node::Leaf { value: (h, r) } => { Node::Leaf { value: (h, r) } => {
trace!("In {:?}, clearing {:?}", root_addr, to_clear); 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 // a partially-pruned branch, and if it's a marked node then it will
// be a level-0 leaf. // be a level-0 leaf.
match to_clear { match to_clear {
[(_, flags)] => Tree::leaf((h.clone(), *r & !*flags)), [(_, flags)] => Ok(Tree::leaf((h.clone(), *r & !*flags))),
_ => { _ => {
panic!("Tree state inconsistent with checkpoints."); 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<_>>(); let to_clear = to_clear.into_iter().collect::<Vec<_>>();
Self { Self {
root_addr: self.root_addr, 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"),
} }
} }