Tests and fixes for a number of prune-flag bugs

benches/prune.rs

@@ -6,87 +6,56 @@ fn main() {
     Divan::from_args().main();
 }
 
-fn fixture(path: &[u8]) -> PathMap<u64> {
-    let mut map = PathMap::new();
-    map.set_val_at(path, 1);
-    map
-}
-
-fn run_remove_val(bencher: Bencher, path: &[u8], root_len: usize, prune: bool) {
-    bencher.with_inputs(|| fixture(path)).bench_local_values(|mut map| {
-        let mut wz = map.write_zipper_at_path(&path[..root_len]);
-        wz.descend_to(&path[root_len..]);
-        black_box(wz.remove_val(prune));
-    });
-}
-
-fn run_remove_branches(bencher: Bencher, path: &[u8], root_len: usize, prune: bool) {
-    bencher.with_inputs(|| fixture(path)).bench_local_values(|mut map| {
-        let focus = &path[..path.len() - 1];
-        let mut wz = map.write_zipper_at_path(&path[..root_len]);
-        wz.descend_to(&focus[root_len..]);
-        black_box(wz.remove_branches(prune));
-    });
-}
-
-fn run_prune_path(bencher: Bencher, path: &[u8], root_len: usize) {
+#[divan::bench]
+fn prune_path(bencher: Bencher) {
     bencher.with_inputs(|| {
         let mut map = PathMap::<u64>::new();
-        map.create_path(path);
+        map.create_path(b"abcd");
         map
     }).bench_local_values(|mut map| {
-        let mut wz = map.write_zipper_at_path(&path[..root_len]);
-        wz.descend_to(&path[root_len..]);
+        let mut wz = map.write_zipper_at_path(b"ab");
+        wz.descend_to(b"cd");
         black_box(wz.prune_path());
     });
 }
 
-#[divan::bench]
-fn prune_path_short_root_at_map_root(bencher: Bencher) {
-    run_prune_path(bencher, b"abcd", 0);
-}
-
-#[divan::bench]
-fn prune_path_short_root_inside_node(bencher: Bencher) {
-    run_prune_path(bencher, b"abcd", 2);
-}
-
-#[divan::bench]
-fn prune_path_long_root_inside_node(bencher: Bencher) {
-    let path: Vec<u8> = (0..100).collect();
-    run_prune_path(bencher, &path, 95);
-}
-
-#[divan::bench(args = [false, true])]
-fn remove_val_short(bencher: Bencher, prune: bool) {
-    run_remove_val(bencher, b"abcd", 2, prune);
-}
-
 #[divan::bench(args = [false, true])]
-fn remove_val_long_root_above_node(bencher: Bencher, prune: bool) {
-    let path: Vec<u8> = (0..100).collect();
-    run_remove_val(bencher, &path, 5, prune);
-}
-
-#[divan::bench(args = [false, true])]
-fn remove_val_long_root_inside_node(bencher: Bencher, prune: bool) {
-    let path: Vec<u8> = (0..100).collect();
-    run_remove_val(bencher, &path, 95, prune);
-}
-
-#[divan::bench(args = [false, true])]
-fn remove_branches_short(bencher: Bencher, prune: bool) {
-    run_remove_branches(bencher, b"abcd", 2, prune);
+fn remove_val(bencher: Bencher, prune: bool) {
+    bencher.with_inputs(|| {
+        let mut map = PathMap::<u64>::new();
+        map.set_val_at(b"abcd", 1);
+        map
+    }).bench_local_values(|mut map| {
+        let mut wz = map.write_zipper_at_path(b"ab");
+        wz.descend_to(b"cd");
+        black_box(wz.remove_val(prune));
+    });
 }
 
 #[divan::bench(args = [false, true])]
-fn remove_branches_long_root_above_node(bencher: Bencher, prune: bool) {
-    let path: Vec<u8> = (0..100).collect();
-    run_remove_branches(bencher, &path, 5, prune);
+fn remove_branches(bencher: Bencher, prune: bool) {
+    bencher.with_inputs(|| {
+        let mut map = PathMap::<u64>::new();
+        map.set_val_at(b"abcd", 1);
+        map
+    }).bench_local_values(|mut map| {
+        let mut wz = map.write_zipper_at_path(b"ab");
+        wz.descend_to(b"c");
+        black_box(wz.remove_branches(prune));
+    });
 }
 
 #[divan::bench(args = [false, true])]
-fn remove_branches_long_root_inside_node(bencher: Bencher, prune: bool) {
-    let path: Vec<u8> = (0..100).collect();
-    run_remove_branches(bencher, &path, 95, prune);
+fn take_map(bencher: Bencher, prune: bool) {
+    bencher.with_inputs(|| {
+        let mut map = PathMap::<u64>::new();
+        map.set_val_at(b"a", 1);
+        map.set_val_at(b"ab", 2);
+        map.set_val_at(b"ac", 3);
+        map
+    }).bench_local_values(|mut map| {
+        let mut wz = map.write_zipper();
+        wz.descend_to(b"a");
+        black_box(wz.take_map(prune));
+    });
 }

pathmap-book/src/1.02.06_zipper_writing.md

@@ -29,7 +29,7 @@ The [`ZipperWriting`] trait provides methods to modify the trie structure and va
 - [`remove_unmasked_branches`] selectively removes branches based on a byte mask
 
 ## Pruning Behavior
-Many operations accept a `prune` parameter that controls whether dangling paths should be automatically cleaned up. When `prune` is `true`, the operation will remove any path segments that become empty, aka "dangling" (no values, and no further downstream branches) as a result of the modification.
+Many operations accept a `prune` parameter. With `prune=false`, an emptied focus remains dangling. With `prune=true`, the result is the same as calling the operation with `false` and then [`prune_path`] at the focus; this also handles a pre-existing dangling tip. Pruning from the middle of a path with descendants does nothing. For [`join_into_take`], the flag applies to the source zipper's focus.
 
 Automatic pruning helps maintain a compact trie structure by removing unnecessary nodes, but can be counter-productive when you plan to perform additional operations that might reuse those paths.
 

src/dense_byte_node.rs

@@ -185,7 +185,7 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> ByteNode<Cf, A>
             let cf = unsafe { self.values.get_unchecked_mut(ix) };
             let result = cf.take_val();
 
-            if prune_limit == 0 && !cf.has_rec() {
+            if prune_limit == 0 && cf.rec().map(|rec| rec.as_tagged().node_is_empty()).unwrap_or(true) {
                 self.mask.clear_bit(k);
                 self.values.remove(ix);
             }
@@ -1003,10 +1003,11 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
         if self.mask.test_bit(k) {
             let ix = self.mask.index_of(k) as usize;
             let cf = unsafe { self.values.get_unchecked_mut(ix) };
+            let had_branches = cf.rec().is_some_and(|node| !node.as_tagged().node_is_empty());
             match (cf.has_rec(), cf.has_val()) {
                 (true, true) => {
                     cf.set_rec_option(None);
-                    true
+                    had_branches
                 },
                 (true, false) => {
                     if prune_limit == 0 {
@@ -1015,11 +1016,14 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
                     } else {
                         cf.set_rec_option(None);
                     }
-                    true
+                    had_branches
                 },
-                (false, _) => {
+                (false, false) if prune_limit == 0 => {
+                    self.values.remove(ix);
+                    self.mask.clear_bit(k);
                     false
                 },
+                (false, _) => false,
             }
         } else {
             false
@@ -1181,9 +1185,13 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
         (Some(&ALL_BYTES[prefix..=prefix]), cf.rec().map(|cf| cf.as_tagged()))
     }
 
-    fn node_remove_unmasked_branches(&mut self, key: &[u8], mask: ByteMask, _prune_limit: usize) {
+    fn node_remove_unmasked_branches(&mut self, key: &[u8], mask: ByteMask, prune_limit: usize) {
         if key.len() > 0 {
-            //We're in a non-existent path below this node
+            // A one-byte key may name an existing dangling child.  With pruning
+            // enabled, remove it even though it has no branches to mask.
+            if key.len() == 1 && prune_limit == 0 {
+                self.node_remove_dangling(key, prune_limit);
+            }
             return
         }
         // in the future we can use `drain_filter`, but that's experimental