Merge pull request #87 from zcash/proptest-shardtree-insert_tree

`shardtree`: Fix bugs in `ShardTree::insert_tree`

shardtree/Cargo.toml

@@ -36,5 +36,6 @@ test-dependencies = ["proptest", "assert_matches"]
 
 [target.'cfg(unix)'.dev-dependencies]
 pprof = { version = "0.9", features = ["criterion", "flamegraph"] } # MSRV 1.56
-dashmap = ">=5, <5.5.0"
+dashmap = ">=5, <5.5.0" # 5.5 has MSRV > 1.60
 inferno = ">=0.11, <0.11.5" # MSRV 1.59
+tempfile = ">=3, <3.7.0" # 3.7 has MSRV 1.63

shardtree/src/batch.rs

@@ -382,10 +382,16 @@ fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
 
 #[cfg(test)]
 mod tests {
+    use std::iter;
+
     use incrementalmerkletree::{Address, Level, Position, Retention};
 
     use super::{LocatedPrunableTree, RetentionFlags};
-    use crate::tree::tests::{leaf, nil, parent};
+    use crate::{
+        memory::MemoryShardStore,
+        tree::tests::{leaf, nil, parent},
+        BatchInsertionResult, ShardTree,
+    };
 
     #[test]
     fn located_from_iter_non_sibling_adjacent() {
@@ -465,4 +471,83 @@ mod tests {
             .unwrap();
         assert_eq!(complete.subtree.right_filled_root(), Ok("abcd".to_string()));
     }
+
+    #[allow(clippy::type_complexity)]
+    pub(super) fn build_insert_tree_and_batch_insert(
+        leaves: Vec<(String, Retention<usize>)>,
+    ) -> (
+        ShardTree<MemoryShardStore<String, usize>, 6, 3>,
+        ShardTree<MemoryShardStore<String, usize>, 6, 3>,
+    ) {
+        let max_checkpoints = 10;
+        let start = Position::from(0);
+        let end = start + leaves.len() as u64;
+
+        // Construct a tree using `ShardTree::insert_tree`.
+        let mut left = ShardTree::new(MemoryShardStore::empty(), max_checkpoints);
+        if let Some(BatchInsertionResult {
+            subtree,
+            mut remainder,
+            ..
+        }) = LocatedPrunableTree::from_iter(start..end, 0.into(), leaves.clone().into_iter())
+        {
+            assert_eq!(remainder.next(), None);
+            left.insert_tree(subtree).unwrap();
+        }
+
+        // Construct a tree using `ShardTree::batch_insert`.
+        let mut right = ShardTree::new(MemoryShardStore::empty(), max_checkpoints);
+        right.batch_insert(start, leaves.into_iter()).unwrap();
+
+        (left, right)
+    }
+
+    #[test]
+    fn batch_insert_matches_insert_tree() {
+        {
+            let (lhs, rhs) = build_insert_tree_and_batch_insert(vec![]);
+            assert_eq!(lhs.max_leaf_position(0), Ok(None));
+            assert_eq!(rhs.max_leaf_position(0), Ok(None));
+        }
+
+        for i in 0..64 {
+            let num_leaves = i + 1;
+            let leaves = iter::repeat(("a".into(), Retention::Ephemeral))
+                .take(num_leaves)
+                .collect();
+            let expected_root = (0..64)
+                .map(|c| if c < num_leaves { 'a' } else { '_' })
+                .fold(String::with_capacity(64), |mut acc, c| {
+                    acc.push(c);
+                    acc
+                });
+
+            let (lhs, rhs) = build_insert_tree_and_batch_insert(leaves);
+            assert_eq!(lhs.max_leaf_position(0), Ok(Some(Position::from(i as u64))));
+            assert_eq!(rhs.max_leaf_position(0), Ok(Some(Position::from(i as u64))));
+
+            assert_eq!(lhs.root_at_checkpoint(0).unwrap(), expected_root);
+            assert_eq!(rhs.root_at_checkpoint(0).unwrap(), expected_root);
+        }
+    }
+}
+
+#[cfg(test)]
+mod proptests {
+    use proptest::prelude::*;
+
+    use super::tests::build_insert_tree_and_batch_insert;
+    use crate::testing::{arb_char_str, arb_leaves};
+
+    proptest! {
+        #[test]
+        fn batch_insert_matches_insert_tree(
+            leaves in arb_leaves(arb_char_str())
+        ) {
+            let (left, right) = build_insert_tree_and_batch_insert(leaves);
+
+            // Check that the resulting trees are equal.
+            assert_eq!(left.root_at_checkpoint(0), right.root_at_checkpoint(0));
+        }
+    }
 }

shardtree/src/lib.rs

@@ -587,7 +587,20 @@ impl<
     ) -> Result<Vec<IncompleteAt>, ShardTreeError<S::Error>> {
         let mut all_incomplete = vec![];
         for subtree in tree.decompose_to_level(Self::subtree_level()).into_iter() {
-            let root_addr = subtree.root_addr;
+            // `ShardTree::max_leaf_position` relies on the invariant that the last shard
+            // in the subtrees vector is never created without a leaf then being added to
+            // it. `LocatedTree::decompose_to_level` can return a trailing empty subtree
+            // 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() {
+                continue;
+            }
+
+            // `LocatedTree::decompose_to_level` will return the tree as-is if it is
+            // smaller than a shard, so we can't assume that the address of `subtree` is a
+            // valid shard address.
+            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
                 .store

shardtree/src/testing.rs

@@ -66,6 +66,35 @@ where
     )
 }
 
