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

Merge pull request #105 from nuttycom/limitable_pruning 

shardtree: Add `Reference` retention to make it possible to limit pruning.
This commit is contained in:
Kris Nuttycombe 2024-05-30 08:57:43 -06:00 committed by GitHub
parent 33fc3f6013 93512976dd
commit 8b4b1315d6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 612 additions and 265 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
incrementalmerkletree/CHANGELOG.md
View File

@ -7,6 +7,15 @@ and this project adheres to Rust's notion of
## Unreleased
### Added
- `incrementalmerkletree::Marking`
### Changed
- `incrementalmerkletree::Retention`
- Has added variant `Retention::Reference`
- `Retention::Checkpoint::is_marked` has been replaced by `Retention::Checkpoint::marking`
to permit checkpoints with `Reference` retention to be represented.
## [0.5.1] - 2024-03-25
### Added

54
incrementalmerkletree/src/lib.rs
View File

@ -62,35 +62,73 @@ pub mod witness;
#[cfg_attr(docsrs, doc(cfg(feature = "test-dependencies")))]
pub mod testing;
/// An enumeration of the additional marking states that can be applied
/// to leaves with [`Retention::Checkpoint`] retention.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Marking {
/// A checkpoint with `Marked` marking will have `Marked` retention after `Checkpoint`
/// retention is removed.
Marked,
/// A checkpoint with `Reference` marking will have `Reference` retention after `Checkpoint`
/// retention is removed.
Reference,
/// A checkpoint with `None` marking will have `Ephemeral` retention after `Checkpoint`
/// retention is removed.
None,
}
/// A type for metadata that is used to determine when and how a leaf can be pruned from a tree.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Retention<C> {
/// A leaf with `Ephemeral` retention will be pruned whenever its sibling is also a leaf with
/// `Ephemeral` retention.
Ephemeral,
Checkpoint { id: C, is_marked: bool },
/// A leaf with `Checkpoint` retention will have its position retained in the tree
/// during pruning, but its value may be pruned (by merging with its sibling). If
/// `Checkpoint` retention is removed from the leaf, then the retention for the leaf will
/// become either `Ephemeral`, `Marked`, or `Reference` depending upon the value of the
/// `marking` field.
Checkpoint { id: C, marking: Marking },
/// A leaf with `Marked` retention will be retained in the tree during pruning. `Marked`
/// retention may only be explicitly removed.
Marked,
/// A leaf with `Reference` retention will be retained in the tree during pruning. `Reference`
/// retention is removed whenever the associated leaf is overwritten with a tree node having
/// `Ephemeral`, `Checkpoint`, or `Marked` retention.
Reference,
}
impl<C> Retention<C> {
/// Returns whether the associated node has [`Retention::Checkpoint`] retention.
pub fn is_checkpoint(&self) -> bool {
matches!(self, Retention::Checkpoint { .. })
}
/// Returns whether the associated node has [`Retention::Marked`] retention
/// or [`Retention::Checkpoint`] retention with [`Marking::Marked`] marking.
pub fn is_marked(&self) -> bool {
match self {
Retention::Ephemeral => false,
Retention::Checkpoint { is_marked, .. } => *is_marked,
Retention::Marked => true,
}
matches!(
self,
Retention::Marked
| Retention::Checkpoint {
marking: Marking::Marked,
..
}
)
}
/// Applies the provided function to the checkpoint identifier, if any, and returns a new
/// `Retention` value having the same structure with the resulting value used as the checkpoint
/// identifier.
pub fn map<'a, D, F: Fn(&'a C) -> D>(&'a self, f: F) -> Retention<D> {
match self {
Retention::Ephemeral => Retention::Ephemeral,
Retention::Checkpoint { id, is_marked } => Retention::Checkpoint {
Retention::Checkpoint { id, marking } => Retention::Checkpoint {
id: f(id),
is_marked: *is_marked,
marking: *marking,
},
Retention::Marked => Retention::Marked,
Retention::Reference => Retention::Reference,
}
}
}

42
incrementalmerkletree/src/testing.rs
View File

@ -5,7 +5,7 @@ use core::marker::PhantomData;
use proptest::prelude::*;
use std::collections::BTreeSet;
use crate::{Hashable, Level, Position, Retention};
use crate::{Hashable, Level, Marking, Position, Retention};
pub mod complete_tree;
@ -208,10 +208,22 @@ impl<H: Hashable + Clone, C: Clone> Operation<H, C> {
}
}
/// Returns a strategy for creating a uniformly-distributed [`Marking`]
/// value.
pub fn arb_marking() -> impl Strategy<Value = Marking> {
prop_oneof![
Just(Marking::Marked),
Just(Marking::Reference),
Just(Marking::None)
]
}
/// Returns a strategy for creating a uniformly-distributed [`Retention`]
/// value.
pub fn arb_retention() -> impl Strategy<Value = Retention<()>> {
prop_oneof![
Just(Retention::Ephemeral),
any::<bool>().prop_map(|is_marked| Retention::Checkpoint { id: (), is_marked }),
arb_marking().prop_map(|marking| Retention::Checkpoint { id: (), marking }),
Just(Retention::Marked),
]
}
@ -557,7 +569,7 @@ pub fn check_root_hashes<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: F
0,
Retention::Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
);
for _ in 0..3 {
@ -735,7 +747,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
0,
Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
);
assert!(tree.rewind());
@ -769,7 +781,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
6,
Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
);
tree.assert_append(7, Ephemeral);
@ -838,7 +850,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
H::from_u64(3),
Checkpoint {
id: C::from_u64(1),
is_marked: true,
marking: Marking::Marked,
},
),
Append(H::from_u64(4), Marked),
@ -847,21 +859,21 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
H::from_u64(5),
Checkpoint {
id: C::from_u64(3),
is_marked: false,
marking: Marking::None,
},
),
Append(
H::from_u64(6),
Checkpoint {
id: C::from_u64(4),
is_marked: false,
marking: Marking::None,
},
),
Append(
H::from_u64(7),
Checkpoint {
id: C::from_u64(5),
is_marked: false,
marking: Marking::None,
},
),
Witness(3u64.into(), 5),
@ -890,7 +902,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
H::from_u64(0),
Checkpoint {
id: C::from_u64(1),
is_marked: true,
marking: Marking::Marked,
},
),
Append(H::from_u64(0), Ephemeral),
@ -900,7 +912,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
H::from_u64(0),
Checkpoint {
id: C::from_u64(2),
is_marked: false,
marking: Marking::None,
},
),
Append(H::from_u64(0), Ephemeral),
@ -939,7 +951,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
H::from_u64(0),
Checkpoint {
id: C::from_u64(4),
is_marked: false,
marking: Marking::None,
},
),
Rewind,
@ -958,14 +970,14 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
H::from_u64(0),
Checkpoint {
id: C::from_u64(1),
is_marked: true,
marking: Marking::Marked,
},
),
Append(
H::from_u64(0),
Checkpoint {
id: C::from_u64(4),
is_marked: false,
marking: Marking::None,
},
),
Witness(Position(2), 2),
@ -1034,7 +1046,7 @@ pub fn check_remove_mark<C: TestCheckpoint, T: Tree<String, C>, F: Fn(usize) ->
"a",
Retention::Checkpoint {
id: C::from_u64(1),
is_marked: true,
marking: Marking::Marked,
},
),
witness(1, 1),

7
incrementalmerkletree/src/testing/complete_tree.rs
View File

@ -2,6 +2,7 @@
use std::cmp::min;
use std::collections::{BTreeMap, BTreeSet};
use crate::Marking;
use crate::{testing::Tree, Hashable, Level, Position, Retention};
const MAX_COMPLETE_SIZE_ERROR: &str = "Positions of a `CompleteTree` must fit into the platform word size, because larger complete trees are not representable.";
@ -104,11 +105,11 @@ impl<H: Hashable, C: Clone + Ord + core::fmt::Debug, const DEPTH: u8> CompleteTr
append(&mut self.leaves, value, DEPTH)?;
self.mark();
}
Retention::Checkpoint { id, is_marked } => {
Retention::Checkpoint { id, marking } => {
let latest_checkpoint = self.checkpoints.keys().rev().next();
if Some(&id) > latest_checkpoint {
append(&mut self.leaves, value, DEPTH)?;
if is_marked {
if marking == Marking::Marked {
self.mark();
}
self.checkpoint(id, self.current_position());
@ -119,7 +120,7 @@ impl<H: Hashable, C: Clone + Ord + core::fmt::Debug, const DEPTH: u8> CompleteTr
});
}
}
Retention::Ephemeral => {
Retention::Ephemeral | Retention::Reference => {
append(&mut self.leaves, value, DEPTH)?;
}
}

18
shardtree/CHANGELOG.md
View File

@ -7,6 +7,24 @@ and this project adheres to Rust's notion of
## Unreleased
### Added
- `shardtree::tree::Tree::{is_leaf, map, try_map, empty_pruned}`
- `shardtree::tree::LocatedTree::{map, try_map}`
- `shardtree::prunable::PrunableTree::{has_computable_root}`
### Changed
- `shardtree::tree::Node` has additional variant `Node::Pruned`.
### Removed
- `shardtree::tree::Tree::is_complete` as it is no longer well-defined in the
presence of `Pruned` nodes. Use `PrunableTree::has_computable_root` to
determine whether it is possible to compute the root of a tree.
### Fixed
- Fixes an error that could occur if an inserted `Frontier` node was
interpreted as a node that had actually had its value observed as though it
had been inserted using the ordinary tree insertion methods.
## [0.3.1] - 2024-04-03
### Fixed

6
shardtree/src/legacy.rs
View File

@ -1,6 +1,8 @@
use std::fmt;
use incrementalmerkletree::{witness::IncrementalWitness, Address, Hashable, Level, Retention};
use incrementalmerkletree::{
witness::IncrementalWitness, Address, Hashable, Level, Marking, Retention,
};
use crate::{
store::ShardStore, InsertionError, LocatedPrunableTree, LocatedTree, PrunableTree,
@ -198,7 +200,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
c.ommers_iter().cloned(),
&Retention::Checkpoint {
id: checkpoint_id,
is_marked: false,
marking: Marking::None,
},
self.root_addr.level(),
)

199
shardtree/src/lib.rs
View File

