Skip to content

Commit

Permalink
shardtree: Use Tree constructors instead of struct creation
Browse files Browse the repository at this point in the history 
This makes it easier for us to track where new tree nodes are created,
and how the node kinds are used within the `ShardTree` data model.
  • Loading branch information
@str4d
str4d committed Jun 26, 2024
1 parent 8b4b131 commit 49c2524
Show file tree
Hide file tree
Showing 6 changed files with 58 additions and 83 deletions.
16 changes: 5 additions & 11 deletions shardtree/src/batch.rs
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
//! Helpers for inserting many leaves into a tree at once.
use std::{collections::BTreeMap, fmt, ops::Range, sync::Arc};
use std::{collections::BTreeMap, fmt, ops::Range};

use incrementalmerkletree::{Address, Hashable, Level, Position, Retention};
use tracing::trace;

use crate::{
error::{InsertionError, ShardTreeError},
store::{Checkpoint, ShardStore},
IncompleteAt, LocatedPrunableTree, LocatedTree, Node, RetentionFlags, ShardTree, Tree,
IncompleteAt, LocatedPrunableTree, LocatedTree, RetentionFlags, ShardTree, Tree,
};

impl<
Expand Down Expand Up @@ -203,9 +203,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
let rflags = RetentionFlags::from(retention);
let mut subtree = LocatedTree {
root_addr: Address::from(position),
root: Tree(Node::Leaf {
value: (value.clone(), rflags),
}),
root: Tree::leaf((value.clone(), rflags)),
};

if position.is_right_child() {
Expand Down Expand Up @@ -268,11 +266,7 @@ fn unite<H: Hashable + Clone + PartialEq>(
root: if lroot.root_addr.level() < prune_below {
Tree::unite(lroot.root_addr.level(), None, lroot.root, rroot.root)
} else {
Tree(Node::Parent {
ann: None,
left: Arc::new(lroot.root),
right: Arc::new(rroot.root),
})
Tree::parent(None, lroot.root, rroot.root)
},
}
}
Expand All @@ -297,7 +291,7 @@ fn combine_with_empty<H: Hashable + Clone + PartialEq>(
});
let sibling = LocatedTree {
root_addr: sibling_addr,
root: Tree(Node::Nil),
root: Tree::empty(),
};
let (lroot, rroot) = if root.root_addr.is_left_child() {
(root, sibling)
Expand Down
10 changes: 6 additions & 4 deletions shardtree/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -424,7 +424,7 @@ impl<
l_addr.level(),
ann.clone(),
new_left,
Tree(Node::Nil),
Tree::empty(),
),
pos,
)
Expand All @@ -444,9 +444,7 @@ impl<
)
}
Node::Leaf { value: (h, r) } => Some((
Tree(Node::Leaf {
value: (h.clone(), *r | RetentionFlags::CHECKPOINT),
}),
Tree::leaf((h.clone(), *r | RetentionFlags::CHECKPOINT)),
root_addr.max_position(),
)),
Node::Nil | Node::Pruned => None,
Expand Down Expand Up @@ -822,6 +820,10 @@ impl<
(None, None) => unreachable!(),
};

// We don't use the `Tree::parent` constructor here, because it
// creates `Arc`s for the child nodes internally, but if we don't
// have a new child then we want to use the `Arc` for the existing
// child.
let new_parent = Tree(Node::Parent {
ann: new_left
.as_ref()
Expand Down
85 changes: 35 additions & 50 deletions shardtree/src/prunable.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -246,9 +246,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
(Tree(Node::Leaf { value: vl }), Tree(Node::Leaf { value: vr })) => {
if vl.0 == vr.0 {
// Merge the flags together.
Ok(Tree(Node::Leaf {
value: (vl.0, vl.1 | vr.1),
}))
Ok(Tree::leaf((vl.0, vl.1 | vr.1)))
} else {
trace!(left = ?vl.0, right = ?vr.0, "Merge conflict for leaves");
Err(addr)
Expand Down Expand Up @@ -314,26 +312,20 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
/// `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)) => Tree(Node::Nil),
(Tree(Node::Nil), Tree(Node::Nil)) => 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(
Node::Leaf {
value: (H::combine(level, &lv.0, &rv.0), rv.1),
},
)
Tree::leaf((H::combine(level, &lv.0, &rv.0), rv.1))
}
(left, right) => Tree(
Node::Parent {
(left, right) => Tree::parent(
ann,
left: Arc::new(left),
right: Arc::new(right),
},
left,
right,
),
}
}
Expand Down Expand Up @@ -514,7 +506,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// to the left to truncate the left child and then reconstruct the
// node with `Nil` as the right sibling
go(position, l_child, left.as_ref()).map(|left| {
Tree::unite(l_child.level(), ann.clone(), left, Tree(Node::Nil))
Tree::unite(l_child.level(), ann.clone(), left, Tree::empty())
})
} else {
// we are truncating within the range of the right node, so recurse
Expand Down Expand Up @@ -584,37 +576,32 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
// 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,
})
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 = if pruned {
Tree::empty_pruned()
} else {
Tree(Node::Parent {
ann: None,
left: full,
right: empty,
})
},
};
}
(node.root.reannotate_root(ann), incomplete)
};
Tree::empty()
};
let full = node.root;
node = LocatedTree {
root_addr: node.root_addr.parent(),
root: if node.root_addr.is_right_child() {
Tree::parent(None, empty, full)
} else {
Tree::parent(None, full, empty)
},
};
}
(node.root.reannotate_root(ann), incomplete)
};

