This changes how we generate the trees by using the recursive strateg…

…y once and then mapping file-nodes to empty directories. This should be more performant than our old implementation, as that lead to the inner strategy (which defines # of nodes, etc ... for the whole tree) to be run multiple times, causing a degradation in performance of proptest generation. Signed-off-by: Christian Hagemeier <[email protected]>
awslabs · Jan 14, 2025 · 1410e86 · 1410e86
1 parent ab77aaa
commit 1410e86
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/mountpoint-s3/tests/reftests/generators.rs b/mountpoint-s3/tests/reftests/generators.rs
@@ -127,21 +127,24 @@ impl Debug for TreeNode {
 /// Leaves are always [TreeNode::File] or an empty [TreeNode::Directory].
 /// Parents are always [TreeNode::Directory].
 pub fn gen_tree(depth: u32, max_size: u32, max_items: u32, max_width: usize) -> impl Strategy<Value = TreeNode> {
-    let leaf = any::<FileContent>().prop_map(TreeNode::File);
-
-    let strategy = leaf.prop_recursive(
+    let leaf = prop_oneof![any::<FileContent>().prop_map(TreeNode::File),];
+    leaf.prop_recursive(
         depth,     // levels
         max_size,  // max number of nodes
         max_items, // number of items per collection
         move |inner| {
-            // Take the inner strategy and make the recursive cases.
-            // Since the size of the tree could be zero, this also introduces directories as leaves.
-            prop::collection::btree_map(any::<Name>(), inner, 0..max_width).prop_map(TreeNode::Directory)
+            prop_oneof![
+                // Take the inner strategy and make the recursive cases.
+                // Since the size of the tree could be zero, this also introduces directories as leaves.
+                prop::collection::btree_map(any::<Name>(), inner, 0..max_width).prop_map(TreeNode::Directory),
+            ]
         },
-    );
-
-    // Ensure the root is always a directory by wrapping the inner strategy
-    prop::collection::btree_map(any::<Name>(), strategy, 0..max_width).prop_map(TreeNode::Directory)
+    )
+    // Ensure the root is always a directory by transforming invalid nodes (i.e. file as root) into empty directories
+    .prop_map(|x| match x {
+        TreeNode::File(_) => TreeNode::Directory(BTreeMap::from([])),
+        _ => x,
+    })
 }
 
 /// Take a generated tree and create the corresponding S3 namespace (list of keys)