@ -25,9 +25,10 @@
use core::fmt::Debug;
use either::Either;
use incrementalmerkletree::frontier::Frontier;
use incrementalmerkletree::Marking;
use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use tracing::trace;
use tracing::{debug, trace};
use incrementalmerkletree::{
frontier::NonEmptyFrontier, Address, Hashable, Level, MerklePath, Position, Retention,
@ -233,16 +234,17 @@ impl<
let (append_result, position, checkpoint_id) =
if let Some(subtree) = self.store.last_shard().map_err(ShardTreeError::Storage)? {
if subtree.root.is_complete() {
let addr = subtree.root_addr;
if addr.index() < Self::max_subtree_index() {
LocatedTree::empty(addr.next_at_level()).append(value, retention)?
} else {
return Err(InsertionError::TreeFull.into());
match subtree.max_position() {
// If the subtree is full, then construct a successor tree.
Some(pos) if pos == subtree.root_addr.max_position() => {
let addr = subtree.root_addr;
if subtree.root_addr.index() < Self::max_subtree_index() {
LocatedTree::empty(addr.next_at_level()).append(value, retention)?
} else {
return Err(InsertionError::TreeFull.into());
}
}
} else {
subtree.append(value, retention)?
_ => subtree.append(value, retention)?,
}
} else {
let root_addr = Address::from_parts(Self::subtree_level(), 0);
@ -271,6 +273,7 @@ impl<
///
/// This method may be used to add a checkpoint for the empty tree; note that
/// [`Retention::Marked`] is invalid for the empty tree.
#[tracing::instrument(skip(self, frontier, leaf_retention))]
pub fn insert_frontier(
&mut self,
frontier: Frontier<H, DEPTH>,
@ -280,18 +283,19 @@ impl<
self.insert_frontier_nodes(nonempty_frontier, leaf_retention)
} else {
match leaf_retention {
Retention::Ephemeral => Ok(()),
Retention::Ephemeral | Retention::Reference => Ok(()),
Retention::Checkpoint {
id,
is_marked: false,
marking: Marking::None | Marking::Reference,
} => self
.store
.add_checkpoint(id, Checkpoint::tree_empty())
.map_err(ShardTreeError::Storage),
Retention::Checkpoint {
is_marked: true, ..
}
| Retention::Marked => Err(ShardTreeError::Insert(
Retention::Marked
| Retention::Checkpoint {
marking: Marking::Marked,
..
} => Err(ShardTreeError::Insert(
InsertionError::MarkedRetentionInvalid,
)),
}
@ -301,6 +305,7 @@ impl<
/// Add the leaf and ommers of the provided frontier as nodes within the subtree corresponding
/// to the frontier's position, and update the cap to include the ommer nodes at levels greater
/// than or equal to the shard height.
#[tracing::instrument(skip(self, frontier, leaf_retention))]
pub fn insert_frontier_nodes(
&mut self,
frontier: NonEmptyFrontier<H>,
@ -308,15 +313,24 @@ impl<
) -> Result<(), ShardTreeError<S::Error>> {
let leaf_position = frontier.position();
let subtree_root_addr = Address::above_position(Self::subtree_level(), leaf_position);
trace!("Subtree containing nodes: {:?}", subtree_root_addr);
let (updated_subtree, supertree) = self
let current_shard = self
.store
.get_shard(subtree_root_addr)
.map_err(ShardTreeError::Storage)?
.unwrap_or_else(|| LocatedTree::empty(subtree_root_addr))
.insert_frontier_nodes(frontier, &leaf_retention)?;
.unwrap_or_else(|| LocatedTree::empty(subtree_root_addr));
trace!(
max_position = ?current_shard.max_position(),
subtree = ?current_shard,
"Current shard");
let (updated_subtree, supertree) =
current_shard.insert_frontier_nodes(frontier, &leaf_retention)?;
trace!(
max_position = ?updated_subtree.max_position(),
subtree = ?updated_subtree,
"Replacement shard"
);
self.store
.put_shard(updated_subtree)
.map_err(ShardTreeError::Storage)?;
@ -333,7 +347,7 @@ impl<
.map_err(ShardTreeError::Storage)?;
}
if let Retention::Checkpoint { id, is_marked: _ } = leaf_retention {
if let Retention::Checkpoint { id, .. } = leaf_retention {
trace!("Adding checkpoint {:?} at {:?}", id, leaf_position);
self.store
.add_checkpoint(id, Checkpoint::at_position(leaf_position))
@ -346,6 +360,7 @@ impl<
/// Insert a tree by decomposing it into its `SHARD_HEIGHT` or smaller parts (if necessary)
/// and inserting those at their appropriate locations.
#[tracing::instrument(skip(self, tree, checkpoints))]
pub fn insert_tree(
&mut self,
tree: LocatedPrunableTree<H>,
@ -359,6 +374,7 @@ impl<
// for some inputs, and given that it is always correct to not insert an empty
// subtree into `self`, we maintain the invariant by skipping empty subtrees.
if subtree.root().is_empty() {
debug!("Subtree with root {:?} is empty.", subtree.root_addr);
continue;
}
@ -368,14 +384,15 @@ impl<
let root_addr = Self::subtree_addr(subtree.root_addr.position_range_start());
let contains_marked = subtree.root.contains_marked();
let (new_subtree, mut incomplete) = self
let current_shard = self
.store
.get_shard(root_addr)
.map_err(ShardTreeError::Storage)?
.unwrap_or_else(|| LocatedTree::empty(root_addr))
.insert_subtree(subtree, contains_marked)?;
.unwrap_or_else(|| LocatedTree::empty(root_addr));
let (replacement_shard, mut incomplete) =
current_shard.insert_subtree(subtree, contains_marked)?;
self.store
.put_shard(new_subtree)
.put_shard(replacement_shard)
.map_err(ShardTreeError::Storage)?;
all_incomplete.append(&mut incomplete);
}
@ -396,8 +413,8 @@ impl<
root_addr: Address,
root: &PrunableTree<H>,
) -> Option<(PrunableTree<H>, Position)> {
match root {
Tree(Node::Parent { ann, left, right }) => {
match &root.0 {
Node::Parent { ann, left, right } => {
let (l_addr, r_addr) = root_addr.children().unwrap();
go(r_addr, right).map_or_else(
|| {
@ -426,13 +443,13 @@ impl<
},
)
}
Tree(Node::Leaf { value: (h, r) }) => Some((
Node::Leaf { value: (h, r) } => Some((
Tree(Node::Leaf {
value: (h.clone(), *r | RetentionFlags::CHECKPOINT),
}),
root_addr.max_position(),
)),
Tree(Node::Nil) => None,
Node::Nil | Node::Pruned => None,
}
}
@ -478,6 +495,7 @@ impl<
Ok(true)
}
#[tracing::instrument(skip(self))]
fn prune_excess_checkpoints(&mut self) -> Result<(), ShardTreeError<S::Error>> {
let checkpoint_count = self
.store
@ -559,12 +577,19 @@ impl<
// Prune each affected subtree
for (subtree_addr, positions) in clear_positions.into_iter() {
let cleared = self
let to_clear = self
.store
.get_shard(subtree_addr)
.map_err(ShardTreeError::Storage)?
.map(|subtree| subtree.clear_flags(positions));
if let Some(cleared) = cleared {
.map_err(ShardTreeError::Storage)?;
if let Some(to_clear) = to_clear {
let pre_clearing_max_position = to_clear.max_position();
let cleared = to_clear.clear_flags(positions);
// Clearing flags should not modify the max position of leaves represented
// in the shard.
assert!(cleared.max_position() == pre_clearing_max_position);
self.store
.put_shard(cleared)
.map_err(ShardTreeError::Storage)?;
@ -740,8 +765,8 @@ impl<
// roots.
truncate_at: Position,
) -> Result<(H, Option<PrunableTree<H>>), ShardTreeError<S::Error>> {
match &cap.root {
Tree(Node::Parent { ann, left, right }) => {
match &cap.root.0 {
Node::Parent { ann, left, right } => {
match ann {
Some(cached_root) if target_addr.contains(&cap.root_addr) => {
Ok((cached_root.as_ref().clone(), None))
@ -815,7 +840,7 @@ impl<
}
}
}
Tree(Node::Leaf { value }) => {
Node::Leaf { value } => {
if truncate_at >= cap.root_addr.position_range_end()
&& target_addr.contains(&cap.root_addr)
{
@ -840,7 +865,7 @@ impl<
))
}
}
Tree(Node::Nil) => {
Node::Nil | Node::Pruned => {
if cap.root_addr == target_addr
|| cap.root_addr.level() == ShardTree::<S, DEPTH, SHARD_HEIGHT>::subtree_level()
{
@ -850,7 +875,7 @@ impl<
Ok((
root.clone(),
if truncate_at >= cap.root_addr.position_range_end() {
// return the compute root as a new leaf to be cached if it contains no
// return the computed root as a new leaf to be cached if it contains no
// empty hashes due to truncation
Some(Tree::leaf((root, RetentionFlags::EPHEMERAL)))
} else {
@ -1284,21 +1309,24 @@ impl<
#[cfg(test)]
mod tests {
use std::convert::Infallible;
use assert_matches::assert_matches;
use proptest::prelude::*;
use incrementalmerkletree::{
frontier::NonEmptyFrontier,
frontier::{Frontier, NonEmptyFrontier},
testing::{
arb_operation, check_append, check_checkpoint_rewind, check_operations,
check_remove_mark, check_rewind_remove_mark, check_root_hashes,
check_witness_consistency, check_witnesses, complete_tree::CompleteTree, CombinedTree,
SipHashable,
},
Address, Hashable, Level, Position, Retention,
Address, Hashable, Level, Marking, MerklePath, Position, Retention,
};
use crate::{
error::{QueryError, ShardTreeError},
store::memory::MemoryShardStore,
testing::{
arb_char_str, arb_shardtree, check_shard_sizes, check_shardtree_insertion,
@ -1396,13 +1424,102 @@ mod tests {
'd'.to_string(),
Retention::Checkpoint {
id: 11,
is_marked: false
marking: Marking::None
},
),
Ok(()),
);
}
#[test]
fn avoid_pruning_reference() {
fn test_with_marking(
frontier_marking: Marking,
) -> Result<MerklePath<String, 6>, ShardTreeError<Infallible>> {
let mut tree = ShardTree::<MemoryShardStore<String, usize>, 6, 3>::new(
MemoryShardStore::empty(),
5,
);
let frontier_end = Position::from((1 << 3) - 3);
let mut f0 = Frontier::<String, 6>::empty();
for c in 'a'..='f' {
f0.append(c.to_string());
}
let frontier = Frontier::from_parts(
frontier_end,
"f".to_owned(),
vec!["e".to_owned(), "abcd".to_owned()],
)
.unwrap();
// Insert a frontier two leaves from the end of the first shard, checkpointed,
// with the specified marking.
tree.insert_frontier(
frontier,
Retention::Checkpoint {
id: 1,
marking: frontier_marking,
},
)?;
// Insert a few leaves beginning at the subsequent position, so as to cross the shard
// boundary.
tree.batch_insert(
frontier_end + 1,
('g'..='j')
.into_iter()
.map(|c| (c.to_string(), Retention::Ephemeral)),
)?;
// Trigger pruning by adding 5 more checkpoints
for i in 2..7 {
tree.checkpoint(i).unwrap();
}
// Insert nodes that require the pruned nodes for witnessing
tree.batch_insert(
frontier_end - 1,
('e'..='f')
.into_iter()
.map(|c| (c.to_string(), Retention::Marked)),
)?;
// Compute the witness
tree.witness_at_checkpoint_id(frontier_end, &6)
}
// If we insert the frontier with Marking::None, the frontier nodes are treated
// as ephemeral nodes and are pruned, leaving an incomplete tree.
assert_matches!(
test_with_marking(Marking::None),
Err(ShardTreeError::Query(QueryError::TreeIncomplete(_)))
);
// If we insert the frontier with Marking::Reference, the frontier nodes will
// not be pruned on completion of the subtree, and thus we'll be able to compute
// the witness.
let expected_witness = MerklePath::from_parts(
[
"e",
"gh",
"abcd",
"ij______",
"________________",
"________________________________",
]
.iter()
.map(|s| s.to_string())
.collect(),
Position::from(5),
)
.unwrap();
let witness = test_with_marking(Marking::Reference).unwrap();
assert_eq!(witness, expected_witness);
}
// Combined tree tests
#[allow(clippy::type_complexity)]
fn new_combined_tree<H: Hashable + Ord + Clone + core::fmt::Debug>(

380
shardtree/src/prunable.rs
View File

@ -4,10 +4,11 @@ use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use bitflags::bitflags;
use incrementalmerkletree::Marking;
use incrementalmerkletree::{
frontier::NonEmptyFrontier, Address, Hashable, Level, Position, Retention,
};
use tracing::trace;
use tracing::{trace, warn};
use crate::error::{InsertionError, QueryError};
use crate::{LocatedTree, Node, Tree};
@ -30,6 +31,12 @@ bitflags! {
/// 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;
}
}
@ -47,14 +54,17 @@ impl<'a, C> From<&'a Retention<C>> for RetentionFlags {
fn from(retention: &'a Retention<C>) -> Self {
match retention {
Retention::Ephemeral => RetentionFlags::EPHEMERAL,
Retention::Checkpoint { is_marked, .. } => {
if *is_marked {
RetentionFlags::CHECKPOINT | RetentionFlags::MARKED
} else {
RetentionFlags::CHECKPOINT
}
}
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,
}
}
}
@ -90,12 +100,25 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
.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 | Node::Pruned => false,
}
}
/// 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,
Node::Nil | Node::Pruned => false,
}
}
@ -112,8 +135,8 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
// so no need to inspect the tree
Ok(H::empty_root(root_addr.level()))
} else {
match self {
Tree(Node::Parent { ann, left, right }) => ann
match &self.0 {
Node::Parent { ann, left, right } => ann
.as_ref()
.filter(|_| truncate_at >= root_addr.position_range_end())
.map_or_else(
@ -135,7 +158,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
Ok(rc.as_ref().clone())
},
),
Tree(Node::Leaf { value }) => {
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())
@ -147,7 +170,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
Err(vec![root_addr])
}
}
Tree(Node::Nil) => Err(vec![root_addr]),
Node::Nil | Node::Pruned => Err(vec![root_addr]),
}
}
}
@ -182,7 +205,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
}
result
}
Node::Nil => BTreeSet::new(),
Node::Nil | Node::Pruned => BTreeSet::new(),
}
}
@ -206,6 +229,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
/// 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, Address> {
#[allow(clippy::type_complexity)]
fn go<H: Hashable + Clone + PartialEq>(
@ -218,6 +242,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
let no_default_fill = addr.position_range_end();
match (t0, t1) {
(Tree(Node::Nil), other) | (other, Tree(Node::Nil)) => Ok(other),
(Tree(Node::Pruned), other) | (other, Tree(Node::Pruned)) => Ok(other),
(Tree(Node::Leaf { value: vl }), Tree(Node::Leaf { value: vr })) => {
if vl.0 == vr.0 {
// Merge the flags together.
@ -279,12 +304,11 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
}
}
trace!(this = ?self, other = ?other, "Merging subtrees");
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 root by hashing leaf values together (or a
/// 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.
@ -292,10 +316,11 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
match (left, right) {
(Tree(Node::Nil), Tree(Node::Nil)) => Tree(Node::Nil),
(Tree(Node::Leaf { value: lv }), Tree(Node::Leaf { value: rv }))
// we can prune right-hand leaves that are not marked; 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::EPHEMERAL =>
// 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(
Node::Leaf {
@ -507,7 +532,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
None
}
}
Node::Nil => None,
Node::Nil | Node::Pruned => None,
}
}
@ -544,129 +569,187 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
root_addr: Address,
into: &PrunableTree<H>,
subtree: LocatedPrunableTree<H>,
is_complete: bool,
contains_marked: bool,
) -> Result<(PrunableTree<H>, Vec<IncompleteAt>), 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<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,
});
node = LocatedTree {
root_addr: node.root_addr.parent(),
root: if node.root_addr.is_right_child() {
Tree(Node::Parent {
ann: None,
left: Arc::new(Tree(Node::Nil)),
right: Arc::new(node.root),
})
} else {
Tree(Node::Parent {
ann: None,
left: Arc::new(node.root),
right: Arc::new(Tree(Node::Nil)),
})
},
};
}
(node.root.reannotate_root(ann), incomplete)
};
trace!(
"Node at {:?} contains subtree at {:?}",
root_addr,
subtree.root_addr(),
root_addr = ?root_addr,
max_position = ?LocatedTree::max_position_internal(root_addr, into),
to_insert = ?subtree.root_addr(),
"Subtree insert"
);
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.node_value().iter().all(|v| v == &value) {
// 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>,
pruned: bool| {
// 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 empty = Arc::new(Tree(if pruned { Node::Pruned } else { Node::Nil }));
let full = Arc::new(node.root);
node = LocatedTree {
root_addr: node.root_addr.parent(),
root: if node.root_addr.is_right_child() {
Tree(Node::Parent {
ann: None,
left: empty,
right: full,
})
} else {
Tree(Node::Parent {
ann: None,
left: full,
right: empty,
})
},
};
}
(node.root.reannotate_root(ann), incomplete)
};
match into {
Tree(Node::Nil) => Ok(replacement(None, subtree, false)),
Tree(Node::Pruned) => Ok(replacement(None, subtree, true)),
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,
// The subtree being inserted may have its root at some level lower
// than the next level down. The siblings of nodes that will be
// generated while descending to the subtree root level will be
// `Nil` nodes (indicating that the value of these nodes have never
// been observed) if the leaf being replaced has `REFERENCE`
// retention. Any other leaf without `REFERENCE` retention will
// have been produced by pruning of previously observed node
// values, so in those cases we use `Pruned` nodes for the absent
// siblings. This allows us to retain the distinction between what
// parts of the tree have been directly observed and what parts
// have not.
!retention.contains(RetentionFlags::REFERENCE),
))
}
}
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, contains_marked)?;
Ok((
// at this point we statically know the root to be a parent
subtree.root.reannotate_root(Some(Arc::new(value.clone()))),
vec![],
Tree::unite(
root_addr.level() - 1,
ann.clone(),
new_left,
right.as_ref().clone(),
),
incomplete,
))
} else {
trace!(
cur_root = ?value,
new_root = ?subtree.root.node_value(),
"Insertion conflict",
);
Err(InsertionError::Conflict(root_addr))
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,
))
}
} 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::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 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) {
let complete = subtree.root.is_complete();
go(*root_addr, root, subtree, complete, contains_marked).map(|(root, incomplete)| {
(
LocatedTree {
root_addr: *root_addr,
root,
},
incomplete,
)
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))
@ -725,9 +808,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
split_at: Level,
) -> (Self, Option<Self>) {
let mut addr = Address::from(position);
let mut subtree = Tree(Node::Leaf {
value: (leaf, leaf_retention.into()),
});
let mut subtree = Tree::leaf((leaf, leaf_retention.into()));
while addr.level() < split_at {
if addr.is_left_child() {
@ -829,13 +910,13 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// nothing to do, so we just return the root
root.clone()
} else {
match root {
Tree(Node::Parent { ann, left, right }) => {
match &root.0 {
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());
trace!(
"In {:?}, partitioned: {:?} {:?}",
"Tree::unite at {:?}, partitioned: {:?} {:?}",
root_addr,
&to_clear[0..p],
&to_clear[p..],
@ -847,7 +928,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
go(&to_clear[p..], r_addr, right),
)
}
Tree(Node::Leaf { value: (h, r) }) => {
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
@ -855,18 +936,16 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// 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),
})
}
[(_, flags)] => Tree(Node::Leaf {
value: (h.clone(), *r & !*flags),
}),
_ => {
panic!("Tree state inconsistent with checkpoints.");
}
}
}
Tree(Node::Nil) => Tree(Node::Nil),
Node::Nil => Tree(Node::Nil),
Node::Pruned => Tree(Node::Pruned),
}
}
}
@ -905,7 +984,7 @@ mod tests {
use super::{LocatedPrunableTree, PrunableTree, RetentionFlags};
use crate::{
error::QueryError,
error::{InsertionError, QueryError},
tree::{
tests::{leaf, nil, parent},
LocatedTree,
@ -1068,27 +1147,22 @@ mod tests {
}
#[test]
fn located_insert_subtree_leaf_overwrites() {
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: Address::from_parts(1.into(), 2),
root_addr: conflict_addr,
root: leaf(("b".to_string(), RetentionFlags::EPHEMERAL)),
},
false,
),
Ok((
LocatedTree {
root_addr: Address::from_parts(2.into(), 1),
root: parent(leaf(("b".to_string(), RetentionFlags::EPHEMERAL)), nil()),
},
vec![],
)),
Err(InsertionError::Conflict(conflict_addr)),
);
}

