Merge pull request #81 from zcash/refactor
Refactor codebase into submodules
This commit is contained in:
Kris Nuttycombe
GitHub
commit
2d86458e1e
2599
shardtree/src/lib.rs
2599
shardtree/src/lib.rs
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,160 @@
|
|||
//! Implementation of an in-memory shard store with no persistence.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::convert::{Infallible, TryFrom};
|
||||
|
||||
use incrementalmerkletree::Address;
|
||||
|
||||
use crate::{Checkpoint, LocatedPrunableTree, LocatedTree, Node, PrunableTree, ShardStore, Tree};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct MemoryShardStore<H, C: Ord> {
|
||||
shards: Vec<LocatedPrunableTree<H>>,
|
||||
checkpoints: BTreeMap<C, Checkpoint>,
|
||||
cap: PrunableTree<H>,
|
||||
}
|
||||
|
||||
impl<H, C: Ord> MemoryShardStore<H, C> {
|
||||
pub fn empty() -> Self {
|
||||
Self {
|
||||
shards: vec![],
|
||||
checkpoints: BTreeMap::new(),
|
||||
cap: PrunableTree::empty(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<H: Clone, C: Clone + Ord> ShardStore for MemoryShardStore<H, C> {
|
||||
type H = H;
|
||||
type CheckpointId = C;
|
||||
type Error = Infallible;
|
||||
|
||||
fn get_shard(
|
||||
&self,
|
||||
shard_root: Address,
|
||||
) -> Result<Option<LocatedPrunableTree<H>>, Self::Error> {
|
||||
let shard_idx =
|
||||
usize::try_from(shard_root.index()).expect("SHARD_HEIGHT > 64 is unsupported");
|
||||
Ok(self.shards.get(shard_idx).cloned())
|
||||
}
|
||||
|
||||
fn last_shard(&self) -> Result<Option<LocatedPrunableTree<H>>, Self::Error> {
|
||||
Ok(self.shards.last().cloned())
|
||||
}
|
||||
|
||||
fn put_shard(&mut self, subtree: LocatedPrunableTree<H>) -> Result<(), Self::Error> {
|
||||
let subtree_addr = subtree.root_addr;
|
||||
for subtree_idx in
|
||||
self.shards.last().map_or(0, |s| s.root_addr.index() + 1)..=subtree_addr.index()
|
||||
{
|
||||
self.shards.push(LocatedTree {
|
||||
root_addr: Address::from_parts(subtree_addr.level(), subtree_idx),
|
||||
root: Tree(Node::Nil),
|
||||
})
|
||||
}
|
||||
|
||||
let shard_idx =
|
||||
usize::try_from(subtree_addr.index()).expect("SHARD_HEIGHT > 64 is unsupported");
|
||||
self.shards[shard_idx] = subtree;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_shard_roots(&self) -> Result<Vec<Address>, Self::Error> {
|
||||
Ok(self.shards.iter().map(|s| s.root_addr).collect())
|
||||
}
|
||||
|
||||
fn truncate(&mut self, from: Address) -> Result<(), Self::Error> {
|
||||
let shard_idx = usize::try_from(from.index()).expect("SHARD_HEIGHT > 64 is unsupported");
|
||||
self.shards.truncate(shard_idx);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_cap(&self) -> Result<PrunableTree<H>, Self::Error> {
|
||||
Ok(self.cap.clone())
|
||||
}
|
||||
|
||||
fn put_cap(&mut self, cap: PrunableTree<H>) -> Result<(), Self::Error> {
|
||||
self.cap = cap;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn add_checkpoint(
|
||||
&mut self,
|
||||
checkpoint_id: C,
|
||||
checkpoint: Checkpoint,
|
||||
) -> Result<(), Self::Error> {
|
||||
self.checkpoints.insert(checkpoint_id, checkpoint);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn checkpoint_count(&self) -> Result<usize, Self::Error> {
|
||||
Ok(self.checkpoints.len())
|
||||
}
|
||||
|
||||
fn get_checkpoint(
|
||||
&self,
|
||||
checkpoint_id: &Self::CheckpointId,
|
||||
) -> Result<Option<Checkpoint>, Self::Error> {
|
||||
Ok(self.checkpoints.get(checkpoint_id).cloned())
|
||||
}
|
||||
|
||||
fn get_checkpoint_at_depth(
|
||||
&self,
|
||||
checkpoint_depth: usize,
|
||||
) -> Result<Option<(C, Checkpoint)>, Self::Error> {
|
||||
Ok(if checkpoint_depth == 0 {
|
||||
None
|
||||
} else {
|
||||
self.checkpoints
|
||||
.iter()
|
||||
.rev()
|
||||
.nth(checkpoint_depth - 1)
|
||||
.map(|(id, c)| (id.clone(), c.clone()))
|
||||
})
|
||||
}
|
||||
|
||||
fn min_checkpoint_id(&self) -> Result<Option<C>, Self::Error> {
|
||||
Ok(self.checkpoints.keys().next().cloned())
|
||||
}
|
||||
|
||||
fn max_checkpoint_id(&self) -> Result<Option<C>, Self::Error> {
|
||||
Ok(self.checkpoints.keys().last().cloned())
|
||||
}
|
||||
|
||||
fn with_checkpoints<F>(&mut self, limit: usize, mut callback: F) -> Result<(), Self::Error>
|
||||
where
|
||||
F: FnMut(&C, &Checkpoint) -> Result<(), Self::Error>,
|
||||
{
|
||||
for (cid, checkpoint) in self.checkpoints.iter().take(limit) {
|
||||
callback(cid, checkpoint)?
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn update_checkpoint_with<F>(
|
||||
&mut self,
|
||||
checkpoint_id: &C,
|
||||
update: F,
|
||||
) -> Result<bool, Self::Error>
|
||||
where
|
||||
F: Fn(&mut Checkpoint) -> Result<(), Self::Error>,
|
||||
{
|
||||
if let Some(c) = self.checkpoints.get_mut(checkpoint_id) {
|
||||
update(c)?;
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
fn remove_checkpoint(&mut self, checkpoint_id: &C) -> Result<(), Self::Error> {
|
||||
self.checkpoints.remove(checkpoint_id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn truncate_checkpoints(&mut self, checkpoint_id: &C) -> Result<(), Self::Error> {
|
||||
self.checkpoints.split_off(checkpoint_id);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,392 @@
|
|||
use assert_matches::assert_matches;
|
||||
use proptest::bool::weighted;
|
||||
use proptest::collection::vec;
|
||||
use proptest::prelude::*;
|
||||
use proptest::sample::select;
|
||||
|
||||
use incrementalmerkletree::{testing, Hashable};
|
||||
|
||||
use super::*;
|
||||
use crate::memory::MemoryShardStore;
|
||||
|
||||
pub fn arb_retention_flags() -> impl Strategy<Value = RetentionFlags> + Clone {
|
||||
select(vec![
|
||||
RetentionFlags::EPHEMERAL,
|
||||
RetentionFlags::CHECKPOINT,
|
||||
RetentionFlags::MARKED,
|
||||
RetentionFlags::MARKED | RetentionFlags::CHECKPOINT,
|
||||
])
|
||||
}
|
||||
|
||||
pub fn arb_tree<A: Strategy + Clone + 'static, V: Strategy + 'static>(
|
||||
arb_annotation: A,
|
||||
arb_leaf: V,
|
||||
depth: u32,
|
||||
size: u32,
|
||||
) -> impl Strategy<Value = Tree<A::Value, V::Value>> + Clone
|
||||
where
|
||||
A::Value: Clone + 'static,
|
||||
V::Value: Clone + 'static,
|
||||
{
|
||||
let leaf = prop_oneof![
|
||||
Just(Tree(Node::Nil)),
|
||||
arb_leaf.prop_map(|value| Tree(Node::Leaf { value }))
|
||||
];
|
||||
|
||||
leaf.prop_recursive(depth, size, 2, move |inner| {
|
||||
(arb_annotation.clone(), inner.clone(), inner).prop_map(|(ann, left, right)| {
|
||||
Tree(if left.is_nil() && right.is_nil() {
|
||||
Node::Nil
|
||||
} else {
|
||||
Node::Parent {
|
||||
ann,
|
||||
left: Rc::new(left),
|
||||
right: Rc::new(right),
|
||||
}
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn arb_prunable_tree<H: Strategy + Clone + 'static>(
|
||||
arb_leaf: H,
|
||||
depth: u32,
|
||||
size: u32,
|
||||
) -> impl Strategy<Value = PrunableTree<H::Value>> + Clone
|
||||
where
|
||||
H::Value: Clone + 'static,
|
||||
{
|
||||
arb_tree(
|
||||
proptest::option::of(arb_leaf.clone().prop_map(Rc::new)),
|
||||
(arb_leaf, arb_retention_flags()),
|
||||
depth,
|
||||
size,
|
||||
)
|
||||
}
|
||||
|
||||
/// Constructs a random shardtree of size up to 2^6 with shards of size 2^3. Returns the tree,
|
||||
/// along with vectors of the checkpoint and mark positions.
|
||||
pub fn arb_shardtree<H: Strategy + Clone>(
|
||||
arb_leaf: H,
|
||||
) -> impl Strategy<
|
||||
Value = (
|
||||
ShardTree<MemoryShardStore<H::Value, usize>, 6, 3>,
|
||||
Vec<Position>,
|
||||
Vec<Position>,
|
||||
),
|
||||
>
|
||||
where
|
||||
H::Value: Hashable + Clone + PartialEq,
|
||||
{
|
||||
vec(
|
||||
(arb_leaf, weighted(0.1), weighted(0.2)),
|
||||
0..=(2usize.pow(6)),
|
||||
)
|
||||
.prop_map(|leaves| {
|
||||
let mut tree = ShardTree::new(MemoryShardStore::empty(), 10);
|
||||
let mut checkpoint_positions = vec![];
|
||||
let mut marked_positions = vec![];
|
||||
tree.batch_insert(
|
||||
Position::from(0),
|
||||
leaves
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(id, (leaf, is_marked, is_checkpoint))| {
|
||||
(
|
||||
leaf,
|
||||
match (is_checkpoint, is_marked) {
|
||||
(false, false) => Retention::Ephemeral,
|
||||
(true, is_marked) => {
|
||||
let pos = Position::try_from(id).unwrap();
|
||||
checkpoint_positions.push(pos);
|
||||
if is_marked {
|
||||
marked_positions.push(pos);
|
||||
}
|
||||
Retention::Checkpoint { id, is_marked }
|
||||
}
|
||||
(false, true) => {
|
||||
marked_positions.push(Position::try_from(id).unwrap());
|
||||
Retention::Marked
|
||||
}
|
||||
},
|
||||
)
|
||||
}),
|
||||
)
|
||||
.unwrap();
|
||||
(tree, checkpoint_positions, marked_positions)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn arb_char_str() -> impl Strategy<Value = String> + Clone {
|
||||
let chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
|
||||
(0usize..chars.len()).prop_map(move |i| chars.get(i..=i).unwrap().to_string())
|
||||
}
|
||||
|
||||
impl<
|
||||
H: Hashable + Ord + Clone + core::fmt::Debug,
|
||||
C: Clone + Ord + core::fmt::Debug,
|
||||
S: ShardStore<H = H, CheckpointId = C>,
|
||||
const DEPTH: u8,
|
||||
const SHARD_HEIGHT: u8,
|
||||
> testing::Tree<H, C> for ShardTree<S, DEPTH, SHARD_HEIGHT>
|
||||
where
|
||||
S::Error: std::fmt::Debug,
|
||||
{
|
||||
fn depth(&self) -> u8 {
|
||||
DEPTH
|
||||
}
|
||||
|
||||
fn append(&mut self, value: H, retention: Retention<C>) -> bool {
|
||||
match ShardTree::append(self, value, retention) {
|
||||
Ok(_) => true,
|
||||
Err(ShardTreeError::Insert(InsertionError::TreeFull)) => false,
|
||||
Err(other) => panic!("append failed due to error: {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
fn current_position(&self) -> Option<Position> {
|
||||
match ShardTree::max_leaf_position(self, 0) {
|
||||
Ok(v) => v,
|
||||
Err(err) => panic!("current position query failed: {:?}", err),
|
||||
}
|
||||
}
|
||||
|
||||
fn get_marked_leaf(&self, position: Position) -> Option<H> {
|
||||
match ShardTree::get_marked_leaf(self, position) {
|
||||
Ok(v) => v,
|
||||
Err(err) => panic!("marked leaf query failed: {:?}", err),
|
||||
}
|
||||
}
|
||||
|
||||
fn marked_positions(&self) -> BTreeSet<Position> {
|
||||
match ShardTree::marked_positions(self) {
|
||||
Ok(v) => v,
|
||||
Err(err) => panic!("marked positions query failed: {:?}", err),
|
||||
}
|
||||
}
|
||||
|
||||
fn root(&self, checkpoint_depth: usize) -> Option<H> {
|
||||
match ShardTree::root_at_checkpoint(self, checkpoint_depth) {
|
||||
Ok(v) => Some(v),
|
||||
Err(err) => panic!("root computation failed: {:?}", err),
|
||||
}
|
||||
}
|
||||
|
||||
fn witness(&self, position: Position, checkpoint_depth: usize) -> Option<Vec<H>> {
|
||||
match ShardTree::witness(self, position, checkpoint_depth) {
|
||||
Ok(p) => Some(p.path_elems().to_vec()),
|
||||
Err(ShardTreeError::Query(
|
||||
QueryError::NotContained(_)
|
||||
| QueryError::TreeIncomplete(_)
|
||||
| QueryError::CheckpointPruned,
|
||||
)) => None,
|
||||
Err(err) => panic!("witness computation failed: {:?}", err),
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_mark(&mut self, position: Position) -> bool {
|
||||
let max_checkpoint = self
|
||||
.store
|
||||
.max_checkpoint_id()
|
||||
.unwrap_or_else(|err| panic!("checkpoint retrieval failed: {:?}", err));
|
||||
|
||||
match ShardTree::remove_mark(self, position, max_checkpoint.as_ref()) {
|
||||
Ok(result) => result,
|
||||
Err(err) => panic!("mark removal failed: {:?}", err),
|
||||
}
|
||||
}
|
||||
|
||||
fn checkpoint(&mut self, checkpoint_id: C) -> bool {
|
||||
ShardTree::checkpoint(self, checkpoint_id).unwrap()
|
||||
}
|
||||
|
||||
fn rewind(&mut self) -> bool {
|
||||
ShardTree::truncate_to_depth(self, 1).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn check_shardtree_insertion<
|
||||
E: Debug,
|
||||
S: ShardStore<H = String, CheckpointId = u32, Error = E>,
|
||||
>(
|
||||
mut tree: ShardTree<S, 4, 3>,
|
||||
) {
|
||||
assert_matches!(
|
||||
tree.batch_insert(
|
||||
Position::from(1),
|
||||
vec![
|
||||
("b".to_string(), Retention::Checkpoint { id: 1, is_marked: false }),
|
||||
("c".to_string(), Retention::Ephemeral),
|
||||
("d".to_string(), Retention::Marked),
|
||||
].into_iter()
|
||||
),
|
||||
Ok(Some((pos, incomplete))) if
|
||||
pos == Position::from(3) &&
|
||||
incomplete == vec![
|
||||
IncompleteAt {
|
||||
address: Address::from_parts(Level::from(0), 0),
|
||||
required_for_witness: true
|
||||
},
|
||||
IncompleteAt {
|
||||
address: Address::from_parts(Level::from(2), 1),
|
||||
required_for_witness: true
|
||||
}
|
||||
]
|
||||
);
|
||||
|
||||
assert_matches!(
|
||||
tree.root_at_checkpoint(1),
|
||||
Err(ShardTreeError::Query(QueryError::TreeIncomplete(v))) if v == vec![Address::from_parts(Level::from(0), 0)]
|
||||
);
|
||||
|
||||
assert_matches!(
|
||||
tree.batch_insert(
|
||||
Position::from(0),
|
||||
vec![
|
||||
("a".to_string(), Retention::Ephemeral),
|
||||
].into_iter()
|
||||
),
|
||||
Ok(Some((pos, incomplete))) if
|
||||
pos == Position::from(0) &&
|
||||
incomplete == vec![]
|
||||
);
|
||||
|
||||
assert_matches!(
|
||||
tree.root_at_checkpoint(0),
|
||||
Ok(h) if h == *"abcd____________"
|
||||
);
|
||||
|
||||
assert_matches!(
|
||||
tree.root_at_checkpoint(1),
|
||||
Ok(h) if h == *"ab______________"
|
||||
);
|
||||
|
||||
assert_matches!(
|
||||
tree.batch_insert(
|
||||
Position::from(10),
|
||||
vec![
|
||||
("k".to_string(), Retention::Ephemeral),
|
||||
("l".to_string(), Retention::Checkpoint { id: 2, is_marked: false }),
|
||||
("m".to_string(), Retention::Ephemeral),
|
||||
].into_iter()
|
||||
),
|
||||
Ok(Some((pos, incomplete))) if
|
||||
pos == Position::from(12) &&
|
||||
incomplete == vec![
|
||||
IncompleteAt {
|
||||
address: Address::from_parts(Level::from(0), 13),
|
||||
required_for_witness: false
|
||||
},
|
||||
IncompleteAt {
|
||||
address: Address::from_parts(Level::from(1), 7),
|
||||
required_for_witness: false
|
||||
},
|
||||
IncompleteAt {
|
||||
address: Address::from_parts(Level::from(1), 4),
|
||||
required_for_witness: false
|
||||
},
|
||||
]
|
||||
);
|
||||
|
||||
assert_matches!(
|
||||
tree.root_at_checkpoint(0),
|
||||
// The (0, 13) and (1, 7) incomplete subtrees are
|
||||
// not considered incomplete here because they appear
|
||||
// at the tip of the tree.
|
||||
Err(ShardTreeError::Query(QueryError::TreeIncomplete(xs))) if xs == vec![
|
||||
Address::from_parts(Level::from(2), 1),
|
||||
Address::from_parts(Level::from(1), 4),
|
||||
]
|
||||
);
|
||||
|
||||
assert_matches!(tree.truncate_to_depth(1), Ok(true));
|
||||
|
||||
assert_matches!(
|
||||
tree.batch_insert(
|
||||
Position::from(4),
|
||||
('e'..'k')
|
||||
.into_iter()
|
||||
.map(|c| (c.to_string(), Retention::Ephemeral))
|
||||
),
|
||||
Ok(_)
|
||||
);
|
||||
|
||||
assert_matches!(
|
||||
tree.root_at_checkpoint(0),
|
||||
Ok(h) if h == *"abcdefghijkl____"
|
||||
);
|
||||
|
||||
assert_matches!(
|
||||
tree.root_at_checkpoint(1),
|
||||
Ok(h) if h == *"ab______________"
|
||||
);
|
||||
}
|
||||
|
||||
pub fn check_shard_sizes<E: Debug, S: ShardStore<H = String, CheckpointId = u32, Error = E>>(
|
||||
mut tree: ShardTree<S, 4, 2>,
|
||||
) {
|
||||
for c in 'a'..'p' {
|
||||
tree.append(c.to_string(), Retention::Ephemeral).unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(tree.store.get_shard_roots().unwrap().len(), 4);
|
||||
assert_eq!(
|
||||
tree.store
|
||||
.get_shard(Address::from_parts(Level::from(2), 3))
|
||||
.unwrap()
|
||||
.and_then(|t| t.max_position()),
|
||||
Some(Position::from(14))
|
||||
);
|
||||
}
|
||||
|
||||
pub fn check_witness_with_pruned_subtrees<
|
||||
E: Debug,
|
||||
S: ShardStore<H = String, CheckpointId = u32, Error = E>,
|
||||
>(
|
||||
mut tree: ShardTree<S, 6, 3>,
|
||||
) {
|
||||
// introduce some roots
|
||||
let shard_root_level = Level::from(3);
|
||||
for idx in 0u64..4 {
|
||||
let root = if idx == 3 {
|
||||
"abcdefgh".to_string()
|
||||
} else {
|
||||
idx.to_string()
|
||||
};
|
||||
tree.insert(Address::from_parts(shard_root_level, idx), root)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// simulate discovery of a note
|
||||
tree.batch_insert(
|
||||
Position::from(24),
|
||||
('a'..='h').into_iter().map(|c| {
|
||||
(
|
||||
c.to_string(),
|
||||
match c {
|
||||
'c' => Retention::Marked,
|
||||
'h' => Retention::Checkpoint {
|
||||
id: 3,
|
||||
is_marked: false,
|
||||
},
|
||||
_ => Retention::Ephemeral,
|
||||
},
|
||||
)
|
||||
}),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// construct a witness for the note
|
||||
let witness = tree.witness(Position::from(26), 0).unwrap();
|
||||
assert_eq!(
|
||||
witness.path_elems(),
|
||||
&[
|
||||
"d",
|
||||
"ab",
|
||||
"efgh",
|
||||
"2",
|
||||
"01",
|
||||
"________________________________"
|
||||
]
|
||||
);
|
||||
}
|
||||
|
|
@ -0,0 +1,421 @@
|
|||
use std::ops::Deref;
|
||||
use std::rc::Rc;
|
||||
|
||||
use incrementalmerkletree::{Address, Level, Position};
|
||||
|
||||
/// A "pattern functor" for a single layer of a binary tree.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Node<C, A, V> {
|
||||
/// A parent node in the tree, annotated with a value of type `A` and with left and right
|
||||
/// children of type `C`.
|
||||
Parent { ann: A, left: C, right: C },
|
||||
/// A node of the tree that contains a value (usually a hash, sometimes with additional
|
||||
/// metadata) and that has no children.
|
||||
///
|
||||
/// Note that leaf nodes may appear at any position in the tree; i.e. they may contain computed
|
||||
/// subtree root values and not just level-0 leaves.
|
||||
Leaf { value: V },
|
||||
/// The empty tree; a subtree or leaf for which no information is available.
|
||||
Nil,
|
||||
}
|
||||
|
||||
impl<C, A, V> Node<C, A, V> {
|
||||
/// Returns whether or not this is the `Nil` tree.
|
||||
///
|
||||
/// This is useful for cases where the compiler can automatically dereference an `Rc`, where
|
||||
/// one would otherwise need additional ceremony to make an equality check.
|
||||
pub fn is_nil(&self) -> bool {
|
||||
matches!(self, Node::Nil)
|
||||
}
|
||||
|
||||
/// Returns the contained leaf value, if this is a leaf node.
|
||||
pub fn leaf_value(&self) -> Option<&V> {
|
||||
match self {
|
||||
Node::Parent { .. } => None,
|
||||
Node::Leaf { value } => Some(value),
|
||||
Node::Nil { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn annotation(&self) -> Option<&A> {
|
||||
match self {
|
||||
Node::Parent { ann, .. } => Some(ann),
|
||||
Node::Leaf { .. } => None,
|
||||
Node::Nil => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Replaces the annotation on this node, if it is a `Node::Parent`; otherwise
|
||||
/// returns this node unaltered.
|
||||
pub fn reannotate(self, ann: A) -> Self {
|
||||
match self {
|
||||
Node::Parent { left, right, .. } => Node::Parent { ann, left, right },
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, C: Clone, A: Clone, V: Clone> Node<C, &'a A, &'a V> {
|
||||
pub fn cloned(&self) -> Node<C, A, V> {
|
||||
match self {
|
||||
Node::Parent { ann, left, right } => Node::Parent {
|
||||
ann: (*ann).clone(),
|
||||
left: left.clone(),
|
||||
right: right.clone(),
|
||||
},
|
||||
Node::Leaf { value } => Node::Leaf {
|
||||
value: (*value).clone(),
|
||||
},
|
||||
Node::Nil => Node::Nil,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// An immutable binary tree with each of its nodes tagged with an annotation value.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Tree<A, V>(pub(crate) Node<Rc<Tree<A, V>>, A, V>);
|
||||
|
||||
impl<A, V> Deref for Tree<A, V> {
|
||||
type Target = Node<Rc<Tree<A, V>>, A, V>;
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl<A, V> Tree<A, V> {
|
||||
pub fn empty() -> Self {
|
||||
Tree(Node::Nil)
|
||||
}
|
||||
|
||||
pub fn leaf(value: V) -> Self {
|
||||
Tree(Node::Leaf { value })
|
||||
}
|
||||
|
||||
pub fn parent(ann: A, left: Self, right: Self) -> Self {
|
||||
Tree(Node::Parent {
|
||||
ann,
|
||||
left: Rc::new(left),
|
||||
right: Rc::new(right),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.0.is_nil()
|
||||
}
|
||||
|
||||
/// Replaces the annotation at the root of the tree, if the root is a `Node::Parent`; otherwise
|
||||
/// returns this tree unaltered.
|
||||
pub fn reannotate_root(self, ann: A) -> Self {
|
||||
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 a vector of the addresses of [`Node::Nil`] 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.
|
||||
pub fn incomplete_nodes(&self, root_addr: Address) -> Vec<Address> {
|
||||
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.incomplete_nodes(left_root);
|
||||
let mut right_incomplete = right.incomplete_nodes(right_root);
|
||||
left_incomplete.append(&mut right_incomplete);
|
||||
left_incomplete
|
||||
}
|
||||
Node::Leaf { .. } => vec![],
|
||||
Node::Nil => vec![root_addr],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A binary Merkle tree with its root at the given address.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct LocatedTree<A, V> {
|
||||
pub(crate) root_addr: Address,
|
||||
pub(crate) root: Tree<A, V>,
|
||||
}
|
||||
|
||||
impl<A, V> LocatedTree<A, V> {
|
||||
/// Constructs a new LocatedTree from its constituent parts
|
||||
pub fn from_parts(root_addr: Address, root: Tree<A, V>) -> Self {
|
||||
LocatedTree { root_addr, root }
|
||||
}
|
||||
|
||||
/// Returns the root address of this tree.
|
||||
pub fn root_addr(&self) -> Address {
|
||||
self.root_addr
|
||||
}
|
||||
|
||||
/// Returns a reference to the root of the tree.
|
||||
pub fn root(&self) -> &Tree<A, V> {
|
||||
&self.root
|
||||
}
|
||||
|
||||
/// Consumes this tree and returns its root as an owned value.
|
||||
pub fn take_root(self) -> Tree<A, V> {
|
||||
self.root
|
||||
}
|
||||
|
||||
/// Returns a new [`LocatedTree`] with the provided value replacing the annotation of its root
|
||||
/// node, if that root node is a `Node::Parent`. Otherwise .
|
||||
pub fn reannotate_root(self, value: A) -> Self {
|
||||
LocatedTree {
|
||||
root_addr: self.root_addr,
|
||||
root: self.root.reannotate_root(value),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the set of incomplete subtree roots contained within this tree, ordered by
|
||||
/// increasing position.
|
||||
pub fn incomplete_nodes(&self) -> Vec<Address> {
|
||||
self.root.incomplete_nodes(self.root_addr)
|
||||
}
|
||||
|
||||
/// 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> {
|
||||
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()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
go(self.root_addr, &self.root)
|
||||
}
|
||||
|
||||
/// Returns the value at the specified position, if any.
|
||||
pub fn value_at_position(&self, position: Position) -> Option<&V> {
|
||||
fn go<A, V>(pos: Position, addr: Address, root: &Tree<A, V>) -> Option<&V> {
|
||||
match &root.0 {
|
||||
Node::Parent { left, right, .. } => {
|
||||
let (l_addr, r_addr) = addr.children().unwrap();
|
||||
if l_addr.position_range().contains(&pos) {
|
||||
go(pos, l_addr, left)
|
||||
} else {
|
||||
go(pos, r_addr, right)
|
||||
}
|
||||
}
|
||||
Node::Leaf { value } if addr.level() == Level::from(0) => Some(value),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
if self.root_addr.position_range().contains(&position) {
|
||||
go(position, self.root_addr, &self.root)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
|
||||
/// Constructs a new empty tree with its root at the provided address.
|
||||
pub fn empty(root_addr: Address) -> Self {
|
||||
Self {
|
||||
root_addr,
|
||||
root: Tree(Node::Nil),
|
||||
}
|
||||
}
|
||||
|
||||
/// Constructs a new tree consisting of a single leaf with the provided value, and the
|
||||
/// specified root address.
|
||||
pub fn with_root_value(root_addr: Address, value: V) -> Self {
|
||||
Self {
|
||||
root_addr,
|
||||
root: Tree(Node::Leaf { value }),
|
||||
}
|
||||
}
|
||||
|
||||
/// Traverses this tree to find the child node at the specified address and returns it.
|
||||
///
|
||||
/// Returns `None` if the specified address is not a descendant of this tree's root address, or
|
||||
/// if the tree is terminated by a [`Node::Nil`] or leaf node before the specified address can
|
||||
/// be reached.
|
||||
pub fn subtree(&self, addr: Address) -> Option<Self> {
|
||||
fn go<A: Clone, V: Clone>(
|
||||
root_addr: Address,
|
||||
root: &Tree<A, V>,
|
||||
addr: Address,
|
||||
) -> Option<LocatedTree<A, V>> {
|
||||
if root_addr == addr {
|
||||
Some(LocatedTree {
|
||||
root_addr,
|
||||
root: root.clone(),
|
||||
})
|
||||
} else {
|
||||
match &root.0 {
|
||||
Node::Parent { left, right, .. } => {
|
||||
let (l_addr, r_addr) = root_addr.children().unwrap();
|
||||
if l_addr.contains(&addr) {
|
||||
go(l_addr, left.as_ref(), addr)
|
||||
} else {
|
||||
go(r_addr, right.as_ref(), addr)
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if self.root_addr.contains(&addr) {
|
||||
go(self.root_addr, &self.root, addr)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Decomposes this tree into the vector of its subtrees having height `level + 1`.
|
||||
///
|
||||
/// If this root address of this tree is lower down in the tree than the level specified,
|
||||
/// the entire tree is returned as the sole element of the result vector.
|
||||
pub fn decompose_to_level(self, level: Level) -> Vec<Self> {
|
||||
fn go<A: Clone, V: Clone>(
|
||||
level: Level,
|
||||
root_addr: Address,
|
||||
root: Tree<A, V>,
|
||||
) -> Vec<LocatedTree<A, V>> {
|
||||
if root_addr.level() == level {
|
||||
vec![LocatedTree { root_addr, root }]
|
||||
} else {
|
||||
match root.0 {
|
||||
Node::Parent { left, right, .. } => {
|
||||
let (l_addr, r_addr) = root_addr.children().unwrap();
|
||||
let mut l_decomposed = go(
|
||||
level,
|
||||
l_addr,
|
||||
Rc::try_unwrap(left).unwrap_or_else(|rc| (*rc).clone()),
|
||||
);
|
||||
let mut r_decomposed = go(
|
||||
level,
|
||||
r_addr,
|
||||
Rc::try_unwrap(right).unwrap_or_else(|rc| (*rc).clone()),
|
||||
);
|
||||
l_decomposed.append(&mut r_decomposed);
|
||||
l_decomposed
|
||||
}
|
||||
_ => vec![],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if level >= self.root_addr.level() {
|
||||
vec![self]
|
||||
} else {
|
||||
go(level, self.root_addr, self.root)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod tests {
|
||||
use incrementalmerkletree::{Address, Level};
|
||||
|
||||
use super::{LocatedTree, Node, Tree};
|
||||
|
||||
pub(crate) fn str_leaf<A>(c: &str) -> Tree<A, String> {
|
||||
Tree(Node::Leaf {
|
||||
value: c.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn nil<A, B>() -> Tree<A, B> {
|
||||
Tree::empty()
|
||||
}
|
||||
|
||||
pub(crate) fn leaf<A, B>(value: B) -> Tree<A, B> {
|
||||
Tree::leaf(value)
|
||||
}
|
||||
|
||||
pub(crate) fn parent<A: Default, B>(left: Tree<A, B>, right: Tree<A, B>) -> Tree<A, B> {
|
||||
Tree::parent(A::default(), left, right)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incomplete_nodes() {
|
||||
let t: Tree<(), String> = parent(nil(), str_leaf("a"));
|
||||
assert_eq!(
|
||||
t.incomplete_nodes(Address::from_parts(Level::from(1), 0)),
|
||||
vec![Address::from_parts(Level::from(0), 0)]
|
||||
);
|
||||
|
||||
let t0 = parent(str_leaf("b"), t.clone());
|
||||
assert_eq!(
|
||||
t0.incomplete_nodes(Address::from_parts(Level::from(2), 1)),
|
||||
vec![Address::from_parts(Level::from(0), 6)]
|
||||
);
|
||||
|
||||
let t1 = parent(nil(), t);
|
||||
assert_eq!(
|
||||
t1.incomplete_nodes(Address::from_parts(Level::from(2), 1)),
|
||||
vec![
|
||||
Address::from_parts(Level::from(1), 2),
|
||||
Address::from_parts(Level::from(0), 6)
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn located() {
|
||||
let l = parent(str_leaf("a"), str_leaf("b"));
|
||||
let r = parent(str_leaf("c"), str_leaf("d"));
|
||||
|
||||
let t: LocatedTree<(), String> = LocatedTree {
|
||||
root_addr: Address::from_parts(2.into(), 1),
|
||||
root: parent(l.clone(), r.clone()),
|
||||
};
|
||||
|
||||
assert_eq!(t.max_position(), Some(7.into()));
|
||||
assert_eq!(t.value_at_position(5.into()), Some(&"b".to_string()));
|
||||
assert_eq!(t.value_at_position(8.into()), None);
|
||||
assert_eq!(t.subtree(Address::from_parts(0.into(), 1)), None);
|
||||
assert_eq!(t.subtree(Address::from_parts(3.into(), 0)), None);
|
||||
|
||||
let subtree_addr = Address::from_parts(1.into(), 3);
|
||||
assert_eq!(
|
||||
t.subtree(subtree_addr),
|
||||
Some(LocatedTree {
|
||||
root_addr: subtree_addr,
|
||||
root: r.clone()
|
||||
})
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
t.decompose_to_level(1.into()),
|
||||
vec![
|
||||
LocatedTree {
|
||||
root_addr: Address::from_parts(1.into(), 2),
|
||||
root: l,
|
||||
},
|
||||
LocatedTree {
|
||||
root_addr: Address::from_parts(1.into(), 3),
|
||||
root: r,
|
||||
}
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue