1312 lines
55 KiB
Rust
1312 lines
55 KiB
Rust
//! Helpers for working with trees that support pruning unneeded leaves and branches.
|
|
|
|
use std::collections::{BTreeMap, BTreeSet};
|
|
use std::fmt;
|
|
use std::sync::Arc;
|
|
|
|
use bitflags::bitflags;
|
|
use incrementalmerkletree::Marking;
|
|
use incrementalmerkletree::{
|
|
frontier::NonEmptyFrontier, Address, Hashable, Level, Position, Retention,
|
|
};
|
|
use tracing::{trace, warn};
|
|
|
|
use crate::error::{InsertionError, QueryError};
|
|
use crate::{LocatedTree, Node, Tree};
|
|
|
|
bitflags! {
|
|
/// Flags storing the [`Retention`] state of a leaf.
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub struct RetentionFlags: u8 {
|
|
/// An leaf with `EPHEMERAL` retention can be pruned as soon as we are certain that it is not part
|
|
/// of the witness for a leaf with [`CHECKPOINT`] or [`MARKED`] retention.
|
|
///
|
|
/// [`CHECKPOINT`]: RetentionFlags::CHECKPOINT
|
|
/// [`MARKED`]: RetentionFlags::MARKED
|
|
const EPHEMERAL = 0b00000000;
|
|
|
|
/// A leaf with `CHECKPOINT` retention can be pruned when there are more than `max_checkpoints`
|
|
/// additional checkpoint leaves, if it is not also a marked leaf.
|
|
const CHECKPOINT = 0b00000001;
|
|
|
|
/// A leaf with `MARKED` retention can be pruned only as a consequence of an explicit deletion
|
|
/// action.
|
|
const MARKED = 0b00000010;
|
|
|
|
/// A leaf with `REFERENCE` retention will not be considered prunable until the `REFERENCE`
|
|
/// flag is removed from the leaf. The `REFERENCE` flag will be removed at any point that
|
|
/// the leaf is overwritten without `REFERENCE` retention, and `REFERENCE` retention cannot
|
|
/// be added to an existing leaf.
|
|
const REFERENCE = 0b00000100;
|
|
}
|
|
}
|
|
|
|
impl RetentionFlags {
|
|
pub fn is_checkpoint(&self) -> bool {
|
|
(*self & RetentionFlags::CHECKPOINT) == RetentionFlags::CHECKPOINT
|
|
}
|
|
|
|
pub fn is_marked(&self) -> bool {
|
|
(*self & RetentionFlags::MARKED) == RetentionFlags::MARKED
|
|
}
|
|
}
|
|
|
|
impl<'a, C> From<&'a Retention<C>> for RetentionFlags {
|
|
fn from(retention: &'a Retention<C>) -> Self {
|
|
match retention {
|
|
Retention::Ephemeral => RetentionFlags::EPHEMERAL,
|
|
Retention::Checkpoint {
|
|
marking: Marking::Marked,
|
|
..
|
|
} => RetentionFlags::CHECKPOINT | RetentionFlags::MARKED,
|
|
Retention::Checkpoint {
|
|
marking: Marking::Reference,
|
|
..
|
|
} => RetentionFlags::CHECKPOINT | RetentionFlags::REFERENCE,
|
|
Retention::Checkpoint { .. } => RetentionFlags::CHECKPOINT,
|
|
Retention::Marked => RetentionFlags::MARKED,
|
|
Retention::Reference => RetentionFlags::REFERENCE,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<C> From<Retention<C>> for RetentionFlags {
|
|
fn from(retention: Retention<C>) -> Self {
|
|
RetentionFlags::from(&retention)
|
|
}
|
|
}
|
|
|
|
/// A [`Tree`] annotated with Merkle hashes.
|
|
pub type PrunableTree<H> = Tree<Option<Arc<H>>, (H, RetentionFlags)>;
|
|
|
|
/// Errors that can occur when merging trees.
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub enum MergeError {
|
|
Conflict(Address),
|
|
TreeMalformed(Address),
|
|
}
|
|
|
|
impl From<MergeError> for InsertionError {
|
|
fn from(value: MergeError) -> Self {
|
|
match value {
|
|
MergeError::Conflict(addr) => InsertionError::Conflict(addr),
|
|
MergeError::TreeMalformed(addr) => InsertionError::InputMalformed(addr),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl fmt::Display for MergeError {
|
|
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
|
match &self {
|
|
MergeError::Conflict(addr) => write!(
|
|
f,
|
|
"Inserted root conflicts with existing root at address {:?}",
|
|
addr
|
|
),
|
|
MergeError::TreeMalformed(addr) => {
|
|
write!(f, "Merge input malformed at address {:?}", addr)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
|
|
/// Returns the the value if this is a leaf.
|
|
pub fn leaf_value(&self) -> Option<&H> {
|
|
self.0.leaf_value().map(|(h, _)| h)
|
|
}
|
|
|
|
/// Returns the cached root value with which the tree has been annotated for this node if it is
|
|
/// available, otherwise return the value if this is a leaf.
|
|
pub fn node_value(&self) -> Option<&H> {
|
|
self.0.annotation().map_or_else(
|
|
|| self.leaf_value(),
|
|
|rc_opt| rc_opt.as_ref().map(|rc| rc.as_ref()),
|
|
)
|
|
}
|
|
|
|
/// Returns whether or not this tree is a leaf with `Marked` retention.
|
|
pub fn is_marked_leaf(&self) -> bool {
|
|
self.0
|
|
.leaf_value()
|
|
.map_or(false, |(_, retention)| retention.is_marked())
|
|
}
|
|
|
|
/// Returns `true` if it is possible to compute or retrieve the Merkle root of this
|
|
/// tree.
|
|
pub fn has_computable_root(&self) -> bool {
|
|
match &self.0 {
|
|
Node::Parent { ann, left, right } => {
|
|
ann.is_some()
|
|
|| (left.as_ref().has_computable_root() && right.as_ref().has_computable_root())
|
|
}
|
|
Node::Leaf { .. } => true,
|
|
Node::Nil => false,
|
|
}
|
|
}
|
|
|
|
/// Returns `true` if no additional leaves can be appended to the right hand side of this tree
|
|
/// without over-filling it given its current depth.
|
|
pub fn is_full(&self) -> bool {
|
|
match &self.0 {
|
|
Node::Nil => false,
|
|
Node::Leaf { .. } => true,
|
|
Node::Parent { ann, right, .. } => ann.is_some() || right.is_full(),
|
|
}
|
|
}
|
|
|
|
/// Determines whether a tree has any [`Retention::Marked`] nodes.
|
|
pub fn contains_marked(&self) -> bool {
|
|
match &self.0 {
|
|
Node::Parent { left, right, .. } => left.contains_marked() || right.contains_marked(),
|
|
Node::Leaf { value: (_, r) } => r.is_marked(),
|
|
Node::Nil => false,
|
|
}
|
|
}
|
|
|
|
/// Returns the Merkle root of this tree, given the address of the root node, or
|
|
/// a vector of the addresses of `Nil` nodes that inhibited the computation of
|
|
/// such a root.
|
|
///
|
|
/// # Parameters:
|
|
/// * `truncate_at` An inclusive lower bound on positions in the tree beyond which all leaf
|
|
/// values will be treated as `Nil`.
|
|
pub fn root_hash(&self, root_addr: Address, truncate_at: Position) -> Result<H, Vec<Address>> {
|
|
if truncate_at <= root_addr.position_range_start() {
|
|
// we are in the part of the tree where we're generating empty roots,
|
|
// so no need to inspect the tree
|
|
Ok(H::empty_root(root_addr.level()))
|
|
} else {
|
|
match &self.0 {
|
|
Node::Parent { ann, left, right } => ann
|
|
.as_ref()
|
|
.filter(|_| truncate_at >= root_addr.position_range_end())
|
|
.map_or_else(
|
|
|| {
|
|
// Compute the roots of the left and right children and hash them
|
|
// together.
|
|
let (l_addr, r_addr) = root_addr
|
|
.children()
|
|
.expect("The root address of a parent node must have children.");
|
|
accumulate_result_with(
|
|
left.root_hash(l_addr, truncate_at),
|
|
right.root_hash(r_addr, truncate_at),
|
|
|left_root, right_root| {
|
|
H::combine(l_addr.level(), &left_root, &right_root)
|
|
},
|
|
)
|
|
},
|
|
|rc| {
|
|
// Since we have an annotation on the root, and we are not truncating
|
|
// within this subtree, we can just use the cached value.
|
|
Ok(rc.as_ref().clone())
|
|
},
|
|
),
|
|
Node::Leaf { value } => {
|
|
if truncate_at >= root_addr.position_range_end() {
|
|
// no truncation of this leaf is necessary, just use it
|
|
Ok(value.0.clone())
|
|
} else {
|
|
// we have a leaf value that is a subtree root created by hashing together
|
|
// the roots of child subtrees, but truncation would require that that leaf
|
|
// value be "split" into its constituent parts, which we can't do so we
|
|
// return an error
|
|
Err(vec![root_addr])
|
|
}
|
|
}
|
|
Node::Nil => Err(vec![root_addr]),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Returns a vector of the positions of [`Node::Leaf`] values in the tree having
|
|
/// [`MARKED`](RetentionFlags::MARKED) retention.
|
|
///
|
|
/// Computing the set of marked positions requires a full traversal of the tree, and so should
|
|
/// be considered to be a somewhat expensive operation.
|
|
pub fn marked_positions(&self, root_addr: Address) -> BTreeSet<Position> {
|
|
match &self.0 {
|
|
Node::Parent { left, right, .. } => {
|
|
// We should never construct parent nodes where both children are Nil.
|
|
// While we could handle that here, if we encountered that case it would
|
|
// be indicative of a programming error elsewhere and so we assert instead.
|
|
assert!(!(left.0.is_nil() && right.0.is_nil()));
|
|
let (left_root, right_root) = root_addr
|
|
.children()
|
|
.expect("A parent node cannot appear at level 0");
|
|
|
|
let mut left_incomplete = left.marked_positions(left_root);
|
|
let mut right_incomplete = right.marked_positions(right_root);
|
|
left_incomplete.append(&mut right_incomplete);
|
|
left_incomplete
|
|
}
|
|
Node::Leaf {
|
|
value: (_, retention),
|
|
} => {
|
|
let mut result = BTreeSet::new();
|
|
if root_addr.level() == 0.into() && retention.is_marked() {
|
|
result.insert(Position::from(root_addr.index()));
|
|
}
|
|
result
|
|
}
|
|
Node::Nil => BTreeSet::new(),
|
|
}
|
|
}
|
|
|
|
/// Prunes the tree by hashing together ephemeral sibling nodes.
|
|
///
|
|
/// `level` must be the level of the root of the node being pruned.
|
|
pub fn prune(self, level: Level) -> Self {
|
|
match self {
|
|
Tree(Node::Parent { ann, left, right }) => Tree::unite(
|
|
level,
|
|
ann,
|
|
left.as_ref().clone().prune(level - 1),
|
|
right.as_ref().clone().prune(level - 1),
|
|
),
|
|
other => other,
|
|
}
|
|
}
|
|
|
|
/// Merge two subtrees having the same root address.
|
|
///
|
|
/// The merge operation is checked to be strictly additive and returns an error if merging
|
|
/// would cause information loss or if a conflict between root hashes occurs at a node. The
|
|
/// returned error contains the address of the node where such a conflict occurred.
|
|
#[tracing::instrument()]
|
|
pub fn merge_checked(self, root_addr: Address, other: Self) -> Result<Self, MergeError> {
|
|
/// Pre-condition: `root_addr` must be the address of `t0` and `t1`.
|
|
#[allow(clippy::type_complexity)]
|
|
fn go<H: Hashable + Clone + PartialEq>(
|
|
addr: Address,
|
|
t0: PrunableTree<H>,
|
|
t1: PrunableTree<H>,
|
|
) -> Result<PrunableTree<H>, MergeError> {
|
|
// Require that any roots the we compute will not be default-filled by picking
|
|
// a starting valid fill point that is outside the range of leaf positions.
|
|
let no_default_fill = addr.position_range_end();
|
|
match (t0, t1) {
|
|
(Tree(Node::Nil), other) | (other, Tree(Node::Nil)) => Ok(other),
|
|
(Tree(Node::Leaf { value: vl }), Tree(Node::Leaf { value: vr })) => {
|
|
if vl.0 == vr.0 {
|
|
// Merge the flags together.
|
|
Ok(Tree::leaf((vl.0, vl.1 | vr.1)))
|
|
} else {
|
|
trace!(left = ?vl.0, right = ?vr.0, "Merge conflict for leaves");
|
|
Err(MergeError::Conflict(addr))
|
|
}
|
|
}
|
|
(Tree(Node::Leaf { value }), parent @ Tree(Node::Parent { .. }))
|
|
| (parent @ Tree(Node::Parent { .. }), Tree(Node::Leaf { value })) => {
|
|
let parent_hash = parent.root_hash(addr, no_default_fill);
|
|
if parent_hash.iter().all(|r| r == &value.0) {
|
|
Ok(parent.reannotate_root(Some(Arc::new(value.0))))
|
|
} else {
|
|
trace!(leaf = ?value, node = ?parent_hash, "Merge conflict for leaf into node");
|
|
Err(MergeError::Conflict(addr))
|
|
}
|
|
}
|
|
(lparent, rparent) => {
|
|
let lroot = lparent.root_hash(addr, no_default_fill).ok();
|
|
let rroot = rparent.root_hash(addr, no_default_fill).ok();
|
|
// If both parents share the same root hash (or if one of them is absent),
|
|
// they can be merged
|
|
if lroot.iter().zip(&rroot).all(|(l, r)| l == r) {
|
|
// using `if let` here to bind variables; we need to borrow the trees for
|
|
// root hash calculation but binding the children of the parent node
|
|
// interferes with binding a reference to the parent.
|
|
if let (
|
|
Tree(Node::Parent {
|
|
ann: lann,
|
|
left: ll,
|
|
right: lr,
|
|
}),
|
|
Tree(Node::Parent {
|
|
ann: rann,
|
|
left: rl,
|
|
right: rr,
|
|
}),
|
|
) = (lparent, rparent)
|
|
{
|
|
let (l_addr, r_addr) =
|
|
addr.children().ok_or(MergeError::TreeMalformed(addr))?;
|
|
Ok(Tree::unite(
|
|
addr.level() - 1,
|
|
lann.or(rann),
|
|
go(l_addr, ll.as_ref().clone(), rl.as_ref().clone())?,
|
|
go(r_addr, lr.as_ref().clone(), rr.as_ref().clone())?,
|
|
))
|
|
} else {
|
|
unreachable!()
|
|
}
|
|
} else {
|
|
trace!(left = ?lroot, right = ?rroot, "Merge conflict for nodes");
|
|
Err(MergeError::Conflict(addr))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
go(root_addr, self, other)
|
|
}
|
|
|
|
/// Unite two nodes by either constructing a new parent node, or, if both nodes are ephemeral
|
|
/// leaves or Nil, constructing a replacement parent by hashing leaf values together (or a
|
|
/// replacement `Nil` value).
|
|
///
|
|
/// `level` must be the level of the two nodes that are being joined.
|
|
pub(crate) fn unite(level: Level, ann: Option<Arc<H>>, left: Self, right: Self) -> Self {
|
|
match (left, right) {
|
|
(Tree(Node::Nil), Tree(Node::Nil)) if ann.is_none() => Tree::empty(),
|
|
(Tree(Node::Leaf { value: lv }), Tree(Node::Leaf { value: rv }))
|
|
// we can prune right-hand leaves that are not marked or reference leaves; if a
|
|
// leaf is a checkpoint then that information will be propagated to the replacement
|
|
// leaf
|
|
if lv.1 == RetentionFlags::EPHEMERAL &&
|
|
(rv.1 & (RetentionFlags::MARKED | RetentionFlags::REFERENCE)) == RetentionFlags::EPHEMERAL =>
|
|
{
|
|
Tree::leaf((H::combine(level, &lv.0, &rv.0), rv.1))
|
|
}
|
|
(left, right) => Tree::parent(ann, left, right),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A [`LocatedTree`] annotated with Merkle hashes.
|
|
pub type LocatedPrunableTree<H> = LocatedTree<Option<Arc<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.
|
|
///
|
|
/// [`MARKED`]: RetentionFlags::MARKED
|
|
pub required_for_witness: bool,
|
|
}
|
|
|
|
impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
|
|
/// Returns the maximum position at which a non-`Nil` leaf has been observed in the tree.
|
|
///
|
|
/// Note that no actual leaf value may exist at this position, as it may have previously been
|
|
/// pruned.
|
|
pub fn max_position(&self) -> Option<Position> {
|
|
/// Pre-condition: `addr` must be the address of `root`.
|
|
fn go<H>(
|
|
addr: Address,
|
|
root: &Tree<Option<Arc<H>>, (H, RetentionFlags)>,
|
|
) -> Option<Position> {
|
|
match &root.0 {
|
|
Node::Nil => None,
|
|
Node::Leaf { .. } => Some(addr.position_range_end() - 1),
|
|
Node::Parent { ann, left, right } => {
|
|
if ann.is_some() {
|
|
Some(addr.max_position())
|
|
} else {
|
|
let (l_addr, r_addr) = addr
|
|
.children()
|
|
.expect("has children because we checked `root` is a parent");
|
|
go(r_addr, right.as_ref()).or_else(|| go(l_addr, left.as_ref()))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
go(self.root_addr, &self.root)
|
|
}
|
|
|
|
/// 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>> {
|
|
let truncate_at = self
|
|
.max_position()
|
|
.map_or_else(|| self.root_addr.position_range_start(), |pos| pos + 1);
|
|
|
|
self.root_hash(truncate_at)
|
|
}
|
|
|
|
/// Returns the positions of marked leaves in the tree.
|
|
pub fn marked_positions(&self) -> BTreeSet<Position> {
|
|
/// Pre-condition: `root_addr` must be the address of `root`.
|
|
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()
|
|
.expect("has children because we checked `root` is a parent");
|
|
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 roots
|
|
/// corresponding to later positions will be filled by the [`Hashable::empty_root`]
|
|
/// implementation for `H`.
|
|
///
|
|
/// 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.
|
|
//
|
|
/// Pre-condition: `root_addr` must be the address of `root`.
|
|
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()
|
|
.expect("has children because we checked `root` is a parent");
|
|
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_right_child() {
|
|
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 we are truncating at
|
|
// a position to the left of the current position
|
|
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. Returns the truncated tree if a leaf or
|
|
/// subtree root with the specified position as its maximum position exists, or `None`
|
|
/// otherwise.
|
|
pub fn truncate_to_position(&self, position: Position) -> Option<Self> {
|
|
/// Pre-condition: `root_addr` must be the address of `root`.
|
|
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()
|
|
.expect("has children because we checked `root` is a parent");
|
|
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::empty())
|
|
})
|
|
} 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<IncompleteAt>), 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.
|
|
///
|
|
/// Pre-condition: `root_addr` must be the address of `into`.
|
|
#[allow(clippy::type_complexity)]
|
|
fn go<H: Hashable + Clone + PartialEq>(
|
|
root_addr: Address,
|
|
into: &PrunableTree<H>,
|
|
subtree: LocatedPrunableTree<H>,
|
|
contains_marked: bool,
|
|
) -> Result<(PrunableTree<H>, Vec<IncompleteAt>), InsertionError> {
|
|
// An empty tree cannot replace any other type of tree.
|
|
if subtree.root().is_nil() {
|
|
Ok((into.clone(), vec![]))
|
|
} else {
|
|
// In the case that we are replacing a node entirely, we need to extend the
|
|
// subtree up to the level of the node being replaced, adding Nil siblings
|
|
// and recording the presence of those incomplete nodes when necessary
|
|
let replacement = |ann: Option<Arc<H>>, mut node: LocatedPrunableTree<H>| {
|
|
// construct the replacement node bottom-up
|
|
let mut incomplete = vec![];
|
|
while node.root_addr.level() < root_addr.level() {
|
|
incomplete.push(IncompleteAt {
|
|
address: node.root_addr.sibling(),
|
|
required_for_witness: contains_marked,
|
|
});
|
|
let full = node.root;
|
|
node = LocatedTree {
|
|
root_addr: node.root_addr.parent(),
|
|
root: if node.root_addr.is_right_child() {
|
|
Tree::parent(None, Tree::empty(), full)
|
|
} else {
|
|
Tree::parent(None, full, Tree::empty())
|
|
},
|
|
};
|
|
}
|
|
(node.root.reannotate_root(ann), incomplete)
|
|
};
|
|
|
|
match into {
|
|
Tree(Node::Nil) => Ok(replacement(None, subtree)),
|
|
Tree(Node::Leaf {
|
|
value: (value, retention),
|
|
}) => {
|
|
if root_addr == subtree.root_addr {
|
|
// The current leaf is at the location we wish to transplant the root
|
|
// of the subtree being inserted, so we either replace the leaf
|
|
// entirely with the subtree, or reannotate the root so as to avoid
|
|
// discarding the existing leaf value.
|
|
|
|
if subtree.root.has_computable_root() {
|
|
Ok((
|
|
if subtree.root.is_leaf() {
|
|
// When replacing a leaf with a leaf, `REFERENCE` retention
|
|
// will be discarded unless both leaves have `REFERENCE`
|
|
// retention.
|
|
subtree
|
|
.root
|
|
.try_map::<(H, RetentionFlags), InsertionError, _>(
|
|
&|(v0, ret0)| {
|
|
if v0 == value {
|
|
let retention_result: RetentionFlags =
|
|
((*retention | *ret0)
|
|
- RetentionFlags::REFERENCE)
|
|
| (RetentionFlags::REFERENCE
|
|
& *retention
|
|
& *ret0);
|
|
Ok((value.clone(), retention_result))
|
|
} else {
|
|
Err(InsertionError::Conflict(root_addr))
|
|
}
|
|
},
|
|
)?
|
|
} else {
|
|
// It is safe to replace the existing root unannotated, because we
|
|
// can always recompute the root from the subtree.
|
|
subtree.root
|
|
},
|
|
vec![],
|
|
))
|
|
} else if subtree.root.node_value().iter().all(|v| v == &value) {
|
|
Ok((
|
|
// at this point we statically know the root to be a parent
|
|
subtree.root.reannotate_root(Some(Arc::new(value.clone()))),
|
|
vec![],
|
|
))
|
|
} else {
|
|
warn!(
|
|
cur_root = ?value,
|
|
new_root = ?subtree.root.node_value(),
|
|
"Insertion conflict",
|
|
);
|
|
Err(InsertionError::Conflict(root_addr))
|
|
}
|
|
} else {
|
|
Ok(replacement(Some(Arc::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::from)
|
|
.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()
|
|
.expect("has children because we checked `into` is a parent");
|
|
if l_addr.contains(&subtree.root_addr) {
|
|
let (new_left, incomplete) =
|
|
go(l_addr, left.as_ref(), subtree, contains_marked)?;
|
|
Ok((
|
|
Tree::unite(
|
|
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, contains_marked)?;
|
|
Ok((
|
|
Tree::unite(
|
|
root_addr.level() - 1,
|
|
ann.clone(),
|
|
left.as_ref().clone(),
|
|
new_right,
|
|
),
|
|
incomplete,
|
|
))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let max_position = self.max_position();
|
|
trace!(
|
|
max_position = ?max_position,
|
|
tree = ?self,
|
|
to_insert = ?subtree,
|
|
"Current shard"
|
|
);
|
|
let LocatedTree { root_addr, root } = self;
|
|
if root_addr.contains(&subtree.root_addr) {
|
|
go(*root_addr, root, subtree, contains_marked).map(|(root, incomplete)| {
|
|
let new_tree = LocatedTree {
|
|
root_addr: *root_addr,
|
|
root,
|
|
};
|
|
assert!(new_tree.max_position() >= max_position);
|
|
(new_tree, incomplete)
|
|
})
|
|
} else {
|
|
Err(InsertionError::NotContained(subtree.root_addr))
|
|
}
|
|
}
|
|
|
|
/// 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>(
|
|
&self,
|
|
value: H,
|
|
retention: Retention<C>,
|
|
) -> Result<(Self, Position, Option<C>), 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, checkpoint_id))
|
|
}
|
|
})
|
|
}
|
|
|
|
// Constructs a pair of trees that contain the leaf and ommers of the given frontier. The first
|
|
// element of the result is a tree with its root at a level less than or equal to `split_at`;
|
|
// the second element is a tree with its leaves at level `split_at` that is only returned if
|
|
// the frontier contains sufficient data to fill the first tree to the `split_at` level.
|
|
fn from_frontier<C>(
|
|
frontier: NonEmptyFrontier<H>,
|
|
leaf_retention: &Retention<C>,
|
|
split_at: Level,
|
|
) -> (Self, Option<Self>) {
|
|
let (position, leaf, ommers) = frontier.into_parts();
|
|
Self::from_frontier_parts(position, leaf, ommers.into_iter(), leaf_retention, split_at)
|
|
}
|
|
|
|
// Permits construction of a subtree from legacy `CommitmentTree` data that may
|
|
// have inaccurate position information (e.g. in the case that the tree is the
|
|
// cursor for an `IncrementalWitness`).
|
|
pub(crate) fn from_frontier_parts<C>(
|
|
position: Position,
|
|
leaf: H,
|
|
mut ommers: impl Iterator<Item = H>,
|
|
leaf_retention: &Retention<C>,
|
|
split_at: Level,
|
|
) -> (Self, Option<Self>) {
|
|
let mut addr = Address::from(position);
|
|
let mut subtree = Tree::leaf((leaf, leaf_retention.into()));
|
|
|
|
while addr.level() < split_at {
|
|
if addr.is_left_child() {
|
|
// the current address is a left child, so create a parent with
|
|
// an empty right-hand tree
|
|
subtree = Tree::parent(None, subtree, Tree::empty());
|
|
} else if let Some(left) = ommers.next() {
|
|
// the current address corresponds to a right child, so create a parent that
|
|
// takes the left sibling to that child from the ommers
|
|
subtree =
|
|
Tree::parent(None, Tree::leaf((left, RetentionFlags::EPHEMERAL)), subtree);
|
|
} else {
|
|
break;
|
|
}
|
|
|
|
addr = addr.parent();
|
|
}
|
|
|
|
let located_subtree = LocatedTree {
|
|
root_addr: addr,
|
|
root: subtree,
|
|
};
|
|
|
|
let located_supertree = if located_subtree.root_addr().level() == split_at {
|
|
let mut addr = located_subtree.root_addr();
|
|
let mut supertree = None;
|
|
for left in ommers {
|
|
// build up the left-biased tree until we get a right-hand node
|
|
while addr.is_left_child() {
|
|
supertree = supertree.map(|t| Tree::parent(None, t, Tree::empty()));
|
|
addr = addr.parent();
|
|
}
|
|
|
|
// once we have a right-hand root, add a parent with the current ommer as the
|
|
// left-hand sibling
|
|
supertree = Some(Tree::parent(
|
|
None,
|
|
Tree::leaf((left, RetentionFlags::EPHEMERAL)),
|
|
supertree.unwrap_or_else(Tree::empty),
|
|
));
|
|
addr = addr.parent();
|
|
}
|
|
|
|
supertree.map(|t| LocatedTree {
|
|
root_addr: addr,
|
|
root: t,
|
|
})
|
|
} else {
|
|
// if there were not enough ommers available from the frontier to reach the address
|
|
// of the root of this tree, there is no contribution to the cap
|
|
None
|
|
};
|
|
|
|
(located_subtree, located_supertree)
|
|
}
|
|
|
|
/// Inserts leaves and subtree roots from the provided frontier into this tree, up to the level
|
|
/// of this tree's root.
|
|
///
|
|
/// Returns the updated tree, along with a [`LocatedPrunableTree`] containing only the remainder
|
|
/// of the frontier's ommers that had addresses at levels greater than the root of this tree.
|
|
///
|
|
/// Returns an error in the following cases:
|
|
/// * the leaf node of `frontier` is at a position that is not contained within this tree's
|
|
/// position range
|
|
/// * a conflict occurs where an ommer of the frontier being inserted does not match the
|
|
/// existing value for that node
|
|
pub fn insert_frontier_nodes<C>(
|
|
&self,
|
|
frontier: NonEmptyFrontier<H>,
|
|
leaf_retention: &Retention<C>,
|
|
) -> Result<(Self, Option<Self>), InsertionError> {
|
|
let subtree_range = self.root_addr.position_range();
|
|
if subtree_range.contains(&frontier.position()) {
|
|
let leaf_is_marked = leaf_retention.is_marked();
|
|
let (subtree, supertree) =
|
|
Self::from_frontier(frontier, leaf_retention, self.root_addr.level());
|
|
|
|
let subtree = self.insert_subtree(subtree, leaf_is_marked)?.0;
|
|
|
|
Ok((subtree, supertree))
|
|
} else {
|
|
Err(InsertionError::OutOfRange(
|
|
frontier.position(),
|
|
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 {
|
|
/// Pre-condition: `root_addr` must be the address of `root`.
|
|
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.0 {
|
|
Node::Parent { ann, left, right } => {
|
|
let (l_addr, r_addr) = root_addr
|
|
.children()
|
|
.expect("has children because we checked `root` is a parent");
|
|
|
|
let p = to_clear.partition_point(|(p, _)| p < &l_addr.position_range_end());
|
|
trace!(
|
|
"Tree::unite at {:?}, partitioned: {:?} {:?}",
|
|
root_addr,
|
|
&to_clear[0..p],
|
|
&to_clear[p..],
|
|
);
|
|
Tree::unite(
|
|
l_addr.level(),
|
|
ann.clone(),
|
|
go(&to_clear[0..p], l_addr, left),
|
|
go(&to_clear[p..], r_addr, right),
|
|
)
|
|
}
|
|
Node::Leaf { value: (h, r) } => {
|
|
trace!("In {:?}, clearing {:?}", root_addr, to_clear);
|
|
// 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 {
|
|
[(_, flags)] => Tree::leaf((h.clone(), *r & !*flags)),
|
|
_ => {
|
|
panic!("Tree state inconsistent with checkpoints.");
|
|
}
|
|
}
|
|
}
|
|
Node::Nil => Tree::empty(),
|
|
}
|
|
}
|
|
}
|
|
|
|
let to_clear = to_clear.into_iter().collect::<Vec<_>>();
|
|
Self {
|
|
root_addr: self.root_addr,
|
|
root: go(&to_clear, self.root_addr, &self.root),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(crate) fn flag_positions(&self) -> BTreeMap<Position, RetentionFlags> {
|
|
fn go<H>(
|
|
root: &PrunableTree<H>,
|
|
root_addr: Address,
|
|
acc: &mut BTreeMap<Position, RetentionFlags>,
|
|
) {
|
|
match &root.0 {
|
|
Node::Parent { left, right, .. } => {
|
|
let (l_addr, r_addr) = root_addr
|
|
.children()
|
|
.expect("A parent node cannot appear at level 0");
|
|
go(left, l_addr, acc);
|
|
go(right, r_addr, acc);
|
|
}
|
|
Node::Leaf { value } if value.1 != RetentionFlags::EPHEMERAL => {
|
|
acc.insert(root_addr.max_position(), value.1);
|
|
}
|
|
_ => (),
|
|
}
|
|
}
|
|
|
|
let mut result = BTreeMap::new();
|
|
go(&self.root, self.root_addr, &mut result);
|
|
result
|
|
}
|
|
}
|
|
|
|
// 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<A, B, C>(
|
|
left: Result<A, Vec<Address>>,
|
|
right: Result<B, Vec<Address>>,
|
|
combine_success: impl FnOnce(A, B) -> C,
|
|
) -> Result<C, Vec<Address>> {
|
|
match (left, right) {
|
|
(Ok(a), Ok(b)) => Ok(combine_success(a, b)),
|
|
(Err(mut xs), Err(mut ys)) => {
|
|
xs.append(&mut ys);
|
|
Err(xs)
|
|
}
|
|
(Ok(_), Err(xs)) => Err(xs),
|
|
(Err(xs), Ok(_)) => Err(xs),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::collections::{BTreeMap, BTreeSet};
|
|
|
|
use incrementalmerkletree::{Address, Level, Position};
|
|
use proptest::proptest;
|
|
|
|
use super::{LocatedPrunableTree, PrunableTree, RetentionFlags};
|
|
use crate::{
|
|
error::{InsertionError, QueryError},
|
|
prunable::MergeError,
|
|
testing::{arb_char_str, arb_prunable_tree},
|
|
tree::{
|
|
tests::{leaf, nil, parent},
|
|
LocatedTree,
|
|
},
|
|
};
|
|
|
|
#[test]
|
|
fn root() {
|
|
let t: PrunableTree<String> = parent(
|
|
leaf(("a".to_string(), RetentionFlags::EPHEMERAL)),
|
|
leaf(("b".to_string(), RetentionFlags::EPHEMERAL)),
|
|
);
|
|
|
|
assert_eq!(
|
|
t.root_hash(Address::from_parts(Level::from(1), 0), Position::from(2)),
|
|
Ok("ab".to_string())
|
|
);
|
|
|
|
let t0 = parent(nil(), t.clone());
|
|
assert_eq!(
|
|
t0.root_hash(Address::from_parts(Level::from(2), 0), Position::from(4)),
|
|
Err(vec![Address::from_parts(Level::from(1), 0)])
|
|
);
|
|
|
|
// Check root computation with truncation
|
|
let t1 = parent(t, nil());
|
|
assert_eq!(
|
|
t1.root_hash(Address::from_parts(Level::from(2), 0), Position::from(2)),
|
|
Ok("ab__".to_string())
|
|
);
|
|
assert_eq!(
|
|
t1.root_hash(Address::from_parts(Level::from(2), 0), Position::from(3)),
|
|
Err(vec![Address::from_parts(Level::from(1), 1)])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn marked_positions() {
|
|
let t: PrunableTree<String> = parent(
|
|
leaf(("a".to_string(), RetentionFlags::EPHEMERAL)),
|
|
leaf(("b".to_string(), RetentionFlags::MARKED)),
|
|
);
|
|
assert_eq!(
|
|
t.marked_positions(Address::from_parts(Level::from(1), 0)),
|
|
BTreeSet::from([Position::from(1)])
|
|
);
|
|
|
|
let t0 = parent(t.clone(), t);
|
|
assert_eq!(
|
|
t0.marked_positions(Address::from_parts(Level::from(2), 1)),
|
|
BTreeSet::from([Position::from(5), Position::from(7)])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn prune() {
|
|
let t: PrunableTree<String> = parent(
|
|
leaf(("a".to_string(), RetentionFlags::EPHEMERAL)),
|
|
leaf(("b".to_string(), RetentionFlags::EPHEMERAL)),
|
|
);
|
|
|
|
assert_eq!(
|
|
t.clone().prune(Level::from(1)),
|
|
leaf(("ab".to_string(), RetentionFlags::EPHEMERAL))
|
|
);
|
|
|
|
let t0 = parent(leaf(("c".to_string(), RetentionFlags::MARKED)), t);
|
|
assert_eq!(
|
|
t0.prune(Level::from(2)),
|
|
parent(
|
|
leaf(("c".to_string(), RetentionFlags::MARKED)),
|
|
leaf(("ab".to_string(), RetentionFlags::EPHEMERAL))
|
|
)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn merge_checked() {
|
|
let t0: PrunableTree<String> =
|
|
parent(leaf(("a".to_string(), RetentionFlags::EPHEMERAL)), nil());
|
|
|
|
let t1: PrunableTree<String> =
|
|
parent(nil(), leaf(("b".to_string(), RetentionFlags::EPHEMERAL)));
|
|
|
|
assert_eq!(
|
|
t0.clone()
|
|
.merge_checked(Address::from_parts(1.into(), 0), t1.clone()),
|
|
Ok(leaf(("ab".to_string(), RetentionFlags::EPHEMERAL)))
|
|
);
|
|
|
|
let t2: PrunableTree<String> =
|
|
parent(leaf(("c".to_string(), RetentionFlags::EPHEMERAL)), nil());
|
|
assert_eq!(
|
|
t0.clone()
|
|
.merge_checked(Address::from_parts(1.into(), 0), t2.clone()),
|
|
Err(MergeError::Conflict(Address::from_parts(0.into(), 0)))
|
|
);
|
|
|
|
let t3: PrunableTree<String> = parent(t0, t2);
|
|
let t4: PrunableTree<String> = parent(t1.clone(), t1);
|
|
|
|
assert_eq!(
|
|
t3.merge_checked(Address::from_parts(2.into(), 0), t4),
|
|
Ok(leaf(("abcb".to_string(), RetentionFlags::EPHEMERAL)))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn merge_checked_flags() {
|
|
let t0: PrunableTree<String> = leaf(("a".to_string(), RetentionFlags::EPHEMERAL));
|
|
let t1: PrunableTree<String> = leaf(("a".to_string(), RetentionFlags::MARKED));
|
|
let t2: PrunableTree<String> = leaf(("a".to_string(), RetentionFlags::CHECKPOINT));
|
|
|
|
assert_eq!(
|
|
t0.merge_checked(Address::from_parts(1.into(), 0), t1.clone()),
|
|
Ok(t1.clone()),
|
|
);
|
|
|
|
assert_eq!(
|
|
t1.merge_checked(Address::from_parts(1.into(), 0), t2),
|
|
Ok(leaf((
|
|
"a".to_string(),
|
|
RetentionFlags::MARKED | RetentionFlags::CHECKPOINT,
|
|
))),
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn located_insert_subtree() {
|
|
let t: LocatedPrunableTree<String> = LocatedTree {
|
|
root_addr: Address::from_parts(3.into(), 1),
|
|
root: parent(
|
|
leaf(("abcd".to_string(), RetentionFlags::EPHEMERAL)),
|
|
parent(nil(), leaf(("gh".to_string(), RetentionFlags::EPHEMERAL))),
|
|
),
|
|
};
|
|
|
|
assert_eq!(
|
|
t.insert_subtree(
|
|
LocatedTree {
|
|
root_addr: Address::from_parts(1.into(), 6),
|
|
root: parent(leaf(("e".to_string(), RetentionFlags::MARKED)), nil())
|
|
},
|
|
true
|
|
),
|
|
Ok((
|
|
LocatedTree {
|
|
root_addr: Address::from_parts(3.into(), 1),
|
|
root: parent(
|
|
leaf(("abcd".to_string(), RetentionFlags::EPHEMERAL)),
|
|
parent(
|
|
parent(leaf(("e".to_string(), RetentionFlags::MARKED)), nil()),
|
|
leaf(("gh".to_string(), RetentionFlags::EPHEMERAL))
|
|
)
|
|
)
|
|
},
|
|
vec![]
|
|
))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn located_insert_subtree_prevents_leaf_overwrite_conflict() {
|
|
let t: LocatedPrunableTree<String> = LocatedTree {
|
|
root_addr: Address::from_parts(2.into(), 1),
|
|
root: parent(leaf(("a".to_string(), RetentionFlags::MARKED)), nil()),
|
|
};
|
|
|
|
let conflict_addr = Address::from_parts(1.into(), 2);
|
|
assert_eq!(
|
|
t.insert_subtree(
|
|
LocatedTree {
|
|
root_addr: conflict_addr,
|
|
root: leaf(("b".to_string(), RetentionFlags::EPHEMERAL)),
|
|
},
|
|
false,
|
|
),
|
|
Err(InsertionError::Conflict(conflict_addr)),
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn located_witness() {
|
|
let t: LocatedPrunableTree<String> = LocatedTree {
|
|
root_addr: Address::from_parts(3.into(), 0),
|
|
root: parent(
|
|
leaf(("abcd".to_string(), RetentionFlags::EPHEMERAL)),
|
|
parent(
|
|
parent(
|
|
leaf(("e".to_string(), RetentionFlags::MARKED)),
|
|
leaf(("f".to_string(), RetentionFlags::EPHEMERAL)),
|
|
),
|
|
leaf(("gh".to_string(), RetentionFlags::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
|
|
)]))
|
|
);
|
|
}
|
|
|
|
proptest! {
|
|
#[test]
|
|
fn clear_flags(
|
|
root in arb_prunable_tree(arb_char_str(), 8, 2^6)
|
|
) {
|
|
let root_addr = Address::from_parts(Level::from(7), 0);
|
|
let tree = LocatedTree::from_parts(root_addr, root).unwrap();
|
|
|
|
let (to_clear, to_retain) = tree.flag_positions().into_iter().enumerate().fold(
|
|
(BTreeMap::new(), BTreeMap::new()),
|
|
|(mut to_clear, mut to_retain), (i, (pos, flags))| {
|
|
if i % 2 == 0 {
|
|
to_clear.insert(pos, flags);
|
|
} else {
|
|
to_retain.insert(pos, flags);
|
|
}
|
|
(to_clear, to_retain)
|
|
}
|
|
);
|
|
|
|
let pre_clearing_max_position = tree.max_position();
|
|
let cleared = tree.clear_flags(to_clear);
|
|
|
|
// Clearing flags should not modify the max position of leaves represented
|
|
// in the shard.
|
|
assert!(cleared.max_position() == pre_clearing_max_position);
|
|
assert_eq!(to_retain, cleared.flag_positions());
|
|
}
|
|
}
|
|
}
|