match into {
Tree(Node::Nil) => Ok(replacement(None, subtree, false)),
Expand Down Expand Up @@ -936,16 +923,14 @@ 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 {
[(_, flags)] => Tree(Node::Leaf {
value: (h.clone(), *r & !*flags),
}),
[(_, flags)] => Tree::leaf((h.clone(), *r & !*flags)),
_ => {
panic!("Tree state inconsistent with checkpoints.");
}
}
}
Node::Nil => Tree(Node::Nil),
Node::Pruned => Tree(Node::Pruned),
Node::Nil => Tree::empty(),
Node::Pruned => Tree::empty_pruned(),
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions shardtree/src/store/memory.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use std::convert::{Infallible, TryFrom};
use incrementalmerkletree::Address;

use super::{Checkpoint, ShardStore};
use crate::{LocatedPrunableTree, LocatedTree, Node, PrunableTree, Tree};
use crate::{LocatedPrunableTree, LocatedTree, PrunableTree, Tree};

/// An implementation of [`ShardStore`] that stores all state in memory.
///
Expand Down Expand Up @@ -54,7 +54,7 @@ impl<H: Clone, C: Clone + Ord> ShardStore for MemoryShardStore<H, C> {
{
self.shards.push(LocatedTree {
root_addr: Address::from_parts(subtree_addr.level(), subtree_idx),
root: Tree(Node::Nil),
root: Tree::empty(),
})
}

Expand Down
16 changes: 6 additions & 10 deletions shardtree/src/testing.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,21 +31,17 @@ where
V::Value: Clone + 'static,
{
let leaf = prop_oneof![
Just(Tree(Node::Nil)),
arb_leaf.prop_map(|value| Tree(Node::Leaf { value }))
Just(Tree::empty()),
arb_leaf.prop_map(|value| Tree::leaf(value))

Check failure on line 35 in shardtree/src/testing.rs

View workflow job for this annotation

GitHub Actions / Clippy (MSRV)

redundant closure 

error: redundant closure
  --> shardtree/src/testing.rs:35:27
   |
35 |         arb_leaf.prop_map(|value| Tree::leaf(value))
   |                           ^^^^^^^^^^^^^^^^^^^^^^^^^ help: replace the closure with the function itself: `Tree::leaf`
   |
   = note: `-D clippy::redundant-closure` implied by `-D warnings`
   = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#redundant_closure

Check failure on line 35 in shardtree/src/testing.rs

View workflow job for this annotation

GitHub Actions / Clippy (MSRV)

redundant closure 

error: redundant closure
  --> shardtree/src/testing.rs:35:27
   |
35 |         arb_leaf.prop_map(|value| Tree::leaf(value))
   |                           ^^^^^^^^^^^^^^^^^^^^^^^^^ help: replace the closure with the function itself: `Tree::leaf`
   |
   = note: `-D clippy::redundant-closure` implied by `-D warnings`
   = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#redundant_closure
];

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
if left.is_nil() && right.is_nil() {
Tree::empty()
} else {
Node::Parent {
ann,
left: Arc::new(left),
right: Arc::new(right),
}
})
Tree::parent(ann, left, right)
}
})
})
}
Expand Down
10 changes: 4 additions & 6 deletions shardtree/src/tree.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,7 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
pub fn empty(root_addr: Address) -> Self {
Self {
root_addr,
root: Tree(Node::Nil),
root: Tree::empty(),
}
}

Expand All @@ -324,7 +324,7 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
pub fn with_root_value(root_addr: Address, value: V) -> Self {
Self {
root_addr,
root: Tree(Node::Leaf { value }),
root: Tree::leaf(value),
}
}

Expand Down Expand Up @@ -412,12 +412,10 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
pub(crate) mod tests {
use incrementalmerkletree::{Address, Level};

use super::{LocatedTree, Node, Tree};
use super::{LocatedTree, Tree};

pub(crate) fn str_leaf<A>(c: &str) -> Tree<A, String> {
Tree(Node::Leaf {
value: c.to_string(),
})
Tree::leaf(c.to_string())
}

pub(crate) fn nil<A, B>() -> Tree<A, B> {
Expand Down

0 comments on commit 49c2524

Please sign in to comment.