Add types & operations for individual shards.
This adds the `LocatedPrunableTree` type, which provides the complete set of operations for individual shards within a larger tree.
This commit is contained in:
parent
34f6bd7ce5
commit
dc5a3ed0e7
|
@ -1,7 +1,7 @@
|
||||||
use core::fmt::Debug;
|
use core::fmt::Debug;
|
||||||
use core::ops::{BitAnd, BitOr, Deref, Not};
|
use core::ops::{BitAnd, BitOr, Deref, Not, Range};
|
||||||
use either::Either;
|
use either::Either;
|
||||||
use std::collections::BTreeSet;
|
use std::collections::{BTreeMap, BTreeSet};
|
||||||
use std::rc::Rc;
|
use std::rc::Rc;
|
||||||
|
|
||||||
use incrementalmerkletree::{Address, Hashable, Level, Position, Retention};
|
use incrementalmerkletree::{Address, Hashable, Level, Position, Retention};
|
||||||
|
@ -688,6 +688,754 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type LocatedPrunableTree<H> = LocatedTree<Option<Rc<H>>, (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<H, C: Ord, I: Iterator<Item = (H, Retention<C>)>> {
|
||||||
|
/// The updated tree after all insertions have been performed.
|
||||||
|
pub subtree: LocatedPrunableTree<H>,
|
||||||
|
/// 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<IncompleteAt>,
|
||||||
|
/// The maximum position at which a leaf was inserted.
|
||||||
|
pub max_insert_position: Option<Position>,
|
||||||
|
/// The positions of all leaves with [`CHECKPOINT`] retention that were inserted.
|
||||||
|
pub checkpoints: BTreeMap<C, Position>,
|
||||||
|
/// 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<S> {
|
||||||
|
/// 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<Position>),
|
||||||
|
/// 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<Address>),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Operations on [`LocatedTree`]s that are annotated with Merkle hashes.
|
||||||
|
impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
|
||||||
|
/// 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<H, Vec<Address>> {
|
||||||
|
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<H, Vec<Address>> {
|
||||||
|
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<Position> {
|
||||||
|
fn go<H: Hashable + Clone + PartialEq>(
|
||||||
|
root_addr: Address,
|
||||||
|
root: &PrunableTree<H>,
|
||||||
|
acc: &mut BTreeSet<Position>,
|
||||||
|
) {
|
||||||
|
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<Vec<H>, QueryError> {
|
||||||
|
// traverse down to the desired leaf position, and then construct
|
||||||
|
// the authentication path on the way back up.
|
||||||
|
fn go<H: Hashable + Clone + PartialEq>(
|
||||||
|
root: &PrunableTree<H>,
|
||||||
|
root_addr: Address,
|
||||||
|
position: Position,
|
||||||
|
truncate_at: Position,
|
||||||
|
) -> Result<Vec<H>, Vec<Address>> {
|
||||||
|
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<Self> {
|
||||||
|
fn go<H: Hashable + Clone + PartialEq>(
|
||||||
|
position: Position,
|
||||||
|
root_addr: Address,
|
||||||
|
root: &PrunableTree<H>,
|
||||||
|
) -> Option<PrunableTree<H>> {
|
||||||
|
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<E>(
|
||||||
|
&self,
|
||||||
|
subtree: Self,
|
||||||
|
contains_marked: bool,
|
||||||
|
) -> Result<(Self, Vec<IncompleteAt>), InsertionError<E>> {
|
||||||
|
// 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<H: Hashable + Clone + PartialEq, E>(
|
||||||
|
root_addr: Address,
|
||||||
|
into: &PrunableTree<H>,
|
||||||
|
subtree: LocatedPrunableTree<H>,
|
||||||
|
is_complete: bool,
|
||||||
|
contains_marked: bool,
|
||||||
|
) -> Result<(PrunableTree<H>, Vec<IncompleteAt>), InsertionError<E>> {
|
||||||
|
// 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<Rc<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,
|
||||||
|
});
|
||||||
|
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<C: Clone + Ord, E>(
|
||||||
|
&self,
|
||||||
|
value: H,
|
||||||
|
retention: Retention<C>,
|
||||||
|
) -> Result<(Self, Position, Option<C>), InsertionError<E>> {
|
||||||
|
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<C: Clone + Ord, I: Iterator<Item = (H, Retention<C>)>, E>(
|
||||||
|
&self,
|
||||||
|
values: I,
|
||||||
|
) -> Result<Option<BatchInsertionResult<H, C, I>>, InsertionError<E>> {
|
||||||
|
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<C: Clone + Ord, I: Iterator<Item = (H, Retention<C>)>>(
|
||||||
|
position_range: Range<Position>,
|
||||||
|
prune_below: Level,
|
||||||
|
mut values: I,
|
||||||
|
) -> Option<BatchInsertionResult<H, C, I>> {
|
||||||
|
// 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<H: Hashable + Clone + PartialEq>(
|
||||||
|
lroot: LocatedPrunableTree<H>,
|
||||||
|
rroot: LocatedPrunableTree<H>,
|
||||||
|
prune_below: Level,
|
||||||
|
) -> LocatedTree<Option<Rc<H>>, (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<H: Hashable + Clone + PartialEq>(
|
||||||
|
mut xs: Vec<(LocatedPrunableTree<H>, bool)>,
|
||||||
|
prune_below: Level,
|
||||||
|
) -> Option<(LocatedPrunableTree<H>, bool, Vec<IncompleteAt>)> {
|
||||||
|
// 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<LocatedPrunableTree<H>>, (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<C, Position> = 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<BatchInsertionResult>)` 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<C: Clone + Ord, I: Iterator<Item = (H, Retention<C>)>, E>(
|
||||||
|
&self,
|
||||||
|
start: Position,
|
||||||
|
values: I,
|
||||||
|
) -> Result<Option<BatchInsertionResult<H, C, I>>, InsertionError<E>> {
|
||||||
|
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<Position, RetentionFlags>) -> Self {
|
||||||
|
fn go<H: Hashable + Clone + PartialEq>(
|
||||||
|
to_clear: &[(Position, RetentionFlags)],
|
||||||
|
root_addr: Address,
|
||||||
|
root: &PrunableTree<H>,
|
||||||
|
) -> PrunableTree<H> {
|
||||||
|
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::<Vec<_>>();
|
||||||
|
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
|
// 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.
|
// accumulate errors, but we don't have one so we just write a special- cased version here.
|
||||||
fn accumulate_result_with<A, B, C>(
|
fn accumulate_result_with<A, B, C>(
|
||||||
|
@ -750,8 +1498,11 @@ pub mod testing {
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use crate::{LocatedTree, Node, PrunableTree, Tree, EPHEMERAL, MARKED};
|
use crate::{
|
||||||
use incrementalmerkletree::{Address, Level, Position};
|
LocatedPrunableTree, LocatedTree, Node, PrunableTree, QueryError, Tree, EPHEMERAL, MARKED,
|
||||||
|
};
|
||||||
|
use core::convert::Infallible;
|
||||||
|
use incrementalmerkletree::{Address, Level, Position, Retention};
|
||||||
use std::collections::BTreeSet;
|
use std::collections::BTreeSet;
|
||||||
use std::rc::Rc;
|
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<String> = 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::<Infallible>(
|
||||||
|
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<String> = 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
|
||||||
|
)]))
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue