Merge pull request #105 from nuttycom/limitable_pruning

shardtree: Add `Reference` retention to make it possible to limit pruning.

incrementalmerkletree/CHANGELOG.md

@@ -7,6 +7,15 @@ and this project adheres to Rust's notion of
 
 ## Unreleased
 
+### Added
+- `incrementalmerkletree::Marking`
+
+### Changed
+- `incrementalmerkletree::Retention`
+  - Has added variant `Retention::Reference`
+  - `Retention::Checkpoint::is_marked` has been replaced by `Retention::Checkpoint::marking`
+    to permit checkpoints with `Reference` retention to be represented.
+
 ## [0.5.1] - 2024-03-25
 
 ### Added

incrementalmerkletree/src/lib.rs

@@ -62,35 +62,73 @@ pub mod witness;
 #[cfg_attr(docsrs, doc(cfg(feature = "test-dependencies")))]
 pub mod testing;
 
+/// An enumeration of the additional marking states that can be applied
+/// to leaves with [`Retention::Checkpoint`] retention.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub enum Marking {
+    /// A checkpoint with `Marked` marking will have `Marked` retention after `Checkpoint`
+    /// retention is removed.
+    Marked,
+    /// A checkpoint with `Reference` marking will have `Reference` retention after `Checkpoint`
+    /// retention is removed.
+    Reference,
+    /// A checkpoint with `None` marking will have `Ephemeral` retention after `Checkpoint`
+    /// retention is removed.
+    None,
+}
+
 /// A type for metadata that is used to determine when and how a leaf can be pruned from a tree.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub enum Retention<C> {
+    /// A leaf with `Ephemeral` retention will be pruned whenever its sibling is also a leaf with
+    /// `Ephemeral` retention.
     Ephemeral,
-    Checkpoint { id: C, is_marked: bool },
+    /// A leaf with `Checkpoint` retention will have its position retained in the tree
+    /// during pruning, but its value may be pruned (by merging with its sibling). If
+    /// `Checkpoint` retention is removed from the leaf, then the retention for the leaf will
+    /// become either `Ephemeral`, `Marked`, or `Reference` depending upon the value of the
+    /// `marking` field.
+    Checkpoint { id: C, marking: Marking },
+    /// A leaf with `Marked` retention will be retained in the tree during pruning. `Marked`
+    /// retention may only be explicitly removed.
     Marked,
+    /// A leaf with `Reference` retention will be retained in the tree during pruning. `Reference`
+    /// retention is removed whenever the associated leaf is overwritten with a tree node having
+    /// `Ephemeral`, `Checkpoint`, or `Marked` retention.
+    Reference,
 }
 
 impl<C> Retention<C> {
+    /// Returns whether the associated node has [`Retention::Checkpoint`] retention.
     pub fn is_checkpoint(&self) -> bool {
         matches!(self, Retention::Checkpoint { .. })
     }
 
+    /// Returns whether the associated node has [`Retention::Marked`] retention
+    /// or [`Retention::Checkpoint`] retention with [`Marking::Marked`] marking.
     pub fn is_marked(&self) -> bool {
-        match self {
+        matches!(
-            Retention::Ephemeral => false,
+            self,
-            Retention::Checkpoint { is_marked, .. } => *is_marked,
+            Retention::Marked
-            Retention::Marked => true,
+                | Retention::Checkpoint {
+                    marking: Marking::Marked,
+                    ..
                 }
+        )
     }
 
+    /// Applies the provided function to the checkpoint identifier, if any, and returns a new
+    /// `Retention` value having the same structure with the resulting value used as the checkpoint
+    /// identifier.
     pub fn map<'a, D, F: Fn(&'a C) -> D>(&'a self, f: F) -> Retention<D> {
         match self {
             Retention::Ephemeral => Retention::Ephemeral,
-            Retention::Checkpoint { id, is_marked } => Retention::Checkpoint {
+            Retention::Checkpoint { id, marking } => Retention::Checkpoint {
                 id: f(id),
-                is_marked: *is_marked,
+                marking: *marking,
             },
             Retention::Marked => Retention::Marked,
+            Retention::Reference => Retention::Reference,
         }
     }
 }

incrementalmerkletree/src/testing.rs

@@ -5,7 +5,7 @@ use core::marker::PhantomData;
 use proptest::prelude::*;
 use std::collections::BTreeSet;
 
-use crate::{Hashable, Level, Position, Retention};
+use crate::{Hashable, Level, Marking, Position, Retention};
 
 pub mod complete_tree;
 
@@ -208,10 +208,22 @@ impl<H: Hashable + Clone, C: Clone> Operation<H, C> {
     }
 }
 
+/// Returns a strategy for creating a uniformly-distributed [`Marking`]
+/// value.
+pub fn arb_marking() -> impl Strategy<Value = Marking> {
+    prop_oneof![
+        Just(Marking::Marked),
+        Just(Marking::Reference),
+        Just(Marking::None)
+    ]
+}
+
+/// Returns a strategy for creating a uniformly-distributed [`Retention`]
+/// value.
 pub fn arb_retention() -> impl Strategy<Value = Retention<()>> {
     prop_oneof![
         Just(Retention::Ephemeral),
-        any::<bool>().prop_map(|is_marked| Retention::Checkpoint { id: (), is_marked }),
+        arb_marking().prop_map(|marking| Retention::Checkpoint { id: (), marking }),
         Just(Retention::Marked),
     ]
 }
@@ -557,7 +569,7 @@ pub fn check_root_hashes<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: F
             0,
             Retention::Checkpoint {
                 id: 1,
-                is_marked: true,
+                marking: Marking::Marked,
             },
         );
         for _ in 0..3 {
@@ -735,7 +747,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
             0,
             Checkpoint {
                 id: 1,
-                is_marked: true,
+                marking: Marking::Marked,
             },
         );
         assert!(tree.rewind());
@@ -769,7 +781,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
             6,
             Checkpoint {
                 id: 1,
-                is_marked: true,
+                marking: Marking::Marked,
             },
         );
         tree.assert_append(7, Ephemeral);
@@ -838,7 +850,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
                 H::from_u64(3),
                 Checkpoint {
                     id: C::from_u64(1),
-                    is_marked: true,
+                    marking: Marking::Marked,
                 },
             ),
             Append(H::from_u64(4), Marked),
@@ -847,21 +859,21 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
                 H::from_u64(5),
                 Checkpoint {
                     id: C::from_u64(3),
-                    is_marked: false,
+                    marking: Marking::None,
                 },
             ),
             Append(
                 H::from_u64(6),
                 Checkpoint {
                     id: C::from_u64(4),
-                    is_marked: false,
+                    marking: Marking::None,
                 },
             ),
             Append(
                 H::from_u64(7),
                 Checkpoint {
                     id: C::from_u64(5),
-                    is_marked: false,
+                    marking: Marking::None,
                 },
             ),
             Witness(3u64.into(), 5),
@@ -890,7 +902,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
                 H::from_u64(0),
                 Checkpoint {
                     id: C::from_u64(1),
-                    is_marked: true,
+                    marking: Marking::Marked,
                 },
             ),
             Append(H::from_u64(0), Ephemeral),
@@ -900,7 +912,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
                 H::from_u64(0),
                 Checkpoint {
                     id: C::from_u64(2),
-                    is_marked: false,
+                    marking: Marking::None,
                 },
             ),
             Append(H::from_u64(0), Ephemeral),
@@ -939,7 +951,7 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
                 H::from_u64(0),
                 Checkpoint {
                     id: C::from_u64(4),
-                    is_marked: false,
+                    marking: Marking::None,
                 },
             ),
             Rewind,
@@ -958,14 +970,14 @@ pub fn check_witnesses<H: TestHashable, C: TestCheckpoint, T: Tree<H, C>, F: Fn(
                 H::from_u64(0),
                 Checkpoint {
                     id: C::from_u64(1),
-                    is_marked: true,
+                    marking: Marking::Marked,
                 },
             ),
             Append(
                 H::from_u64(0),
                 Checkpoint {
                     id: C::from_u64(4),
-                    is_marked: false,
+                    marking: Marking::None,
                 },
             ),
             Witness(Position(2), 2),
@@ -1034,7 +1046,7 @@ pub fn check_remove_mark<C: TestCheckpoint, T: Tree<String, C>, F: Fn(usize) ->
                 "a",
                 Retention::Checkpoint {
                     id: C::from_u64(1),
-                    is_marked: true,
+                    marking: Marking::Marked,
                 },
             ),
             witness(1, 1),

incrementalmerkletree/src/testing/complete_tree.rs

@@ -2,6 +2,7 @@
 use std::cmp::min;
 use std::collections::{BTreeMap, BTreeSet};
 
+use crate::Marking;
 use crate::{testing::Tree, Hashable, Level, Position, Retention};
 
 const MAX_COMPLETE_SIZE_ERROR: &str = "Positions of a `CompleteTree` must fit into the platform word size, because larger complete trees are not representable.";
@@ -104,11 +105,11 @@ impl<H: Hashable, C: Clone + Ord + core::fmt::Debug, const DEPTH: u8> CompleteTr
                 append(&mut self.leaves, value, DEPTH)?;
                 self.mark();
             }
-            Retention::Checkpoint { id, is_marked } => {
+            Retention::Checkpoint { id, marking } => {
                 let latest_checkpoint = self.checkpoints.keys().rev().next();
                 if Some(&id) > latest_checkpoint {
                     append(&mut self.leaves, value, DEPTH)?;
-                    if is_marked {
+                    if marking == Marking::Marked {
                         self.mark();
                     }
                     self.checkpoint(id, self.current_position());
@@ -119,7 +120,7 @@ impl<H: Hashable, C: Clone + Ord + core::fmt::Debug, const DEPTH: u8> CompleteTr
                     });
                 }
             }
-            Retention::Ephemeral => {
+            Retention::Ephemeral | Retention::Reference => {
                 append(&mut self.leaves, value, DEPTH)?;
             }
         }

shardtree/CHANGELOG.md

@@ -7,6 +7,24 @@ and this project adheres to Rust's notion of
 
 ## Unreleased
 
+### Added
+- `shardtree::tree::Tree::{is_leaf, map, try_map, empty_pruned}`
+- `shardtree::tree::LocatedTree::{map, try_map}`
+- `shardtree::prunable::PrunableTree::{has_computable_root}`
+
+### Changed
+- `shardtree::tree::Node` has additional variant `Node::Pruned`.
+
+### Removed
+- `shardtree::tree::Tree::is_complete` as it is no longer well-defined in the
+  presence of `Pruned` nodes. Use `PrunableTree::has_computable_root` to
+  determine whether it is possible to compute the root of a tree.
+
+### Fixed
+- Fixes an error that could occur if an inserted `Frontier` node was
+  interpreted as a node that had actually had its value observed as though it
+  had been inserted using the ordinary tree insertion methods.
+
 ## [0.3.1] - 2024-04-03
 
 ### Fixed

shardtree/src/legacy.rs

@@ -1,6 +1,8 @@
 use std::fmt;
 
-use incrementalmerkletree::{witness::IncrementalWitness, Address, Hashable, Level, Retention};
+use incrementalmerkletree::{
+    witness::IncrementalWitness, Address, Hashable, Level, Marking, Retention,
+};
 
 use crate::{
     store::ShardStore, InsertionError, LocatedPrunableTree, LocatedTree, PrunableTree,
@@ -198,7 +200,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                     c.ommers_iter().cloned(),
                     &Retention::Checkpoint {
                         id: checkpoint_id,
-                        is_marked: false,
+                        marking: Marking::None,
                     },
                     self.root_addr.level(),
                 )

shardtree/src/lib.rs

@@ -25,9 +25,10 @@
 use core::fmt::Debug;
 use either::Either;
 use incrementalmerkletree::frontier::Frontier;
+use incrementalmerkletree::Marking;
 use std::collections::{BTreeMap, BTreeSet};
 use std::sync::Arc;
-use tracing::trace;
+use tracing::{debug, trace};
 
 use incrementalmerkletree::{
     frontier::NonEmptyFrontier, Address, Hashable, Level, MerklePath, Position, Retention,
@@ -233,16 +234,17 @@ impl<
 
         let (append_result, position, checkpoint_id) =
             if let Some(subtree) = self.store.last_shard().map_err(ShardTreeError::Storage)? {
-                if subtree.root.is_complete() {
+                match subtree.max_position() {
+                    // If the subtree is full, then construct a successor tree.
+                    Some(pos) if pos == subtree.root_addr.max_position() => {
                         let addr = subtree.root_addr;
-
+                        if subtree.root_addr.index() < Self::max_subtree_index() {
-                    if addr.index() < Self::max_subtree_index() {
                             LocatedTree::empty(addr.next_at_level()).append(value, retention)?
                         } else {
                             return Err(InsertionError::TreeFull.into());
                         }
-                } else {
+                    }
-                    subtree.append(value, retention)?
+                    _ => subtree.append(value, retention)?,
                 }
             } else {
                 let root_addr = Address::from_parts(Self::subtree_level(), 0);
@@ -271,6 +273,7 @@ impl<
     ///
     /// This method may be used to add a checkpoint for the empty tree; note that
     /// [`Retention::Marked`] is invalid for the empty tree.
+    #[tracing::instrument(skip(self, frontier, leaf_retention))]
     pub fn insert_frontier(
         &mut self,
         frontier: Frontier<H, DEPTH>,
@@ -280,18 +283,19 @@ impl<
             self.insert_frontier_nodes(nonempty_frontier, leaf_retention)
         } else {
             match leaf_retention {
-                Retention::Ephemeral => Ok(()),
+                Retention::Ephemeral | Retention::Reference => Ok(()),
                 Retention::Checkpoint {
                     id,
-                    is_marked: false,
+                    marking: Marking::None | Marking::Reference,
                 } => self
                     .store
                     .add_checkpoint(id, Checkpoint::tree_empty())
                     .map_err(ShardTreeError::Storage),
-                Retention::Checkpoint {
+                Retention::Marked
-                    is_marked: true, ..
+                | Retention::Checkpoint {
-                }
+                    marking: Marking::Marked,
-                | Retention::Marked => Err(ShardTreeError::Insert(
+                    ..
+                } => Err(ShardTreeError::Insert(
                     InsertionError::MarkedRetentionInvalid,
                 )),
             }
@@ -301,6 +305,7 @@ impl<
     /// Add the leaf and ommers of the provided frontier as nodes within the subtree corresponding
     /// to the frontier's position, and update the cap to include the ommer nodes at levels greater
     /// than or equal to the shard height.
+    #[tracing::instrument(skip(self, frontier, leaf_retention))]
     pub fn insert_frontier_nodes(
         &mut self,
         frontier: NonEmptyFrontier<H>,
@@ -308,15 +313,24 @@ impl<
     ) -> Result<(), ShardTreeError<S::Error>> {
         let leaf_position = frontier.position();
         let subtree_root_addr = Address::above_position(Self::subtree_level(), leaf_position);
-        trace!("Subtree containing nodes: {:?}", subtree_root_addr);
 
-        let (updated_subtree, supertree) = self
+        let current_shard = self
             .store
             .get_shard(subtree_root_addr)
             .map_err(ShardTreeError::Storage)?
-            .unwrap_or_else(|| LocatedTree::empty(subtree_root_addr))
+            .unwrap_or_else(|| LocatedTree::empty(subtree_root_addr));
-            .insert_frontier_nodes(frontier, &leaf_retention)?;
 
+        trace!(
+            max_position = ?current_shard.max_position(),
+            subtree = ?current_shard,
+            "Current shard");
+        let (updated_subtree, supertree) =
+            current_shard.insert_frontier_nodes(frontier, &leaf_retention)?;
+        trace!(
+            max_position = ?updated_subtree.max_position(),
+            subtree = ?updated_subtree,
+            "Replacement shard"
+        );
         self.store
             .put_shard(updated_subtree)
             .map_err(ShardTreeError::Storage)?;
@@ -333,7 +347,7 @@ impl<
                 .map_err(ShardTreeError::Storage)?;
         }
 
-        if let Retention::Checkpoint { id, is_marked: _ } = leaf_retention {
+        if let Retention::Checkpoint { id, .. } = leaf_retention {
             trace!("Adding checkpoint {:?} at {:?}", id, leaf_position);
             self.store
                 .add_checkpoint(id, Checkpoint::at_position(leaf_position))
@@ -346,6 +360,7 @@ impl<
 
     /// Insert a tree by decomposing it into its `SHARD_HEIGHT` or smaller parts (if necessary)
     /// and inserting those at their appropriate locations.
+    #[tracing::instrument(skip(self, tree, checkpoints))]
     pub fn insert_tree(
         &mut self,
         tree: LocatedPrunableTree<H>,
@@ -359,6 +374,7 @@ impl<
             // for some inputs, and given that it is always correct to not insert an empty
             // subtree into `self`, we maintain the invariant by skipping empty subtrees.
             if subtree.root().is_empty() {
+                debug!("Subtree with root {:?} is empty.", subtree.root_addr);
                 continue;
             }
 
@@ -368,14 +384,15 @@ impl<
             let root_addr = Self::subtree_addr(subtree.root_addr.position_range_start());
 
             let contains_marked = subtree.root.contains_marked();
-            let (new_subtree, mut incomplete) = self
+            let current_shard = self
                 .store
                 .get_shard(root_addr)
                 .map_err(ShardTreeError::Storage)?
-                .unwrap_or_else(|| LocatedTree::empty(root_addr))
+                .unwrap_or_else(|| LocatedTree::empty(root_addr));
-                .insert_subtree(subtree, contains_marked)?;
+            let (replacement_shard, mut incomplete) =
+                current_shard.insert_subtree(subtree, contains_marked)?;
             self.store
-                .put_shard(new_subtree)
+                .put_shard(replacement_shard)
                 .map_err(ShardTreeError::Storage)?;
             all_incomplete.append(&mut incomplete);
         }
@@ -396,8 +413,8 @@ impl<
             root_addr: Address,
             root: &PrunableTree<H>,
         ) -> Option<(PrunableTree<H>, Position)> {
-            match root {
+            match &root.0 {
-                Tree(Node::Parent { ann, left, right }) => {
+                Node::Parent { ann, left, right } => {
                     let (l_addr, r_addr) = root_addr.children().unwrap();
                     go(r_addr, right).map_or_else(
                         || {
@@ -426,13 +443,13 @@ impl<
                         },
                     )
                 }
-                Tree(Node::Leaf { value: (h, r) }) => Some((
+                Node::Leaf { value: (h, r) } => Some((
                     Tree(Node::Leaf {
                         value: (h.clone(), *r | RetentionFlags::CHECKPOINT),
                     }),
                     root_addr.max_position(),
                 )),
-                Tree(Node::Nil) => None,
+                Node::Nil | Node::Pruned => None,
             }
         }
 
@@ -478,6 +495,7 @@ impl<
         Ok(true)
     }
 
+    #[tracing::instrument(skip(self))]
     fn prune_excess_checkpoints(&mut self) -> Result<(), ShardTreeError<S::Error>> {
         let checkpoint_count = self
             .store
@@ -559,12 +577,19 @@ impl<
 
             // Prune each affected subtree
             for (subtree_addr, positions) in clear_positions.into_iter() {
-                let cleared = self
+                let to_clear = self
                     .store
                     .get_shard(subtree_addr)
-                    .map_err(ShardTreeError::Storage)?
+                    .map_err(ShardTreeError::Storage)?;
-                    .map(|subtree| subtree.clear_flags(positions));
+
-                if let Some(cleared) = cleared {
+                if let Some(to_clear) = to_clear {
+                    let pre_clearing_max_position = to_clear.max_position();
+                    let cleared = to_clear.clear_flags(positions);
+
+                    // Clearing flags should not modify the max position of leaves represented
+                    // in the shard.
+                    assert!(cleared.max_position() == pre_clearing_max_position);
+
                     self.store
                         .put_shard(cleared)
                         .map_err(ShardTreeError::Storage)?;
@@ -740,8 +765,8 @@ impl<
         // roots.
         truncate_at: Position,
     ) -> Result<(H, Option<PrunableTree<H>>), ShardTreeError<S::Error>> {
-        match &cap.root {
+        match &cap.root.0 {
-            Tree(Node::Parent { ann, left, right }) => {
+            Node::Parent { ann, left, right } => {
                 match ann {
                     Some(cached_root) if target_addr.contains(&cap.root_addr) => {
                         Ok((cached_root.as_ref().clone(), None))
@@ -815,7 +840,7 @@ impl<
                     }
                 }
             }
-            Tree(Node::Leaf { value }) => {
+            Node::Leaf { value } => {
                 if truncate_at >= cap.root_addr.position_range_end()
                     && target_addr.contains(&cap.root_addr)
                 {
@@ -840,7 +865,7 @@ impl<
                     ))
                 }
             }
-            Tree(Node::Nil) => {
+            Node::Nil | Node::Pruned => {
                 if cap.root_addr == target_addr
                     || cap.root_addr.level() == ShardTree::<S, DEPTH, SHARD_HEIGHT>::subtree_level()
                 {
@@ -850,7 +875,7 @@ impl<
                     Ok((
                         root.clone(),
                         if truncate_at >= cap.root_addr.position_range_end() {
-                            // return the compute root as a new leaf to be cached if it contains no
+                            // return the computed root as a new leaf to be cached if it contains no
                             // empty hashes due to truncation
                             Some(Tree::leaf((root, RetentionFlags::EPHEMERAL)))
                         } else {
@@ -1284,21 +1309,24 @@ impl<
 
 #[cfg(test)]
 mod tests {
+    use std::convert::Infallible;
+
     use assert_matches::assert_matches;
     use proptest::prelude::*;
 
     use incrementalmerkletree::{
-        frontier::NonEmptyFrontier,
+        frontier::{Frontier, NonEmptyFrontier},
         testing::{
             arb_operation, check_append, check_checkpoint_rewind, check_operations,
             check_remove_mark, check_rewind_remove_mark, check_root_hashes,
             check_witness_consistency, check_witnesses, complete_tree::CompleteTree, CombinedTree,
             SipHashable,
         },
-        Address, Hashable, Level, Position, Retention,
+        Address, Hashable, Level, Marking, MerklePath, Position, Retention,
     };
 
     use crate::{
+        error::{QueryError, ShardTreeError},
         store::memory::MemoryShardStore,
         testing::{
             arb_char_str, arb_shardtree, check_shard_sizes, check_shardtree_insertion,
@@ -1396,13 +1424,102 @@ mod tests {
                 'd'.to_string(),
                 Retention::Checkpoint {
                     id: 11,
-                    is_marked: false
+                    marking: Marking::None
                 },
             ),
             Ok(()),
         );
     }
 
+    #[test]
+    fn avoid_pruning_reference() {
+        fn test_with_marking(
+            frontier_marking: Marking,
+        ) -> Result<MerklePath<String, 6>, ShardTreeError<Infallible>> {
+            let mut tree = ShardTree::<MemoryShardStore<String, usize>, 6, 3>::new(
+                MemoryShardStore::empty(),
+                5,
+            );
+
+            let frontier_end = Position::from((1 << 3) - 3);
+            let mut f0 = Frontier::<String, 6>::empty();
+            for c in 'a'..='f' {
+                f0.append(c.to_string());
+            }
+
+            let frontier = Frontier::from_parts(
+                frontier_end,
+                "f".to_owned(),
+                vec!["e".to_owned(), "abcd".to_owned()],
+            )
+            .unwrap();
+
+            // Insert a frontier two leaves from the end of the first shard, checkpointed,
+            // with the specified marking.
+            tree.insert_frontier(
+                frontier,
+                Retention::Checkpoint {
+                    id: 1,
+                    marking: frontier_marking,
+                },
+            )?;
+
+            // Insert a few leaves beginning at the subsequent position, so as to cross the shard
+            // boundary.
+            tree.batch_insert(
+                frontier_end + 1,
+                ('g'..='j')
+                    .into_iter()
+                    .map(|c| (c.to_string(), Retention::Ephemeral)),
+            )?;
+
+            // Trigger pruning by adding 5 more checkpoints
+            for i in 2..7 {
+                tree.checkpoint(i).unwrap();
+            }
+
+            // Insert nodes that require the pruned nodes for witnessing
+            tree.batch_insert(
+                frontier_end - 1,
+                ('e'..='f')
+                    .into_iter()
+                    .map(|c| (c.to_string(), Retention::Marked)),
+            )?;
+
+            // Compute the witness
+            tree.witness_at_checkpoint_id(frontier_end, &6)
+        }
+
+        // If we insert the frontier with Marking::None, the frontier nodes are treated
+        // as ephemeral nodes and are pruned, leaving an incomplete tree.
+        assert_matches!(
+            test_with_marking(Marking::None),
+            Err(ShardTreeError::Query(QueryError::TreeIncomplete(_)))
+        );
+
+        // If we insert the frontier with Marking::Reference, the frontier nodes will
+        // not be pruned on completion of the subtree, and thus we'll be able to compute
+        // the witness.
+        let expected_witness = MerklePath::from_parts(
+            [
+                "e",
+                "gh",
+                "abcd",
+                "ij______",
+                "________________",
+                "________________________________",
+            ]
+            .iter()
+            .map(|s| s.to_string())
+            .collect(),
+            Position::from(5),
+        )
+        .unwrap();
+
+        let witness = test_with_marking(Marking::Reference).unwrap();
+        assert_eq!(witness, expected_witness);
+    }
+
     // Combined tree tests
     #[allow(clippy::type_complexity)]
     fn new_combined_tree<H: Hashable + Ord + Clone + core::fmt::Debug>(

shardtree/src/prunable.rs

@@ -4,10 +4,11 @@ use std::collections::{BTreeMap, BTreeSet};
 use std::sync::Arc;
 
 use bitflags::bitflags;
+use incrementalmerkletree::Marking;
 use incrementalmerkletree::{
     frontier::NonEmptyFrontier, Address, Hashable, Level, Position, Retention,
 };
-use tracing::trace;
+use tracing::{trace, warn};
 
 use crate::error::{InsertionError, QueryError};
 use crate::{LocatedTree, Node, Tree};
@@ -30,6 +31,12 @@ bitflags! {
         /// A leaf with `MARKED` retention can be pruned only as a consequence of an explicit deletion
         /// action.
         const MARKED = 0b00000010;
+
+        /// A leaf with `REFERENCE` retention will not be considered prunable until the `REFERENCE`
+        /// flag is removed from the leaf. The `REFERENCE` flag will be removed at any point that
+        /// the leaf is overwritten without `REFERENCE` retention, and `REFERENCE` retention cannot
+        /// be added to an existing leaf.
+        const REFERENCE = 0b00000100;
     }
 }
 
@@ -47,14 +54,17 @@ impl<'a, C> From<&'a Retention<C>> for RetentionFlags {
     fn from(retention: &'a Retention<C>) -> Self {
         match retention {
             Retention::Ephemeral => RetentionFlags::EPHEMERAL,
-            Retention::Checkpoint { is_marked, .. } => {
+            Retention::Checkpoint {
-                if *is_marked {
+                marking: Marking::Marked,
-                    RetentionFlags::CHECKPOINT | RetentionFlags::MARKED
+                ..
-                } else {
+            } => RetentionFlags::CHECKPOINT | RetentionFlags::MARKED,
-                    RetentionFlags::CHECKPOINT
+            Retention::Checkpoint {
-                }
+                marking: Marking::Reference,
-            }
+                ..
+            } => RetentionFlags::CHECKPOINT | RetentionFlags::REFERENCE,
+            Retention::Checkpoint { .. } => RetentionFlags::CHECKPOINT,
             Retention::Marked => RetentionFlags::MARKED,
+            Retention::Reference => RetentionFlags::REFERENCE,
         }
     }
 }
@@ -90,12 +100,25 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
             .map_or(false, |(_, retention)| retention.is_marked())
     }
 
+    /// Returns `true` if it is possible to compute or retrieve the Merkle root of this
+    /// tree.
+    pub fn has_computable_root(&self) -> bool {
+        match &self.0 {
+            Node::Parent { ann, left, right } => {
+                ann.is_some()
+                    || (left.as_ref().has_computable_root() && right.as_ref().has_computable_root())
+            }
+            Node::Leaf { .. } => true,
+            Node::Nil | Node::Pruned => false,
+        }
+    }
+
     /// Determines whether a tree has any [`Retention::Marked`] nodes.
     pub fn contains_marked(&self) -> bool {
         match &self.0 {
             Node::Parent { left, right, .. } => left.contains_marked() || right.contains_marked(),
             Node::Leaf { value: (_, r) } => r.is_marked(),
-            Node::Nil => false,
+            Node::Nil | Node::Pruned => false,
         }
     }
 
@@ -112,8 +135,8 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
             // so no need to inspect the tree
             Ok(H::empty_root(root_addr.level()))
         } else {
-            match self {
+            match &self.0 {
-                Tree(Node::Parent { ann, left, right }) => ann
+                Node::Parent { ann, left, right } => ann
                     .as_ref()
                     .filter(|_| truncate_at >= root_addr.position_range_end())
                     .map_or_else(
@@ -135,7 +158,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
                             Ok(rc.as_ref().clone())
                         },
                     ),
-                Tree(Node::Leaf { value }) => {
+                Node::Leaf { value } => {
                     if truncate_at >= root_addr.position_range_end() {
                         // no truncation of this leaf is necessary, just use it
                         Ok(value.0.clone())
@@ -147,7 +170,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
                         Err(vec![root_addr])
                     }
                 }
-                Tree(Node::Nil) => Err(vec![root_addr]),
+                Node::Nil | Node::Pruned => Err(vec![root_addr]),
             }
         }
     }
@@ -182,7 +205,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
                 }
                 result
             }
-            Node::Nil => BTreeSet::new(),
+            Node::Nil | Node::Pruned => BTreeSet::new(),
         }
     }
 
@@ -206,6 +229,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
     /// The merge operation is checked to be strictly additive and returns an error if merging
     /// would cause information loss or if a conflict between root hashes occurs at a node. The
     /// returned error contains the address of the node where such a conflict occurred.
+    #[tracing::instrument()]
     pub fn merge_checked(self, root_addr: Address, other: Self) -> Result<Self, Address> {
         #[allow(clippy::type_complexity)]
         fn go<H: Hashable + Clone + PartialEq>(
@@ -218,6 +242,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
             let no_default_fill = addr.position_range_end();
             match (t0, t1) {
                 (Tree(Node::Nil), other) | (other, Tree(Node::Nil)) => Ok(other),
+                (Tree(Node::Pruned), other) | (other, Tree(Node::Pruned)) => Ok(other),
                 (Tree(Node::Leaf { value: vl }), Tree(Node::Leaf { value: vr })) => {
                     if vl.0 == vr.0 {
                         // Merge the flags together.
@@ -279,12 +304,11 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
             }
         }
 
-        trace!(this = ?self, other = ?other, "Merging subtrees");
         go(root_addr, self, other)
     }
 
     /// Unite two nodes by either constructing a new parent node, or, if both nodes are ephemeral
-    /// leaves or Nil, constructing a replacement root by hashing leaf values together (or a
+    /// leaves or Nil, constructing a replacement parent by hashing leaf values together (or a
     /// replacement `Nil` value).
     ///
     /// `level` must be the level of the two nodes that are being joined.
@@ -292,10 +316,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
+                // we can prune right-hand leaves that are not marked or reference leaves; if a
-                // is a checkpoint then that information will be propagated to
+                // leaf is a checkpoint then that information will be propagated to the replacement
-                // the replacement leaf
+                // leaf
-                if lv.1 == RetentionFlags::EPHEMERAL && (rv.1 & RetentionFlags::MARKED) == RetentionFlags::EPHEMERAL =>
+                if lv.1 == RetentionFlags::EPHEMERAL &&
+                    (rv.1 & (RetentionFlags::MARKED | RetentionFlags::REFERENCE)) == RetentionFlags::EPHEMERAL =>
             {
                 Tree(
                     Node::Leaf {
@@ -507,7 +532,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                         None
                     }
                 }
-                Node::Nil => None,
+                Node::Nil | Node::Pruned => None,
             }
         }
 
@@ -544,13 +569,24 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
             root_addr: Address,
             into: &PrunableTree<H>,
             subtree: LocatedPrunableTree<H>,
-            is_complete: bool,
             contains_marked: bool,
         ) -> Result<(PrunableTree<H>, Vec<IncompleteAt>), InsertionError> {
+            trace!(
+                root_addr = ?root_addr,
+                max_position = ?LocatedTree::max_position_internal(root_addr, into),
+                to_insert = ?subtree.root_addr(),
+                "Subtree insert"
+            );
+            // An empty tree cannot replace any other type of tree.
+            if subtree.root().is_nil() {
+                Ok((into.clone(), vec![]))
+            } else {
                 // 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>| {
+                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() {
@@ -558,19 +594,21 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                             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: Arc::new(Tree(Node::Nil)),
+                                    left: empty,
-                                right: Arc::new(node.root),
+                                    right: full,
                                 })
                             } else {
                                 Tree(Node::Parent {
                                     ann: None,
-                                left: Arc::new(node.root),
+                                    left: full,
-                                right: Arc::new(Tree(Node::Nil)),
+                                    right: empty,
                                 })
                             },
                         };
@@ -578,19 +616,48 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                     (node.root.reannotate_root(ann), incomplete)
                 };
 
-            trace!(
-                "Node at {:?} contains subtree at {:?}",
-                root_addr,
-                subtree.root_addr(),
-            );
                 match into {
-                Tree(Node::Nil) => Ok(replacement(None, subtree)),
+                    Tree(Node::Nil) => Ok(replacement(None, subtree, false)),
-                Tree(Node::Leaf { value: (value, _) }) => {
+                    Tree(Node::Pruned) => Ok(replacement(None, subtree, true)),
+                    Tree(Node::Leaf {
+                        value: (value, retention),
+                    }) => {
                         if root_addr == subtree.root_addr {
-                        if is_complete {
+                            // 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 subtree.root.has_computable_root() {
+                                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.
+                                        // can always recompute the root from the subtree.
-                            Ok((subtree.root, vec![]))
+                                        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
@@ -598,7 +665,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                                     vec![],
                                 ))
                             } else {
-                            trace!(
+                                warn!(
                                     cur_root = ?value,
                                     new_root = ?subtree.root.node_value(),
                                     "Insertion conflict",
@@ -606,7 +673,22 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                                 Err(InsertionError::Conflict(root_addr))
                             }
                         } else {
-                        Ok(replacement(Some(Arc::new(value.clone())), subtree))
+                            Ok(replacement(
+                                Some(Arc::new(value.clone())),
+                                subtree,
+                                // The subtree being inserted may have its root at some level lower
+                                // than the next level down. The siblings of nodes that will be
+                                // generated while descending to the subtree root level will be
+                                // `Nil` nodes (indicating that the value of these nodes have never
+                                // been observed) if the leaf being replaced has `REFERENCE`
+                                // retention. Any other leaf without `REFERENCE` retention will
+                                // have been produced by pruning of previously observed node
+                                // values, so in those cases we use `Pruned` nodes for the absent
+                                // siblings. This allows us to retain the distinction between what
+                                // parts of the tree have been directly observed and what parts
+                                // have not.
+                                !retention.contains(RetentionFlags::REFERENCE),
+                            ))
                         }
                     }
                     parent if root_addr == subtree.root_addr => {
@@ -624,7 +706,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                         let (l_addr, r_addr) = root_addr.children().unwrap();
                         if l_addr.contains(&subtree.root_addr) {
                             let (new_left, incomplete) =
-                            go(l_addr, left.as_ref(), subtree, is_complete, contains_marked)?;
+                                go(l_addr, left.as_ref(), subtree, contains_marked)?;
                             Ok((
                                 Tree::unite(
                                     root_addr.level() - 1,
@@ -635,13 +717,8 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                                 incomplete,
                             ))
                         } else {
-                        let (new_right, incomplete) = go(
+                            let (new_right, incomplete) =
-                            r_addr,
+                                go(r_addr, right.as_ref(), subtree, contains_marked)?;
-                            right.as_ref(),
-                            subtree,
-                            is_complete,
-                            contains_marked,
-                        )?;
                             Ok((
                                 Tree::unite(
                                     root_addr.level() - 1,
@@ -655,18 +732,24 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                     }
                 }
             }
+        }
 
+        let max_position = self.max_position();
+        trace!(
+            max_position = ?max_position,
+            tree = ?self,
+            to_insert = ?subtree,
+            "Current shard"
+        );
         let LocatedTree { root_addr, root } = self;
         if root_addr.contains(&subtree.root_addr) {
-            let complete = subtree.root.is_complete();
+            go(*root_addr, root, subtree, contains_marked).map(|(root, incomplete)| {
-            go(*root_addr, root, subtree, complete, contains_marked).map(|(root, incomplete)| {
+                let new_tree = LocatedTree {
-                (
-                    LocatedTree {
                     root_addr: *root_addr,
                     root,
-                    },
+                };
-                    incomplete,
+                assert!(new_tree.max_position() >= max_position);
-                )
+                (new_tree, incomplete)
             })
         } else {
             Err(InsertionError::NotContained(subtree.root_addr))
@@ -725,9 +808,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
         split_at: Level,
     ) -> (Self, Option<Self>) {
         let mut addr = Address::from(position);
-        let mut subtree = Tree(Node::Leaf {
+        let mut subtree = Tree::leaf((leaf, leaf_retention.into()));
-            value: (leaf, leaf_retention.into()),
-        });
 
         while addr.level() < split_at {
             if addr.is_left_child() {
@@ -829,13 +910,13 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                 // nothing to do, so we just return the root
                 root.clone()
             } else {
-                match root {
+                match &root.0 {
-                    Tree(Node::Parent { ann, left, right }) => {
+                    Node::Parent { ann, left, right } => {
                         let (l_addr, r_addr) = root_addr.children().unwrap();
 
                         let p = to_clear.partition_point(|(p, _)| p < &l_addr.position_range_end());
                         trace!(
-                            "In {:?}, partitioned: {:?} {:?}",
+                            "Tree::unite at {:?}, partitioned: {:?} {:?}",
                             root_addr,
                             &to_clear[0..p],
                             &to_clear[p..],
@@ -847,7 +928,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                             go(&to_clear[p..], r_addr, right),
                         )
                     }
-                    Tree(Node::Leaf { value: (h, r) }) => {
+                    Node::Leaf { value: (h, r) } => {
                         trace!("In {:?}, clearing {:?}", root_addr, to_clear);
                         // When we reach a leaf, we should be down to just a single position
                         // which should correspond to the last level-0 child of the address's
@@ -855,18 +936,16 @@ 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 {
-                            [(pos, flags)] => {
+                            [(_, flags)] => Tree(Node::Leaf {
-                                assert_eq!(*pos, root_addr.max_position());
-                                Tree(Node::Leaf {
                                 value: (h.clone(), *r & !*flags),
-                                })
+                            }),
-                            }
                             _ => {
                                 panic!("Tree state inconsistent with checkpoints.");
                             }
                         }
                     }
-                    Tree(Node::Nil) => Tree(Node::Nil),
+                    Node::Nil => Tree(Node::Nil),
+                    Node::Pruned => Tree(Node::Pruned),
                 }
             }
         }
@@ -905,7 +984,7 @@ mod tests {
 
     use super::{LocatedPrunableTree, PrunableTree, RetentionFlags};
     use crate::{
-        error::QueryError,
+        error::{InsertionError, QueryError},
         tree::{
             tests::{leaf, nil, parent},
             LocatedTree,
@@ -1068,27 +1147,22 @@ mod tests {
     }
 
     #[test]
-    fn located_insert_subtree_leaf_overwrites() {
+    fn located_insert_subtree_prevents_leaf_overwrite_conflict() {
         let t: LocatedPrunableTree<String> = LocatedTree {
             root_addr: Address::from_parts(2.into(), 1),
             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((
+            Err(InsertionError::Conflict(conflict_addr)),
-                LocatedTree {
-                    root_addr: Address::from_parts(2.into(), 1),
-                    root: parent(leaf(("b".to_string(), RetentionFlags::EPHEMERAL)), nil()),
-                },
-                vec![],
-            )),
         );
     }
 

shardtree/src/store/caching.rs

@@ -219,7 +219,7 @@ mod tests {
             append_str, check_operations, unmark, witness, CombinedTree, Operation, TestHashable,
             Tree,
         },
-        Hashable, Position, Retention,
+        Hashable, Marking, Position, Retention,
     };
 
     use super::CachingShardStore;
@@ -300,7 +300,7 @@ mod tests {
                 String::from_u64(0),
                 Checkpoint {
                     id: 1,
-                    is_marked: true,
+                    marking: Marking::Marked,
                 },
             ));
             for _ in 0..3 {
@@ -542,7 +542,7 @@ mod tests {
                 String::from_u64(0),
                 Checkpoint {
                     id: 1,
-                    is_marked: true,
+                    marking: Marking::Marked,
                 },
             ));
             assert!(tree.rewind());
@@ -584,7 +584,7 @@ mod tests {
                 String::from_u64(6),
                 Checkpoint {
                     id: 1,
-                    is_marked: true,
+                    marking: Marking::Marked,
                 },
             ));
             assert!(tree.append(String::from_u64(7), Ephemeral));
@@ -671,7 +671,7 @@ mod tests {
                     String::from_u64(3),
                     Checkpoint {
                         id: 1,
-                        is_marked: true,
+                        marking: Marking::Marked,
                     },
                 ),
                 Append(String::from_u64(4), Marked),
@@ -680,21 +680,21 @@ mod tests {
                     String::from_u64(5),
                     Checkpoint {
                         id: 3,
-                        is_marked: false,
+                        marking: Marking::None,
                     },
                 ),
                 Append(
                     String::from_u64(6),
                     Checkpoint {
                         id: 4,
-                        is_marked: false,
+                        marking: Marking::None,
                     },
                 ),
                 Append(
                     String::from_u64(7),
                     Checkpoint {
                         id: 5,
-                        is_marked: false,
+                        marking: Marking::None,
                     },
                 ),
                 Witness(3u64.into(), 5),
@@ -732,7 +732,7 @@ mod tests {
                     String::from_u64(0),
                     Checkpoint {
                         id: 1,
-                        is_marked: true,
+                        marking: Marking::Marked,
                     },
                 ),
                 Append(String::from_u64(0), Ephemeral),
@@ -742,7 +742,7 @@ mod tests {
                     String::from_u64(0),
                     Checkpoint {
                         id: 2,
-                        is_marked: false,
+                        marking: Marking::None,
                     },
                 ),
                 Append(String::from_u64(0), Ephemeral),
@@ -790,7 +790,7 @@ mod tests {
                     String::from_u64(0),
                     Checkpoint {
                         id: 4,
-                        is_marked: false,
+                        marking: Marking::None,
                     },
                 ),
                 Rewind,
@@ -818,14 +818,14 @@ mod tests {
                     String::from_u64(0),
                     Checkpoint {
                         id: 1,
-                        is_marked: true,
+                        marking: Marking::Marked,
                     },
                 ),
                 Append(
                     String::from_u64(0),
                     Checkpoint {
                         id: 4,
-                        is_marked: false,
+                        marking: Marking::None,
                     },
                 ),
                 Witness(Position::from(2), 2),
@@ -962,7 +962,7 @@ mod tests {
                     "a",
                     Retention::Checkpoint {
                         id: 1,
-                        is_marked: true,
+                        marking: Marking::Marked,
                     },
                 ),
                 witness(1, 1),

shardtree/src/testing.rs

@@ -86,7 +86,14 @@ where
                     leaf,
                     match (is_checkpoint, is_marked) {
                         (false, false) => Retention::Ephemeral,
-                        (true, is_marked) => Retention::Checkpoint { id, is_marked },
+                        (true, is_marked) => Retention::Checkpoint {
+                            id,
+                            marking: if is_marked {
+                                Marking::Marked
+                            } else {
+                                Marking::None
+                            },
+                        },
                         (false, true) => Retention::Marked,
                     },
                 )
@@ -121,10 +128,10 @@ where
                 .map(|(id, (leaf, retention))| {
                     let pos = Position::try_from(id).unwrap();
                     match retention {
-                        Retention::Ephemeral => (),
+                        Retention::Ephemeral | Retention::Reference => (),
-                        Retention::Checkpoint { is_marked, .. } => {
+                        Retention::Checkpoint { marking, .. } => {
                             checkpoint_positions.push(pos);
-                            if is_marked {
+                            if marking == Marking::Marked {
                                 marked_positions.push(pos);
                             }
                         }
@@ -234,7 +241,7 @@ pub fn check_shardtree_insertion<
         tree.batch_insert(
             Position::from(1),
             vec![
-                ("b".to_string(), Retention::Checkpoint { id: 1, is_marked: false }),
+                ("b".to_string(), Retention::Checkpoint { id: 1, marking: Marking::None }),
                 ("c".to_string(), Retention::Ephemeral),
                 ("d".to_string(), Retention::Marked),
             ].into_iter()
@@ -285,7 +292,7 @@ pub fn check_shardtree_insertion<
             Position::from(10),
             vec![
                 ("k".to_string(), Retention::Ephemeral),
-                ("l".to_string(), Retention::Checkpoint { id: 2, is_marked: false }),
+                ("l".to_string(), Retention::Checkpoint { id: 2, marking: Marking::None }),
                 ("m".to_string(), Retention::Ephemeral),
             ].into_iter()
         ),
@@ -386,7 +393,7 @@ pub fn check_witness_with_pruned_subtrees<
                     'c' => Retention::Marked,
                     'h' => Retention::Checkpoint {
                         id: 3,
-                        is_marked: false,
+                        marking: Marking::None,
                     },
                     _ => Retention::Ephemeral,
                 },

shardtree/src/tree.rs

@@ -19,6 +19,9 @@ pub enum Node<C, A, V> {
     Leaf { value: V },
     /// The empty tree; a subtree or leaf for which no information is available.
     Nil,
+    /// An empty node in the tree created as a consequence of partial reinserion of data into a
+    /// subtree after the subtree was previously pruned.
+    Pruned,
 }
 
 impl<C, A, V> Node<C, A, V> {
@@ -35,7 +38,8 @@ impl<C, A, V> Node<C, A, V> {
         match self {
             Node::Parent { .. } => None,
             Node::Leaf { value } => Some(value),
-            Node::Nil { .. } => None,
+            Node::Nil => None,
+            Node::Pruned => None,
         }
     }
 
@@ -45,6 +49,7 @@ impl<C, A, V> Node<C, A, V> {
             Node::Parent { ann, .. } => Some(ann),
             Node::Leaf { .. } => None,
             Node::Nil => None,
+            Node::Pruned => None,
         }
     }
 
@@ -71,6 +76,7 @@ impl<'a, C: Clone, A: Clone, V: Clone> Node<C, &'a A, &'a V> {
                 value: (*value).clone(),
             },
             Node::Nil => Node::Nil,
+            Node::Pruned => Node::Pruned,
         }
     }
 }
@@ -92,6 +98,11 @@ impl<A, V> Tree<A, V> {
         Tree(Node::Nil)
     }
 
+    /// Constructs the empty tree consisting of a single pruned node.
+    pub fn empty_pruned() -> Self {
+        Tree(Node::Pruned)
+    }
+
     /// Constructs a tree containing a single leaf.
     pub fn leaf(value: V) -> Self {
         Tree(Node::Leaf { value })
@@ -117,18 +128,13 @@ impl<A, V> Tree<A, V> {
         Tree(self.0.reannotate(ann))
     }
 
-    /// Returns `true` if no [`Node::Nil`] nodes are present in the tree, `false` otherwise.
+    /// Returns `true` this is a [`Node::Leaf`], `false` otherwise.
-    pub fn is_complete(&self) -> bool {
+    pub fn is_leaf(&self) -> bool {
-        match &self.0 {
+        matches!(&self.0, Node::Leaf { .. })
-            Node::Parent { left, right, .. } => {
-                left.as_ref().is_complete() && right.as_ref().is_complete()
-            }
-            Node::Leaf { .. } => true,
-            Node::Nil { .. } => false,
-        }
     }
 
-    /// Returns a vector of the addresses of [`Node::Nil`] subtree roots within this tree.
+    /// Returns a vector of the addresses of [`Node::Nil`] and [`Node::Pruned`] subtree roots
+    /// within this tree.
     ///
     /// The given address must correspond to the root of this tree, or this method will
     /// yield incorrect results or may panic.
@@ -149,9 +155,46 @@ impl<A, V> Tree<A, V> {
                 left_incomplete
             }
             Node::Leaf { .. } => vec![],
-            Node::Nil => vec![root_addr],
+            Node::Nil | Node::Pruned => 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,
+    {
+        Tree(match &self.0 {
+            Node::Parent { ann, left, right } => Node::Parent {
+                ann: ann.clone(),
+                left: Arc::new(left.map(f)),
+                right: Arc::new(right.map(f)),
+            },
+            Node::Leaf { value } => Node::Leaf { value: f(value) },
+            Node::Nil => Node::Nil,
+            Node::Pruned => Node::Pruned,
+        })
+    }
+
+    /// 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,
+            Node::Pruned => Node::Pruned,
+        }))
+    }
 }
 
 /// A binary Merkle tree with its root at the given address.
@@ -202,18 +245,19 @@ impl<A, V> LocatedTree<A, V> {
     /// 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> {
+        Self::max_position_internal(self.root_addr, &self.root)
-            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)
+    pub(crate) fn max_position_internal(addr: Address, root: &Tree<A, V>) -> Option<Position> {
+        match &root.0 {
+            Node::Nil => None,
+            Node::Leaf { .. } | Node::Pruned => Some(addr.position_range_end() - 1),
+            Node::Parent { left, right, .. } => {
+                let (l_addr, r_addr) = addr.children().unwrap();
+                Self::max_position_internal(r_addr, right.as_ref())
+                    .or_else(|| Self::max_position_internal(l_addr, left.as_ref()))
+            }
+        }
     }
 
     /// Returns the value at the specified position, if any.
@@ -239,6 +283,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> {