28
shardtree/src/store/caching.rs
View File

@ -219,7 +219,7 @@ mod tests {
append_str, check_operations, unmark, witness, CombinedTree, Operation, TestHashable,
Tree,
},
Hashable, Position, Retention,
Hashable, Marking, Position, Retention,
};
use super::CachingShardStore;
@ -300,7 +300,7 @@ mod tests {
String::from_u64(0),
Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
));
for _ in 0..3 {
@ -542,7 +542,7 @@ mod tests {
String::from_u64(0),
Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
));
assert!(tree.rewind());
@ -584,7 +584,7 @@ mod tests {
String::from_u64(6),
Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
));
assert!(tree.append(String::from_u64(7), Ephemeral));
@ -671,7 +671,7 @@ mod tests {
String::from_u64(3),
Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
),
Append(String::from_u64(4), Marked),
@ -680,21 +680,21 @@ mod tests {
String::from_u64(5),
Checkpoint {
id: 3,
is_marked: false,
marking: Marking::None,
},
),
Append(
String::from_u64(6),
Checkpoint {
id: 4,
is_marked: false,
marking: Marking::None,
},
),
Append(
String::from_u64(7),
Checkpoint {
id: 5,
is_marked: false,
marking: Marking::None,
},
),
Witness(3u64.into(), 5),
@ -732,7 +732,7 @@ mod tests {
String::from_u64(0),
Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
),
Append(String::from_u64(0), Ephemeral),
@ -742,7 +742,7 @@ mod tests {
String::from_u64(0),
Checkpoint {
id: 2,
is_marked: false,
marking: Marking::None,
},
),
Append(String::from_u64(0), Ephemeral),
@ -790,7 +790,7 @@ mod tests {
String::from_u64(0),
Checkpoint {
id: 4,
is_marked: false,
marking: Marking::None,
},
),
Rewind,
@ -818,14 +818,14 @@ mod tests {
String::from_u64(0),
Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
),
Append(
String::from_u64(0),
Checkpoint {
id: 4,
is_marked: false,
marking: Marking::None,
},
),
Witness(Position::from(2), 2),
@ -962,7 +962,7 @@ mod tests {
"a",
Retention::Checkpoint {
id: 1,
is_marked: true,
marking: Marking::Marked,
},
),
witness(1, 1),

21
shardtree/src/testing.rs
View File

@ -86,7 +86,14 @@ where
leaf,
match (is_checkpoint, is_marked) {
(false, false) => Retention::Ephemeral,
(true, is_marked) => Retention::Checkpoint { id, is_marked },
(true, is_marked) => Retention::Checkpoint {
id,
marking: if is_marked {
Marking::Marked
} else {
Marking::None
},
},
(false, true) => Retention::Marked,
},
)
@ -121,10 +128,10 @@ where
.map(|(id, (leaf, retention))| {
let pos = Position::try_from(id).unwrap();
match retention {
Retention::Ephemeral => (),
Retention::Checkpoint { is_marked, .. } => {
Retention::Ephemeral | Retention::Reference => (),
Retention::Checkpoint { marking, .. } => {
checkpoint_positions.push(pos);
if is_marked {
if marking == Marking::Marked {
marked_positions.push(pos);
}
}
@ -234,7 +241,7 @@ pub fn check_shardtree_insertion<
tree.batch_insert(
Position::from(1),
vec![
("b".to_string(), Retention::Checkpoint { id: 1, is_marked: false }),
("b".to_string(), Retention::Checkpoint { id: 1, marking: Marking::None }),
("c".to_string(), Retention::Ephemeral),
("d".to_string(), Retention::Marked),
].into_iter()
@ -285,7 +292,7 @@ pub fn check_shardtree_insertion<
Position::from(10),
vec![
("k".to_string(), Retention::Ephemeral),
("l".to_string(), Retention::Checkpoint { id: 2, is_marked: false }),
("l".to_string(), Retention::Checkpoint { id: 2, marking: Marking::None }),
("m".to_string(), Retention::Ephemeral),
].into_iter()
),
@ -386,7 +393,7 @@ pub fn check_witness_with_pruned_subtrees<
'c' => Retention::Marked,
'h' => Retention::Checkpoint {
id: 3,
is_marked: false,
marking: Marking::None,
},
_ => Retention::Ephemeral,
},

113
shardtree/src/tree.rs
View File

@ -19,6 +19,9 @@ pub enum Node<C, A, V> {
Leaf { value: V },
/// The empty tree; a subtree or leaf for which no information is available.
Nil,
/// An empty node in the tree created as a consequence of partial reinserion of data into a
/// subtree after the subtree was previously pruned.
Pruned,
}
impl<C, A, V> Node<C, A, V> {
@ -35,7 +38,8 @@ impl<C, A, V> Node<C, A, V> {
match self {
Node::Parent { .. } => None,
Node::Leaf { value } => Some(value),
Node::Nil { .. } => None,
Node::Nil => None,
Node::Pruned => None,
}
}
@ -45,6 +49,7 @@ impl<C, A, V> Node<C, A, V> {
Node::Parent { ann, .. } => Some(ann),
Node::Leaf { .. } => None,
Node::Nil => None,
Node::Pruned => None,
}
}
@ -71,6 +76,7 @@ impl<'a, C: Clone, A: Clone, V: Clone> Node<C, &'a A, &'a V> {
value: (*value).clone(),
},
Node::Nil => Node::Nil,
Node::Pruned => Node::Pruned,
}
}
}
@ -92,6 +98,11 @@ impl<A, V> Tree<A, V> {
Tree(Node::Nil)
}
/// Constructs the empty tree consisting of a single pruned node.
pub fn empty_pruned() -> Self {
Tree(Node::Pruned)
}
/// Constructs a tree containing a single leaf.
pub fn leaf(value: V) -> Self {
Tree(Node::Leaf { value })
@ -117,18 +128,13 @@ impl<A, V> Tree<A, V> {
Tree(self.0.reannotate(ann))
}
/// Returns `true` if no [`Node::Nil`] nodes are present in the tree, `false` otherwise.
pub fn is_complete(&self) -> bool {
match &self.0 {
Node::Parent { left, right, .. } => {
left.as_ref().is_complete() && right.as_ref().is_complete()
}
Node::Leaf { .. } => true,
Node::Nil { .. } => false,
}
/// Returns `true` this is a [`Node::Leaf`], `false` otherwise.
pub fn is_leaf(&self) -> bool {
matches!(&self.0, Node::Leaf { .. })
}
/// Returns a vector of the addresses of [`Node::Nil`] subtree roots within this tree.
/// Returns a vector of the addresses of [`Node::Nil`] and [`Node::Pruned`] subtree roots
/// within this tree.
///
/// The given address must correspond to the root of this tree, or this method will
/// yield incorrect results or may panic.
@ -149,9 +155,46 @@ impl<A, V> Tree<A, V> {
left_incomplete
}
Node::Leaf { .. } => vec![],
Node::Nil => vec![root_addr],
Node::Nil | Node::Pruned => vec![root_addr],
}
}
/// Applies the provided function to each leaf of the tree and returns
/// a new tree having the same structure as the original.
pub fn map<B, F: Fn(&V) -> B>(&self, f: &F) -> Tree<A, B>
where
A: Clone,
{
Tree(match &self.0 {
Node::Parent { ann, left, right } => Node::Parent {
ann: ann.clone(),
left: Arc::new(left.map(f)),
right: Arc::new(right.map(f)),
},
Node::Leaf { value } => Node::Leaf { value: f(value) },
Node::Nil => Node::Nil,
Node::Pruned => Node::Pruned,
})
}
/// Applies the provided function to each leaf of the tree and returns
/// a new tree having the same structure as the original, or an error
/// if any transformation of the leaf fails.
pub fn try_map<B, E, F: Fn(&V) -> Result<B, E>>(&self, f: &F) -> Result<Tree<A, B>, E>
where
A: Clone,
{
Ok(Tree(match &self.0 {
Node::Parent { ann, left, right } => Node::Parent {
ann: ann.clone(),
left: Arc::new(left.try_map(f)?),
right: Arc::new(right.try_map(f)?),
},
Node::Leaf { value } => Node::Leaf { value: f(value)? },
Node::Nil => Node::Nil,
Node::Pruned => Node::Pruned,
}))
}
}
/// A binary Merkle tree with its root at the given address.
@ -202,18 +245,19 @@ impl<A, V> LocatedTree<A, V> {
/// 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> {
fn go<A, V>(addr: Address, root: &Tree<A, V>) -> Option<Position> {
match &root.0 {
Node::Nil => None,
Node::Leaf { .. } => Some(addr.position_range_end() - 1),
Node::Parent { left, right, .. } => {
let (l_addr, r_addr) = addr.children().unwrap();
go(r_addr, right.as_ref()).or_else(|| go(l_addr, left.as_ref()))
}
Self::max_position_internal(self.root_addr, &self.root)
}
pub(crate) fn max_position_internal(addr: Address, root: &Tree<A, V>) -> Option<Position> {
match &root.0 {
Node::Nil => None,
Node::Leaf { .. } | Node::Pruned => Some(addr.position_range_end() - 1),
Node::Parent { left, right, .. } => {
let (l_addr, r_addr) = addr.children().unwrap();
Self::max_position_internal(r_addr, right.as_ref())
.or_else(|| Self::max_position_internal(l_addr, left.as_ref()))
}
}
go(self.root_addr, &self.root)
}
/// Returns the value at the specified position, if any.
@ -239,6 +283,31 @@ impl<A, V> LocatedTree<A, V> {
None
}
}
/// Applies the provided function to each leaf of the tree and returns
/// a new tree having the same structure as the original.
pub fn map<B, F: Fn(&V) -> B>(&self, f: &F) -> LocatedTree<A, B>
where
A: Clone,
{
LocatedTree {
root_addr: self.root_addr,
root: self.root.map(f),
}
}
/// Applies the provided function to each leaf of the tree and returns
/// a new tree having the same structure as the original, or an error
/// if any transformation of the leaf fails.
pub fn try_map<B, E, F: Fn(&V) -> Result<B, E>>(&self, f: &F) -> Result<LocatedTree<A, B>, E>
where
A: Clone,
{
Ok(LocatedTree {
root_addr: self.root_addr,
root: self.root.try_map(f)?,
})
}
}
impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {