shardtree: Reduce dependence upon LocatedTree::max_position

shardtree/src/lib.rs

@@ -235,17 +235,15 @@ impl<
 
         let (append_result, position, checkpoint_id) =
             if let Some(subtree) = self.store.last_shard().map_err(ShardTreeError::Storage)? {
-                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() {
-                            LocatedTree::empty(addr.next_at_level()).append(value, retention)?
-                        } else {
-                            return Err(InsertionError::TreeFull.into());
-                        }
+                if subtree.root().is_full() {
+                    let addr = subtree.root_addr;
+                    if subtree.root_addr.index() < Self::max_subtree_index() {
+                        LocatedTree::empty(addr.next_at_level()).append(value, retention)?
+                    } else {
+                        return Err(InsertionError::TreeFull.into());
                     }
-                    _ => subtree.append(value, retention)?,
+                } else {
+                    subtree.append(value, retention)?
                 }
             } else {
                 let root_addr = Address::from_parts(Self::subtree_level(), 0);
@@ -573,13 +571,7 @@ impl<
                     .map_err(ShardTreeError::Storage)?;
 
                 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)?;

shardtree/src/prunable.rs

@@ -362,10 +362,11 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
     /// If the tree contains any [`Node::Nil`] nodes that are to the left of filled nodes in the
     /// tree, this will return an error containing the addresses of those nodes.
     pub fn right_filled_root(&self) -> Result<H, Vec<Address>> {
-        self.root_hash(
-            self.max_position()
-                .map_or_else(|| self.root_addr.position_range_start(), |pos| pos + 1),
-        )
+        let truncate_at = self
+            .max_position()
+            .map_or_else(|| self.root_addr.position_range_start(), |pos| pos + 1);
+
+        self.root_hash(truncate_at)
     }
 
     /// Returns the positions of marked leaves in the tree.
@@ -949,6 +950,33 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
             root: go(&to_clear, self.root_addr, &self.root),
         }
     }
+
+    #[cfg(test)]
+    pub(crate) fn flag_positions(&self) -> BTreeMap<Position, RetentionFlags> {
+        fn go<H>(
+            root: &PrunableTree<H>,
+            root_addr: Address,
+            acc: &mut BTreeMap<Position, RetentionFlags>,
+        ) {
+            match &root.0 {
+                Node::Parent { left, right, .. } => {
+                    let (l_addr, r_addr) = root_addr
+                        .children()
+                        .expect("A parent node cannot appear at level 0");
+                    go(&left, l_addr, acc);
+                    go(&right, r_addr, acc);
+                }
+                Node::Leaf { value } if value.1 != RetentionFlags::EPHEMERAL => {
+                    acc.insert(root_addr.max_position(), value.1);
+                }
+                _ => (),
+            }
+        }
+
+        let mut result = BTreeMap::new();
+        go(&self.root, self.root_addr, &mut result);
+        result
+    }
 }
 
 // We need an applicative functor for Result for this function so that we can correctly
@@ -971,13 +999,15 @@ fn accumulate_result_with<A, B, C>(
 
 #[cfg(test)]
 mod tests {
-    use std::collections::BTreeSet;
+    use std::collections::{BTreeMap, BTreeSet};
 
     use incrementalmerkletree::{Address, Level, Position};
+    use proptest::proptest;
 
     use super::{LocatedPrunableTree, PrunableTree, RetentionFlags};
     use crate::{
         error::{InsertionError, QueryError},
+        testing::{arb_char_str, arb_prunable_tree},
         tree::{
             tests::{leaf, nil, parent},
             LocatedTree,
@@ -1197,4 +1227,34 @@ mod tests {
             )]))
         );
     }
+
+    proptest! {
+        #[test]
+        fn clear_flags(
+            root in arb_prunable_tree(arb_char_str(), 8, 2^6)
+        ) {
+            let root_addr = Address::from_parts(Level::from(7), 0);
+            let tree = LocatedTree::from_parts(root_addr, root);
+
+            let (to_clear, to_retain) = tree.flag_positions().into_iter().enumerate().fold(
+                (BTreeMap::new(), BTreeMap::new()),
+                |(mut to_clear, mut to_retain), (i, (pos, flags))| {
+                    if i % 2 == 0 {
+                        to_clear.insert(pos, flags);
+                    } else {
+                        to_retain.insert(pos, flags);
+                    }
+                    (to_clear, to_retain)
+                }
+            );
+
+            let pre_clearing_max_position = tree.max_position();
+            let cleared = tree.clear_flags(to_clear);
+
+            // Clearing flags should not modify the max position of leaves represented
+            // in the shard.
+            assert!(cleared.max_position() == pre_clearing_max_position);
+            assert_eq!(to_retain, cleared.flag_positions());
+        }
+    }
 }

shardtree/src/tree.rs

@@ -148,6 +148,15 @@ impl<A, V> Tree<A, V> {
         matches!(&self.0, Node::Leaf { .. })
     }
 
+    /// Returns `true` if no additional nodes can be appended to this tree.
+    pub fn is_full(&self) -> bool {
+        match &self.0 {
+            Node::Nil => false,
+            Node::Leaf { .. } | Node::Pruned => true,
+            Node::Parent { right, .. } => right.is_full(),
+        }
+    }
+
     /// Returns a vector of the addresses of [`Node::Nil`] and [`Node::Pruned`] subtree roots
     /// within this tree.
     ///
@@ -260,19 +269,18 @@ 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> {
-        Self::max_position_internal(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()))
+        fn go<A, V>(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();
+                    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.