diff --git a/shardtree/CHANGELOG.md b/shardtree/CHANGELOG.md index ec396f4..7dfcda1 100644 --- a/shardtree/CHANGELOG.md +++ b/shardtree/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to Rust's notion of ## Unreleased +### Added +- `shardtree::tree::Tree::{is_leaf, map, try_map}` +- `shardtree::tree::LocatedTree::{map, try_map}` + ## [0.3.1] - 2024-04-03 ### Fixed diff --git a/shardtree/src/prunable.rs b/shardtree/src/prunable.rs index b50e84f..62690be 100644 --- a/shardtree/src/prunable.rs +++ b/shardtree/src/prunable.rs @@ -302,10 +302,11 @@ impl PrunableTree { 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 { @@ -600,12 +601,45 @@ impl LocatedPrunableTree { match into { Tree(Node::Nil) => Ok(replacement(None, subtree)), - Tree(Node::Leaf { value: (value, _) }) => { + 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 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![])) + 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 a complete 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 @@ -932,7 +966,7 @@ mod tests { use super::{LocatedPrunableTree, PrunableTree, RetentionFlags}; use crate::{ - error::QueryError, + error::{InsertionError, QueryError}, tree::{ tests::{leaf, nil, parent}, LocatedTree, @@ -1101,21 +1135,16 @@ mod tests { 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)), ); } diff --git a/shardtree/src/tree.rs b/shardtree/src/tree.rs index 54ad57e..179ef01 100644 --- a/shardtree/src/tree.rs +++ b/shardtree/src/tree.rs @@ -117,6 +117,11 @@ impl Tree { Tree(self.0.reannotate(ann)) } + /// Returns `true` this is a [`Node::Leaf`], `false` otherwise. + pub fn is_leaf(&self) -> bool { + matches!(&self.0, Node::Leaf { .. }) + } + /// Returns `true` if no [`Node::Nil`] nodes are present in the tree, `false` otherwise. pub fn is_complete(&self) -> bool { match &self.0 { @@ -152,6 +157,41 @@ impl Tree { Node::Nil => 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>(&self, f: &F) -> Tree + where + A: Clone, + { + match &self.0 { + Node::Parent { ann, left, right } => Tree(Node::Parent { + ann: ann.clone(), + left: Arc::new(left.map(f)), + right: Arc::new(right.map(f)), + }), + Node::Leaf { value } => Tree(Node::Leaf { value: f(value) }), + Node::Nil => Tree(Node::Nil), + } + } + + /// 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 Result>(&self, f: &F) -> Result, 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, + })) + } } /// A binary Merkle tree with its root at the given address. @@ -240,6 +280,31 @@ impl LocatedTree { 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>(&self, f: &F) -> LocatedTree + 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 Result>(&self, f: &F) -> Result, E> + where + A: Clone, + { + Ok(LocatedTree { + root_addr: self.root_addr, + root: self.root.try_map(f)?, + }) + } } impl LocatedTree {