Move `LocatedTree` into `tree` submodule

Co-authored-by: Kris Nuttycombe <kris@nutty.land>
This commit is contained in:
Jack Grigg 2023-07-05 19:04:40 +00:00
parent 0f15a57cb5
commit 0d531d9309
2 changed files with 229 additions and 229 deletions

View File

@ -14,196 +14,11 @@ use incrementalmerkletree::{
use incrementalmerkletree::witness::IncrementalWitness;
mod tree;
pub use self::tree::{Node, Tree};
pub use self::tree::{LocatedTree, Node, Tree};
mod prunable;
pub use self::prunable::{PrunableTree, RetentionFlags};
/// A binary Merkle tree with its root at the given address.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LocatedTree<A, V> {
root_addr: Address,
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)
}
}
}
pub type LocatedPrunableTree<H> = LocatedTree<Option<Rc<H>>, (H, RetentionFlags)>;
/// A data structure describing the nature of a [`Node::Nil`] node in the tree that was introduced
@ -3327,7 +3142,7 @@ mod tests {
arb_char_str, arb_shardtree, check_shard_sizes, check_shardtree_insertion,
check_witness_with_pruned_subtrees,
},
tree::tests::{leaf, nil, parent, str_leaf},
tree::tests::{leaf, nil, parent},
InsertionError, LocatedPrunableTree, LocatedTree, MemoryShardStore, QueryError,
RetentionFlags, ShardTree,
};
@ -3411,46 +3226,6 @@ mod tests {
);
}
#[test]
fn located_tree() {
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,
}
]
);
}
#[test]
fn located_prunable_tree_insert() {
let tree = LocatedPrunableTree::empty(Address::from_parts(Level::from(2), 0));

View File

@ -1,7 +1,7 @@
use std::ops::Deref;
use std::rc::Rc;
use incrementalmerkletree::Address;
use incrementalmerkletree::{Address, Level, Position};
/// A "pattern functor" for a single layer of a binary tree.
#[derive(Clone, Debug, PartialEq, Eq)]
@ -146,11 +146,196 @@ impl<A, V> Tree<A, V> {
}
}
/// 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::{Node, Tree};
use super::{LocatedTree, Node, Tree};
pub(crate) fn str_leaf<A>(c: &str) -> Tree<A, String> {
Tree(Node::Leaf {
@ -193,4 +378,44 @@ pub(crate) mod tests {
]
);
}
#[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,
}
]
);
}
}