diff --git a/shardtree/src/lib.rs b/shardtree/src/lib.rs index df5222d..c83eb34 100644 --- a/shardtree/src/lib.rs +++ b/shardtree/src/lib.rs @@ -1,7 +1,7 @@ use core::fmt::Debug; -use core::ops::{BitAnd, BitOr, Deref, Not}; +use core::ops::{BitAnd, BitOr, Deref, Not, Range}; use either::Either; -use std::collections::BTreeSet; +use std::collections::{BTreeMap, BTreeSet}; use std::rc::Rc; use incrementalmerkletree::{Address, Hashable, Level, Position, Retention}; @@ -688,6 +688,754 @@ impl LocatedTree { } } +type LocatedPrunableTree = LocatedTree>, (H, RetentionFlags)>; + +/// A data structure describing the nature of a [`Node::Nil`] node in the tree that was introduced +/// as the consequence of an insertion. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct IncompleteAt { + /// The address of the empty node. + pub address: Address, + /// A flag identifying whether or not the missing node is required in order to construct a + /// witness for a node with [`MARKED`] retention. + pub required_for_witness: bool, +} + +/// A type for the result of a batch insertion operation. +/// +/// This result type contains the newly constructed tree, the addresses any new incomplete internal +/// nodes within that tree that were introduced as a consequence of that insertion, and the +/// remainder of the iterator that provided the inserted values. +#[derive(Debug)] +pub struct BatchInsertionResult)>> { + /// The updated tree after all insertions have been performed. + pub subtree: LocatedPrunableTree, + /// A flag identifying whether the constructed subtree contains a marked node. + pub contains_marked: bool, + /// The vector of addresses of [`Node::Nil`] nodes that were inserted into the tree as part of + /// the insertion operation, for nodes that are required in order to construct a witness for + /// each inserted leaf with [`MARKED`] retention. + pub incomplete: Vec, + /// The maximum position at which a leaf was inserted. + pub max_insert_position: Option, + /// The positions of all leaves with [`CHECKPOINT`] retention that were inserted. + pub checkpoints: BTreeMap, + /// The unconsumed remainder of the iterator from which leaves were inserted, if the tree + /// was completely filled before the iterator was fully consumed. + pub remainder: I, +} + +/// An error prevented the insertion of values into the subtree. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum InsertionError { + /// The caller attempted to insert a subtree into a tree that does not contain + /// the subtree's root address. + NotContained, + /// The start of the range of positions provided for insertion is not included + /// in the range of positions within this subtree. + OutOfRange(Range), + /// An existing root hash conflicts with the root hash of a node being inserted. + Conflict(Address), + /// An out-of-order checkpoint was detected + /// + /// Checkpoint identifiers must be in nondecreasing order relative to tree positions. + CheckpointOutOfOrder, + /// An append operation has exceeded the capacity of the tree. + TreeFull, + /// An error was produced by the underlying [`ShardStore`] + Storage(S), +} + +/// Errors that may be returned in the process of querying a [`ShardTree`] +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum QueryError { + /// The caller attempted to query the value at an address within a tree that does not contain + /// that address. + NotContained(Address), + /// A leaf required by a given checkpoint has been pruned, or is otherwise not accessible in + /// the tree. + CheckpointPruned, + /// It is not possible to compute a root for one or more subtrees because they contain + /// [`Node::Nil`] values at positions that cannot be replaced with default hashes. + TreeIncomplete(Vec
), +} + +/// Operations on [`LocatedTree`]s that are annotated with Merkle hashes. +impl LocatedPrunableTree { + /// Computes the root hash of this tree, truncated to the given position. + /// + /// If the tree contains any [`Node::Nil`] nodes corresponding to positions less than + /// `truncate_at`, this will return an error containing the addresses of those nodes within the + /// tree. + pub fn root_hash(&self, truncate_at: Position) -> Result> { + self.root.root_hash(self.root_addr, truncate_at) + } + + /// Compute the root hash of this subtree, filling empty nodes along the rightmost path of the + /// subtree with the empty root value for the given level. + /// + /// This should only be used for computing roots when it is known that no successor trees + /// exist. + /// + /// If the tree contains any [`Node::Nil`] nodes that are to the left of filled nodes in the + /// tree, this will return an error containing the addresses of those nodes. + pub fn right_filled_root(&self) -> Result> { + self.root_hash( + self.max_position() + .map_or_else(|| self.root_addr.position_range_start(), |pos| pos + 1), + ) + } + + /// Returns the positions of marked leaves in the tree. + pub fn marked_positions(&self) -> BTreeSet { + fn go( + root_addr: Address, + root: &PrunableTree, + acc: &mut BTreeSet, + ) { + match &root.0 { + Node::Parent { left, right, .. } => { + let (l_addr, r_addr) = root_addr.children().unwrap(); + go(l_addr, left.as_ref(), acc); + go(r_addr, right.as_ref(), acc); + } + Node::Leaf { value } => { + if value.1.is_marked() && root_addr.level() == 0.into() { + acc.insert(Position::from(root_addr.index())); + } + } + _ => {} + } + } + + let mut result = BTreeSet::new(); + go(self.root_addr, &self.root, &mut result); + result + } + + /// Compute the witness for the leaf at the specified position. + /// + /// This tree will be truncated to the `truncate_at` position, and then empty + /// empty roots corresponding to later positions will be filled by [`H::empty_root`]. + /// + /// Returns either the witness for the leaf at the specified position, or an error that + /// describes the causes of failure. + pub fn witness(&self, position: Position, truncate_at: Position) -> Result, QueryError> { + // traverse down to the desired leaf position, and then construct + // the authentication path on the way back up. + fn go( + root: &PrunableTree, + root_addr: Address, + position: Position, + truncate_at: Position, + ) -> Result, Vec
> { + match &root.0 { + Node::Parent { left, right, .. } => { + let (l_addr, r_addr) = root_addr.children().unwrap(); + if root_addr.level() > 1.into() { + let r_start = r_addr.position_range_start(); + if position < r_start { + accumulate_result_with( + go(left.as_ref(), l_addr, position, truncate_at), + right.as_ref().root_hash(r_addr, truncate_at), + |mut witness, sibling_root| { + witness.push(sibling_root); + witness + }, + ) + } else { + // if the position we're witnessing is down the right-hand branch then + // we always set the truncation bound outside the range of leaves on the + // left, because we don't allow any empty nodes to the left + accumulate_result_with( + left.as_ref().root_hash(l_addr, r_start), + go(right.as_ref(), r_addr, position, truncate_at), + |sibling_root, mut witness| { + witness.push(sibling_root); + witness + }, + ) + } + } else { + // we handle the level 0 leaves here by adding the sibling of our desired + // leaf to the witness + if position.is_odd() { + if right.is_marked_leaf() { + left.leaf_value() + .map(|v| vec![v.clone()]) + .ok_or_else(|| vec![l_addr]) + } else { + Err(vec![l_addr]) + } + } else if left.is_marked_leaf() { + // If we have the left-hand leaf and the right-hand leaf is empty, we + // can fill it with the empty leaf, but only if `fill_start` is None or + // it is located at `position + 1`. + if truncate_at <= position + 1 { + Ok(vec![H::empty_leaf()]) + } else { + right + .leaf_value() + .map_or_else(|| Err(vec![r_addr]), |v| Ok(vec![v.clone()])) + } + } else { + Err(vec![r_addr]) + } + } + } + _ => { + // if we encounter a nil or leaf node, we were unable to descend + // to the leaf at the desired position. + Err(vec![root_addr]) + } + } + } + + if self.root_addr.position_range().contains(&position) { + go(&self.root, self.root_addr, position, truncate_at) + .map_err(QueryError::TreeIncomplete) + } else { + Err(QueryError::NotContained(self.root_addr)) + } + } + + /// Prunes this tree by replacing all nodes that are right-hand children along the path + /// to the specified position with [`Node::Nil`]. + /// + /// The leaf at the specified position is retained. + pub fn truncate_to_position(&self, position: Position) -> Option { + fn go( + position: Position, + root_addr: Address, + root: &PrunableTree, + ) -> Option> { + match &root.0 { + Node::Parent { ann, left, right } => { + let (l_child, r_child) = root_addr.children().unwrap(); + if position < r_child.position_range_start() { + // we are truncating within the range of the left node, so recurse + // to the left to truncate the left child and then reconstruct the + // 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(Node::Nil)) + }) + } 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) + }) + } + } + Node::Leaf { .. } => { + if root_addr.max_position() <= position { + Some(root.clone()) + } else { + None + } + } + Node::Nil => None, + } + } + + if self.root_addr.position_range().contains(&position) { + go(position, self.root_addr, &self.root).map(|root| LocatedTree { + root_addr: self.root_addr, + root, + }) + } else { + None + } + } + + /// Inserts a descendant subtree into this subtree, creating empty sibling nodes as necessary + /// to fill out the tree. + /// + /// In the case that a leaf node would be replaced by an incomplete subtree, the resulting + /// parent node will be annotated with the existing leaf value. + /// + /// Returns the updated tree, along with the addresses of any [`Node::Nil`] nodes that were + /// inserted in the process of creating the parent nodes down to the insertion point, or an + /// error if the specified subtree's root address is not in the range of valid descendants of + /// the root node of this tree or if the insertion would result in a conflict between computed + /// root hashes of complete subtrees. + pub fn insert_subtree( + &self, + subtree: Self, + contains_marked: bool, + ) -> Result<(Self, Vec), InsertionError> { + // A function to recursively dig into the tree, creating a path downward and introducing + // empty nodes as necessary until we can insert the provided subtree. + #[allow(clippy::type_complexity)] + fn go( + root_addr: Address, + into: &PrunableTree, + subtree: LocatedPrunableTree, + is_complete: bool, + contains_marked: bool, + ) -> Result<(PrunableTree, Vec), InsertionError> { + // In the case that we are replacing a node entirely, we need to extend the + // subtree up to the level of the node being replaced, adding Nil siblings + // and recording the presence of those incomplete nodes when necessary + let replacement = |ann: Option>, mut node: LocatedPrunableTree| { + // construct the replacement node bottom-up + let mut incomplete = vec![]; + while node.root_addr.level() < root_addr.level() { + incomplete.push(IncompleteAt { + address: node.root_addr.sibling(), + required_for_witness: contains_marked, + }); + node = LocatedTree { + root_addr: node.root_addr.parent(), + root: if node.root_addr.is_right_child() { + Tree(Node::Parent { + ann: None, + left: Rc::new(Tree(Node::Nil)), + right: Rc::new(node.root), + }) + } else { + Tree(Node::Parent { + ann: None, + left: Rc::new(node.root), + right: Rc::new(Tree(Node::Nil)), + }) + }, + }; + } + (node.root.reannotate_root(ann), incomplete) + }; + + match into { + Tree(Node::Nil) => Ok(replacement(None, subtree)), + Tree(Node::Leaf { value: (value, _) }) => { + if root_addr == subtree.root_addr { + if is_complete { + // It is safe to replace the existing root unannotated, because we + // can always recompute the root from a complete subtree. + Ok((subtree.root, vec![])) + } else if subtree + .root + .0 + .annotation() + .and_then(|ann| ann.as_ref()) + .iter() + .all(|v| v.as_ref() == value) + { + Ok(( + // at this point we statically know the root to be a parent + subtree.root.reannotate_root(Some(Rc::new(value.clone()))), + vec![], + )) + } else { + Err(InsertionError::Conflict(root_addr)) + } + } else { + Ok(replacement(Some(Rc::new(value.clone())), subtree)) + } + } + parent if root_addr == subtree.root_addr => { + // Merge the existing subtree with the subtree being inserted. + // A merge operation can't introduce any new incomplete roots. + parent + .clone() + .merge_checked(root_addr, subtree.root) + .map_err(InsertionError::Conflict) + .map(|tree| (tree, vec![])) + } + Tree(Node::Parent { ann, left, right }) => { + // In this case, we have an existing parent but we need to dig down farther + // before we can insert the subtree that we're carrying for insertion. + let (l_addr, r_addr) = root_addr.children().unwrap(); + if l_addr.contains(&subtree.root_addr) { + let (new_left, incomplete) = + go(l_addr, left.as_ref(), subtree, is_complete, contains_marked)?; + Ok(( + Tree::unite( + root_addr.level() - 1, + ann.clone(), + new_left, + right.as_ref().clone(), + ), + incomplete, + )) + } else { + let (new_right, incomplete) = go( + r_addr, + right.as_ref(), + subtree, + is_complete, + contains_marked, + )?; + Ok(( + Tree::unite( + root_addr.level() - 1, + ann.clone(), + left.as_ref().clone(), + new_right, + ), + incomplete, + )) + } + } + } + } + + let LocatedTree { root_addr, root } = self; + if root_addr.contains(&subtree.root_addr) { + let complete = subtree.root.reduce(&is_complete); + go(*root_addr, root, subtree, complete, contains_marked).map(|(root, incomplete)| { + ( + LocatedTree { + root_addr: *root_addr, + root, + }, + incomplete, + ) + }) + } else { + Err(InsertionError::NotContained) + } + } + + /// Append a single value at the first available position in the tree. + /// + /// Prefer to use [`Self::batch_append`] or [`Self::batch_insert`] when appending multiple + /// values, as these operations require fewer traversals of the tree than are necessary when + /// performing multiple sequential calls to [`Self::append`]. + pub fn append( + &self, + value: H, + retention: Retention, + ) -> Result<(Self, Position, Option), InsertionError> { + let checkpoint_id = if let Retention::Checkpoint { id, .. } = &retention { + Some(id.clone()) + } else { + None + }; + + self.batch_append(Some((value, retention)).into_iter()) + // We know that the max insert position will have been incremented by one. + .and_then(|r| { + let mut r = r.expect("We know the iterator to have been nonempty."); + if r.remainder.next().is_some() { + Err(InsertionError::TreeFull) + } else { + Ok((r.subtree, r.max_insert_position.unwrap(), checkpoint_id)) + } + }) + } + + /// Append a values from an iterator, beginning at the first available position in the tree. + /// + /// Returns an error if the tree is full. If the position at the end of the iterator is outside + /// of the subtree's range, the unconsumed part of the iterator will be returned as part of + /// the result. + pub fn batch_append)>, E>( + &self, + values: I, + ) -> Result>, InsertionError> { + let append_position = self + .max_position() + .map(|p| p + 1) + .unwrap_or_else(|| self.root_addr.position_range_start()); + self.batch_insert(append_position, values) + } + + /// Builds a [`LocatedPrunableTree`] from an iterator of level-0 leaves. + /// + /// This may be used in conjunction with [`ShardTree::insert_tree`] to support + /// partially-parallelizable tree construction. Multiple subtrees may be constructed in + /// parallel from iterators over (preferably, though not necessarily) disjoint leaf ranges, and + /// [`ShardTree::insert_tree`] may be used to insert those subtrees into the `ShardTree` in + /// arbitrary order. + /// + /// * `position_range` - The range of leaf positions at which values will be inserted. This + /// range is also used to place an upper bound on the number of items that will be consumed + /// from the `values` iterator. + /// * `prune_below` - Nodes with [`EPHEMERAL`] retention that are not required to be retained + /// in order to construct a witness for a marked node or to make it possible to rewind to a + /// checkpointed node may be pruned so long as their address is at less than the specified + /// level. + /// * `values` The iterator of `(H, Retention)` pairs from which to construct the tree. + pub fn from_iter)>>( + position_range: Range, + prune_below: Level, + mut values: I, + ) -> Option> { + // Unite two subtrees by either adding a parent node, or a leaf containing the Merkle root + // of such a parent if both nodes are ephemeral leaves. + // + // `unite` is only called when both root addrs have the same parent. `batch_insert` never + // constructs Nil nodes, so we don't create any incomplete root information here. + fn unite( + lroot: LocatedPrunableTree, + rroot: LocatedPrunableTree, + prune_below: Level, + ) -> LocatedTree>, (H, RetentionFlags)> { + 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) + } else { + Tree(Node::Parent { + ann: None, + left: Rc::new(lroot.root), + right: Rc::new(rroot.root), + }) + }, + } + } + + // Builds a single tree from the provided stack of subtrees, which must be non-overlapping + // 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. + fn build_minimal_tree( + mut xs: Vec<(LocatedPrunableTree, bool)>, + prune_below: Level, + ) -> Option<(LocatedPrunableTree, bool, Vec)> { + // 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() { + let sibling_addr = cur.root_addr.sibling(); + incomplete.push(IncompleteAt { + address: sibling_addr, + required_for_witness: top_marked, + }); + cur = unite( + cur, + LocatedTree { + root_addr: sibling_addr, + root: Tree(Node::Nil), + }, + 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); + } 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)); + let sibling_addr = cur.root_addr.sibling(); + incomplete.push(IncompleteAt { + address: sibling_addr, + required_for_witness: top_marked, + }); + cur = unite( + cur, + LocatedTree { + root_addr: sibling_addr, + root: Tree(Node::Nil), + }, + prune_below, + ); + break; + } + } else { + // top.root_addr.level < cur.root_addr.level, so we've merged as much as we + // can from the right and now need to work from the left. + xs.push((top, top_marked)); + break; + } + } + + // push our accumulated max-height right hand node back on to the stack. + xs.push((cur, contains_marked)); + + // 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( + None, + |acc: Option>, (next_tree, next_marked)| { + if let Some(mut prev_tree) = acc { + // add nil branches to build up the left tree until we can merge it + // with the right + while prev_tree.root_addr.level() < next_tree.root_addr.level() { + let sibling_addr = prev_tree.root_addr.sibling(); + contains_marked = contains_marked || next_marked; + incomplete.push(IncompleteAt { + address: sibling_addr, + required_for_witness: next_marked, + }); + prev_tree = unite( + LocatedTree { + root_addr: sibling_addr, + root: Tree(Node::Nil), + }, + prev_tree, + prune_below, + ); + } + + // at this point, prev_tree.level == next_tree.level + Some(unite(prev_tree, next_tree, prune_below)) + } else { + Some(next_tree) + } + }, + ); + + res_tree.map(|t| (t, contains_marked, incomplete)) + } else { + None + } + } + + // A stack of complete subtrees to be inserted as descendants into the subtree labeled + // with the addresses at which they will be inserted, along with their root hashes. + let mut fragments: Vec<(Self, bool)> = vec![]; + let mut position = position_range.start; + let mut checkpoints: BTreeMap = BTreeMap::new(); + while position < position_range.end { + if let Some((value, retention)) = values.next() { + if let Retention::Checkpoint { id, .. } = &retention { + checkpoints.insert(id.clone(), position); + } + + let rflags = RetentionFlags::from(retention); + let mut subtree = LocatedTree { + root_addr: Address::from(position), + root: Tree(Node::Leaf { + value: (value.clone(), rflags), + }), + }; + + if position.is_odd() { + // At odd positions, we are completing a subtree and so we unite 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); + } else { + // this is not a sibling node, so we push it back on to the stack + // and are done + fragments.push((potential_sibling, marked)); + break; + } + } + } + + fragments.push((subtree, rflags.is_marked())); + position += 1; + } else { + break; + } + } + + build_minimal_tree(fragments, prune_below).map( + |(to_insert, contains_marked, incomplete)| BatchInsertionResult { + subtree: to_insert, + contains_marked, + incomplete, + max_insert_position: Some(position - 1), + checkpoints, + remainder: values, + }, + ) + } + + /// Put a range of values into the subtree by consuming the given iterator, starting at the + /// specified position. + /// + /// The start position must exist within the position range of this subtree. If the position at + /// the end of the iterator is outside of the subtree's range, the unconsumed part of the + /// iterator will be returned as part of the result. + /// + /// Returns `Ok(None)` if the provided iterator is empty, `Ok(Some)` if + /// values were successfully inserted, or an error if the start position provided is outside + /// of this tree's position range or if a conflict with an existing subtree root is detected. + pub fn batch_insert)>, E>( + &self, + start: Position, + values: I, + ) -> Result>, InsertionError> { + let subtree_range = self.root_addr.position_range(); + let contains_start = subtree_range.contains(&start); + if contains_start { + let position_range = Range { + start, + end: subtree_range.end, + }; + Self::from_iter(position_range, self.root_addr.level(), values) + .map(|mut res| { + let (subtree, mut incomplete) = self + .clone() + .insert_subtree(res.subtree, res.contains_marked)?; + res.subtree = subtree; + res.incomplete.append(&mut incomplete); + Ok(res) + }) + .transpose() + } else { + Err(InsertionError::OutOfRange(subtree_range)) + } + } + + /// Clears the specified retention flags at all positions specified, pruning any branches + /// that no longer need to be retained. + pub fn clear_flags(&self, to_clear: BTreeMap) -> Self { + fn go( + to_clear: &[(Position, RetentionFlags)], + root_addr: Address, + root: &PrunableTree, + ) -> PrunableTree { + if to_clear.is_empty() { + // nothing to do, so we just return the root + root.clone() + } else { + match root { + Tree(Node::Parent { ann, left, right }) => { + let (l_addr, r_addr) = root_addr.children().unwrap(); + + let p = to_clear.partition_point(|(p, _)| p < &l_addr.position_range_end()); + Tree::unite( + l_addr.level(), + ann.clone(), + go(&to_clear[0..p], l_addr, left), + go(&to_clear[p..], r_addr, right), + ) + } + Tree(Node::Leaf { value: (h, r) }) => { + // When we reach a leaf, we should be down to just a single position + // which should correspond to the last level-0 child of the address's + // subtree range; if it's a checkpoint this will always be the case for + // a partially-pruned branch, and if it's a marked node then it will + // be a level-0 leaf. + match to_clear { + [(pos, flags)] => { + assert_eq!(*pos, root_addr.max_position()); + Tree(Node::Leaf { + value: (h.clone(), *r & !*flags), + }) + } + _ => { + panic!("Tree state inconsistent with checkpoints."); + } + } + } + Tree(Node::Nil) => Tree(Node::Nil), + } + } + } + + let to_clear = to_clear.into_iter().collect::>(); + Self { + root_addr: self.root_addr, + root: go(&to_clear, self.root_addr, &self.root), + } + } +} + // We need an applicative functor for Result for this function so that we can correctly // accumulate errors, but we don't have one so we just write a special- cased version here. fn accumulate_result_with( @@ -750,8 +1498,11 @@ pub mod testing { #[cfg(test)] mod tests { - use crate::{LocatedTree, Node, PrunableTree, Tree, EPHEMERAL, MARKED}; - use incrementalmerkletree::{Address, Level, Position}; + use crate::{ + LocatedPrunableTree, LocatedTree, Node, PrunableTree, QueryError, Tree, EPHEMERAL, MARKED, + }; + use core::convert::Infallible; + use incrementalmerkletree::{Address, Level, Position, Retention}; use std::collections::BTreeSet; use std::rc::Rc; @@ -938,4 +1689,126 @@ mod tests { ] ); } + + #[test] + fn located_prunable_tree_insert() { + let tree = LocatedPrunableTree::empty(Address::from_parts(Level::from(2), 0)); + let (base, _, _) = tree + .append::<(), Infallible>("a".to_string(), Retention::Ephemeral) + .unwrap(); + assert_eq!(base.right_filled_root(), Ok("a___".to_string())); + + // Perform an in-order insertion. + let (in_order, pos, _) = base + .append::<(), Infallible>("b".to_string(), Retention::Ephemeral) + .unwrap(); + assert_eq!(pos, 1.into()); + assert_eq!(in_order.right_filled_root(), Ok("ab__".to_string())); + + // On the same tree, perform an out-of-order insertion. + let out_of_order = base + .batch_insert::<(), _, Infallible>( + Position::from(3), + vec![("d".to_string(), Retention::Ephemeral)].into_iter(), + ) + .unwrap() + .unwrap(); + assert_eq!( + out_of_order.subtree, + LocatedPrunableTree { + root_addr: Address::from_parts(2.into(), 0), + root: parent( + parent(leaf(("a".to_string(), EPHEMERAL)), nil()), + parent(nil(), leaf(("d".to_string(), EPHEMERAL))) + ) + } + ); + + let complete = out_of_order + .subtree + .batch_insert::<(), _, Infallible>( + Position::from(1), + vec![ + ("b".to_string(), Retention::Ephemeral), + ("c".to_string(), Retention::Ephemeral), + ] + .into_iter(), + ) + .unwrap() + .unwrap(); + assert_eq!(complete.subtree.right_filled_root(), Ok("abcd".to_string())); + } + + #[test] + fn located_prunable_tree_insert_subtree() { + let t: LocatedPrunableTree = LocatedTree { + root_addr: Address::from_parts(3.into(), 1), + root: parent( + leaf(("abcd".to_string(), EPHEMERAL)), + parent(nil(), leaf(("gh".to_string(), EPHEMERAL))), + ), + }; + + assert_eq!( + t.insert_subtree::( + LocatedTree { + root_addr: Address::from_parts(1.into(), 6), + root: parent(leaf(("e".to_string(), MARKED)), nil()) + }, + true + ), + Ok(( + LocatedTree { + root_addr: Address::from_parts(3.into(), 1), + root: parent( + leaf(("abcd".to_string(), EPHEMERAL)), + parent( + parent(leaf(("e".to_string(), MARKED)), nil()), + leaf(("gh".to_string(), EPHEMERAL)) + ) + ) + }, + vec![] + )) + ); + } + + #[test] + fn located_prunable_tree_witness() { + let t: LocatedPrunableTree = LocatedTree { + root_addr: Address::from_parts(3.into(), 0), + root: parent( + leaf(("abcd".to_string(), EPHEMERAL)), + parent( + parent( + leaf(("e".to_string(), MARKED)), + leaf(("f".to_string(), EPHEMERAL)), + ), + leaf(("gh".to_string(), EPHEMERAL)), + ), + ), + }; + + assert_eq!( + t.witness(4.into(), 8.into()), + Ok(vec!["f", "gh", "abcd"] + .into_iter() + .map(|s| s.to_string()) + .collect()) + ); + assert_eq!( + t.witness(4.into(), 6.into()), + Ok(vec!["f", "__", "abcd"] + .into_iter() + .map(|s| s.to_string()) + .collect()) + ); + assert_eq!( + t.witness(4.into(), 7.into()), + Err(QueryError::TreeIncomplete(vec![Address::from_parts( + 1.into(), + 3 + )])) + ); + } }