shardtree: Discard `REFERENCE` retention in leaf overwrites.

Also, disallow value-conflicted overwrites at the leaf level.

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

shardtree/src/prunable.rs

@@ -302,10 +302,11 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
         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<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
 
                 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)),
         );
     }
 

shardtree/src/tree.rs

@@ -117,6 +117,11 @@ impl<A, V> Tree<A, V> {
         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<A, V> Tree<A, V> {
             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, F: Fn(&V) -> B>(&self, f: &F) -> Tree<A, B>
+    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<B, E, F: Fn(&V) -> Result<B, E>>(&self, f: &F) -> Result<Tree<A, B>, 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<A, V> LocatedTree<A, V> {
             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, F: Fn(&V) -> B>(&self, f: &F) -> LocatedTree<A, B>
+    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<B, E, F: Fn(&V) -> Result<B, E>>(&self, f: &F) -> Result<LocatedTree<A, B>, E>
+    where
+        A: Clone,
+    {
+        Ok(LocatedTree {
+            root_addr: self.root_addr,
+            root: self.root.try_map(f)?,
+        })
+    }
 }
 
 impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {