Fixes for a bunch of prune issues

.github/scripts/bench_ab.py

@@ -165,6 +165,11 @@ def compare_bench(self, bench, rounds_so_far):
         cmp = self.fmt.compare_fields(avg['base'], avg['head'], 'median_ns')
         self.results[bench] = {f'{g}/{c}': {'base': r['base'], 'head': r['other'], 'pct': r['pct']} for (g, c), r in cmp.items()}
         table = re.sub(r'\x1b\[[0-9;]*m', '', self.fmt.render_divan_table(cmp))
+        # a PR that adds, removes or renames cases leaves them on one side only; say so rather than drop them silently
+        for side, mine, theirs in (('base', avg['base'], avg['head']), ('head', avg['head'], avg['base'])):
+            only = sorted(f'{g}/{c}' for g, c in mine.keys() - theirs.keys())
+            if only:
+                table += f'\nonly in {side} ({len(only)}): ' + ', '.join(only)
         text = (f'{bench}  (base {self.short(self.base_sha)}  head {self.short(self.head_sha)}'
                 f'  rounds {rounds_so_far}  median ns)\n{table}\n\n')
         (self.out / f'cmp-{bench}.txt').write_text(text)

Cargo.toml

@@ -137,6 +137,10 @@ harness = false
 name = "zipper_head_owned"
 harness = false
 
+[[bench]]
+name = "prune"
+harness = false
+
 [[bench]]
 name = "act_paths"
 harness = false

benches/divan_fmt.py

@@ -119,6 +119,8 @@ def render_divan_table(data):
     for group in grouped:
         grouped[group].sort(key=lambda item: case_sort_key(item[0]))
 
+    if not grouped:
+        return "(no cases)"
     (_, first_record) = next(iter(grouped.values()))[0]
     lines = []
 

benches/prune.rs

@@ -0,0 +1,92 @@
+use divan::{Bencher, Divan, black_box};
+use pathmap::PathMap;
+use pathmap::zipper::*;
+
+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) {
+    bencher.with_inputs(|| {
+        let mut map = PathMap::<u64>::new();
+        map.create_path(path);
+        map
+    }).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.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);
+}
+
+#[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);
+}
+
+#[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);
+}

benches/zipper_head_owned.rs

@@ -2,6 +2,8 @@ use divan::{Bencher, Divan, black_box};
 use pathmap::PathMap;
 use pathmap::zipper::*;
 
+const REPEATS: usize = 100;
+
 fn main() {
     Divan::from_args().sample_count(100).main();
 }
@@ -16,82 +18,84 @@ fn zipper_head_fixture() -> PathMap<usize> {
     map
 }
 
-fn bench_read_creation<F>(bencher: Bencher, repeats: usize, mut read_child_count: F)
+fn bench_head_read_creation<'trie, H>(bencher: Bencher, head: &H)
 where
-    F: FnMut() -> usize,
+    H: ZipperCreation<'trie, usize>,
 {
+    let path = [7u8];
     bencher.bench_local(|| {
         let mut observed = 0usize;
-        for _ in 0..repeats {
-            observed += read_child_count();
+        for _ in 0..REPEATS {
+            let reader = head.read_zipper_at_borrowed_path(black_box(&path)).unwrap();
+            observed += reader.child_count();
         }
         black_box(observed);
     });
 }
 
-fn bench_write_creation_cleanup<F>(bencher: Bencher, repeats: usize, mut create_and_cleanup: F)
+fn bench_head_write_creation_cleanup<'trie, H, const CHECKED: bool>(bencher: Bencher, head: &H)
 where
-    F: FnMut([u8; 2]) -> usize,
+    H: ZipperCreation<'trie, usize>,
 {
     bencher.bench_local(|| {
+        let mut writers = Vec::with_capacity(REPEATS);
         let mut observed = 0usize;
-        for i in 0..repeats {
-            observed += create_and_cleanup([240u8, i as u8]);
+        for i in 0..REPEATS {
+            let path = black_box([240u8, i as u8]);
+            // The paths are disjoint. All writers stay live until cleanup below.
+            let writer = if CHECKED {
+                head.write_zipper_at_exclusive_path(path).unwrap()
+            } else {
+                unsafe { head.write_zipper_at_exclusive_path_unchecked(path) }
+            };
+            observed += writer.path_exists() as usize;
+            writers.push(writer);
+        }
+        for writer in writers {
+            head.cleanup_write_zipper(writer);
         }
         black_box(observed);
     });
 }
 
-fn bench_head_read_creation<'trie, H>(bencher: Bencher, repeats: usize, head: &H)
-where
-    H: ZipperCreation<'trie, usize>,
-{
-    let path = [7u8];
-
-    bench_read_creation(bencher, repeats, || {
-        let reader = head.read_zipper_at_borrowed_path(black_box(&path)).unwrap();
-        reader.child_count()
-    });
+#[divan::bench]
+fn borrowed_head_read_creation(bencher: Bencher) {
+    let mut map = zipper_head_fixture();
+    let head = black_box(&mut map).zipper_head();
+    bench_head_read_creation(bencher, &head);
 }
 
-fn bench_head_write_creation_cleanup<'trie, H>(bencher: Bencher, repeats: usize, head: &H)
-where
-    H: ZipperCreation<'trie, usize>,
-{
-    bench_write_creation_cleanup(bencher, repeats, |path| {
-        let writer = head
-            .write_zipper_at_exclusive_path(black_box(path))
-            .unwrap();
-        let observed = writer.path_exists() as usize;
-        head.cleanup_write_zipper(writer);
-        observed
-    });
+#[divan::bench]
+fn owned_head_read_creation(bencher: Bencher) {
+    let map = zipper_head_fixture();
+    let head = black_box(map).into_zipper_head([]);
+    bench_head_read_creation(bencher, &head);
 }
 
-#[divan::bench(args = [1usize, 10, 100])]
-fn borrowed_head_read_creation(bencher: Bencher, repeats: usize) {
+#[divan::bench]
+fn borrowed_head_write_creation_cleanup(bencher: Bencher) {
     let mut map = zipper_head_fixture();
     let head = black_box(&mut map).zipper_head();
-    bench_head_read_creation(bencher, repeats, &head);
+    bench_head_write_creation_cleanup::<_, true>(bencher, &head);
 }
 
-#[divan::bench(args = [1usize, 10, 100])]
-fn owned_head_read_creation(bencher: Bencher, repeats: usize) {
+#[divan::bench]
+fn owned_head_write_creation_cleanup(bencher: Bencher) {
     let map = zipper_head_fixture();
     let head = black_box(map).into_zipper_head([]);
-    bench_head_read_creation(bencher, repeats, &head);
+    bench_head_write_creation_cleanup::<_, true>(bencher, &head);
 }
 
-#[divan::bench(args = [1usize, 10, 100])]
-fn borrowed_head_write_creation_cleanup(bencher: Bencher, repeats: usize) {
+#[divan::bench]
+fn borrowed_head_write_creation_cleanup_unchecked(bencher: Bencher) {
     let mut map = zipper_head_fixture();
     let head = black_box(&mut map).zipper_head();
-    bench_head_write_creation_cleanup(bencher, repeats, &head);
+    bench_head_write_creation_cleanup::<_, false>(bencher, &head);
 }
 
-#[divan::bench(args = [1usize, 10, 100])]
-fn owned_head_write_creation_cleanup(bencher: Bencher, repeats: usize) {
+#[divan::bench]
+fn owned_head_write_creation_cleanup_unchecked(bencher: Bencher) {
     let map = zipper_head_fixture();
     let head = black_box(map).into_zipper_head([]);
-    bench_head_write_creation_cleanup(bencher, repeats, &head);
+    bench_head_write_creation_cleanup::<_, false>(bencher, &head);
 }

src/dense_byte_node.rs

@@ -178,14 +178,14 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> ByteNode<Cf, A>
     }
 
     #[inline]
-    pub fn remove_val(&mut self, k: u8, prune: bool) -> Option<V> {
+    pub fn remove_val(&mut self, k: u8, prune_limit: usize) -> Option<V> {
         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 result = cf.take_val();
 
-            if prune && !cf.has_rec() {
+            if prune_limit == 0 && !cf.has_rec() {
                 self.mask.clear_bit(k);
                 self.values.remove(ix);
             }
@@ -895,9 +895,9 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
             }
         }
     }
-    fn node_remove_val(&mut self, key: &[u8], prune: bool) -> Option<V> {
+    fn node_remove_val(&mut self, key: &[u8], prune_limit: usize) -> Option<V> {
         if key.len() == 1 {
-            self.remove_val(key[0], prune)
+            self.remove_val(key[0], prune_limit)
         } else {
             None
         }
@@ -929,8 +929,11 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
             }
         }
     }
-    fn node_remove_dangling(&mut self, key: &[u8]) -> usize {
+    fn node_remove_dangling(&mut self, key: &[u8], min_keep_len: usize) -> usize {
         debug_assert!(key.len() > 0);
+        if min_keep_len >= key.len() {
+            return 0;
+        }
         if key.len() == 1 {
             let k = key[0];
             if self.mask.test_bit(k) {
@@ -994,7 +997,7 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
             }
         }
     }
-    fn node_remove_all_branches(&mut self, key: &[u8], prune: bool) -> bool {
+    fn node_remove_all_branches(&mut self, key: &[u8], prune_limit: usize) -> bool {
         if key.len() > 1 {
             return false;
         }
@@ -1009,7 +1012,7 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
                     true
                 },
                 (true, false) => {
-                    if prune {
+                    if prune_limit == 0 {
                         self.values.remove(ix);
                         self.mask.clear_bit(k);
                     } else {
@@ -1181,7 +1184,7 @@ 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: bool) {
+    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
             return
@@ -1298,7 +1301,7 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
         }
     }
 
-    fn take_node_at_key(&mut self, key: &[u8], prune: bool) -> Option<TrieNodeODRc<V, A>> {
+    fn take_node_at_key(&mut self, key: &[u8], prune_limit: usize) -> Option<TrieNodeODRc<V, A>> {
         if key.len() < 2 {
             debug_assert!(key.len() == 1);
             let k = key[0];
@@ -1308,7 +1311,7 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
                 let cf = unsafe { self.values.get_unchecked_mut(ix) };
                 let result = cf.take_rec();
 
-                if prune && !cf.has_val() {
+                if prune_limit == 0 && !cf.has_val() {
                     self.mask.clear_bit(k);
                     self.values.remove(ix);
                 }

src/empty_node.rs

@@ -35,13 +35,13 @@ impl<V: Clone + Send + Sync, A: Allocator> TrieNode<V, A> for EmptyNode {
     fn node_get_val(&self, _key: &[u8]) -> Option<&V> {
         None
     }
-    fn node_remove_val(&mut self, _key: &[u8], _prune: bool) -> Option<V> {
+    fn node_remove_val(&mut self, _key: &[u8], _prune_limit: usize) -> Option<V> {
         unreachable!()
     }
     fn node_create_dangling(&mut self, _key: &[u8]) -> Result<(bool, bool), TrieNodeODRc<V, A>> {
         unreachable!()
     }
-    fn node_remove_dangling(&mut self, _key: &[u8]) -> usize {
+    fn node_remove_dangling(&mut self, _key: &[u8], _min_keep_len: usize) -> usize {
         unreachable!()
     }
     fn node_get_val_mut(&mut self, _key: &[u8]) -> Option<&mut V> {
@@ -53,10 +53,10 @@ impl<V: Clone + Send + Sync, A: Allocator> TrieNode<V, A> for EmptyNode {
     fn node_set_branch(&mut self, _key: &[u8], _new_node: TrieNodeODRc<V, A>) -> Result<bool, TrieNodeODRc<V, A>> {
         unreachable!() //we should head this off upstream
     }
-    fn node_remove_all_branches(&mut self, _key: &[u8], _prune: bool) -> bool {
+    fn node_remove_all_branches(&mut self, _key: &[u8], _prune_limit: usize) -> bool {
         false
     }
-    fn node_remove_unmasked_branches(&mut self, _key: &[u8], _mask: ByteMask, _prune: bool) {}
+    fn node_remove_unmasked_branches(&mut self, _key: &[u8], _mask: ByteMask, _prune_limit: usize) {}
     fn node_is_empty(&self) -> bool { true }
     fn new_iter_token(&self) -> IterToken {
         0
@@ -115,7 +115,7 @@ impl<V: Clone + Send + Sync, A: Allocator> TrieNode<V, A> for EmptyNode {
     fn get_node_at_key(&self, _key: &[u8]) -> AbstractNodeRef<'_, V, A> {
         AbstractNodeRef::None
     }
-    fn take_node_at_key(&mut self, _key: &[u8], _prune: bool) -> Option<TrieNodeODRc<V, A>> {
+    fn take_node_at_key(&mut self, _key: &[u8], _prune_limit: usize) -> Option<TrieNodeODRc<V, A>> {
         None
     }
     fn pjoin_dyn(&self, other: TaggedNodeRef<V, A>) -> AlgebraicResult<TrieNodeODRc<V, A>> where V: Lattice {