+/// Constructs a random sequence of leaves that form a tree of size up to 2^6.
+pub fn arb_leaves<H: Strategy + Clone>(
+    arb_leaf: H,
+) -> impl Strategy<Value = Vec<(H::Value, Retention<usize>)>>
+where
+    H::Value: Hashable + Clone + PartialEq,
+{
+    vec(
+        (arb_leaf, weighted(0.1), weighted(0.2)),
+        0..=(2usize.pow(6)),
+    )
+    .prop_map(|leaves| {
+        leaves
+            .into_iter()
+            .enumerate()
+            .map(|(id, (leaf, is_marked, is_checkpoint))| {
+                (
+                    leaf,
+                    match (is_checkpoint, is_marked) {
+                        (false, false) => Retention::Ephemeral,
+                        (true, is_marked) => Retention::Checkpoint { id, is_marked },
+                        (false, true) => Retention::Marked,
+                    },
+                )
+            })
+            .collect()
+    })
+}
+
 /// Constructs a random shardtree of size up to 2^6 with shards of size 2^3. Returns the tree,
 /// along with vectors of the checkpoint and mark positions.
 pub fn arb_shardtree<H: Strategy + Clone>(
@@ -80,11 +109,7 @@ pub fn arb_shardtree<H: Strategy + Clone>(
 where
     H::Value: Hashable + Clone + PartialEq,
 {
-    vec(
-        (arb_leaf, weighted(0.1), weighted(0.2)),
-        0..=(2usize.pow(6)),
-    )
-    .prop_map(|leaves| {
+    arb_leaves(arb_leaf).prop_map(|leaves| {
         let mut tree = ShardTree::new(MemoryShardStore::empty(), 10);
         let mut checkpoint_positions = vec![];
         let mut marked_positions = vec![];
@@ -93,25 +118,19 @@ where
             leaves
                 .into_iter()
                 .enumerate()
-                .map(|(id, (leaf, is_marked, is_checkpoint))| {
-                    (
-                        leaf,
-                        match (is_checkpoint, is_marked) {
-                            (false, false) => Retention::Ephemeral,
-                            (true, is_marked) => {
-                                let pos = Position::try_from(id).unwrap();
-                                checkpoint_positions.push(pos);
-                                if is_marked {
-                                    marked_positions.push(pos);
-                                }
-                                Retention::Checkpoint { id, is_marked }
+                .map(|(id, (leaf, retention))| {
+                    let pos = Position::try_from(id).unwrap();
+                    match retention {
+                        Retention::Ephemeral => (),
+                        Retention::Checkpoint { is_marked, .. } => {
+                            checkpoint_positions.push(pos);
+                            if is_marked {
+                                marked_positions.push(pos);
                             }
-                            (false, true) => {
-                                marked_positions.push(Position::try_from(id).unwrap());
-                                Retention::Marked
-                            }
-                        },
-                    )
+                        }
+                        Retention::Marked => marked_positions.push(pos),
+                    }
+                    (leaf, retention)
                 }),
         )
         .unwrap();