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

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.
This commit is contained in:
Kris Nuttycombe 2025-02-17 13:44:22 -07:00
parent 561a6dc29b
commit 1d7cd427f0
4 changed files with 254 additions and 171 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
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 {
@ -313,27 +318,27 @@ 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 +351,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 +359,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 +373,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)
} }
} }

50
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(
PrunableTree::parent_checked(
None, None,
supertree.unwrap_or_else(PrunableTree::empty), supertree.unwrap_or_else(PrunableTree::empty),
Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)), 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,10 +204,15 @@ 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
.cursor()
.as_ref()
.filter(|c| c.size() > 0)
.map(|c| {
Self::from_frontier_parts( Self::from_frontier_parts(
witness.tip_position(), witness.tip_position(),
c.leaf() c.leaf()
@ -212,7 +225,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
}, },
self.root_addr.level(), 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!(

32
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,
), )
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)), 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,

220
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,31 +600,49 @@ 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)
.expect("truncation does not introduce invalid subtree roots")
.map(|root| LocatedTree {
root_addr: self.root_addr, root_addr: self.root_addr,
root, root,
}) })
@ -619,6 +651,33 @@ 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
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 +703,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 +766,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 +787,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((
Tree::unite(
root_addr.level() - 1, root_addr.level() - 1,
ann.clone(), ann.clone(),
new_left, new_left,
right.as_ref().clone(), right.as_ref().clone(),
), )
incomplete, .map_err(|_| InsertionError::InputMalformed(root_addr))
)) .map(|t| (t, 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((
Tree::unite(
root_addr.level() - 1, root_addr.level() - 1,
ann.clone(), ann.clone(),
left.as_ref().clone(), left.as_ref().clone(),
new_right, new_right,
), )
incomplete, .map_err(|_| InsertionError::InputMalformed(root_addr))
)) .map(|t| (t, incomplete))
} }
} }
} }
@ -787,13 +821,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 +879,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 +893,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 +901,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 +930,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(
PrunableTree::parent_checked(
None, None,
Tree::leaf((left, RetentionFlags::EPHEMERAL)), Tree::leaf((left, RetentionFlags::EPHEMERAL)),
supertree.unwrap_or_else(Tree::empty), supertree.unwrap_or_else(Tree::empty),
)); )
.map_err(|_| addr)?,
);
addr = addr.parent(); addr = addr.parent();
} }
@ -906,7 +961,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 +984,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 +1006,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 +1019,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 +1040,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 +1054,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"),
} }
} }