Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

shardtree: Use Tree constructors instead of struct creation #106

Merged
merged 1 commit into from
Jun 26, 2024
Merged

shardtree: Use Tree constructors instead of struct creation #106

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 5 additions & 11 deletions shardtree/src/batch.rs
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder why this doesn't format as a single line?

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
Open in desktop
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
17 changes: 5 additions & 12 deletions shardtree/src/testing.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,22 +30,15 @@ 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 }))
];
let leaf = prop_oneof![Just(Tree::empty()), arb_leaf.prop_map(Tree::leaf)];

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
Open in desktop
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
Loading