Fix graft_child_maps below a long root path

benches/binary_keys.rs

@@ -45,6 +45,94 @@ fn binary_insert(bencher: Bencher, n: u64) {
     divan::black_box_drop(out)
 }
 
+// Every branch in these fixtures has at most two children. Short paths use
+// all eight three-byte binary keys; long paths branch at four spaced bytes.
+fn short_key(mask: u8) -> [u8; 3] {
+    [
+        b'0' + ((mask >> 2) & 1),
+        b'0' + ((mask >> 1) & 1),
+        b'0' + (mask & 1),
+    ]
+}
+
+fn seed_val(map: &mut PathMap<u64>, key: &[u8], val: u64) {
+    map.write_zipper_at_path(key).set_val(val);
+}
+
+fn short_map(target_len: usize, create: bool) -> PathMap<u64> {
+    let target = short_key(7);
+    let mut map = PathMap::new();
+    for mask in 0..8 {
+        let key = short_key(mask);
+        if !create || !key.starts_with(&target[..target_len]) {
+            seed_val(&mut map, &key, mask as u64);
+        }
+    }
+    if !create && target_len < target.len() {
+        seed_val(&mut map, &target[..target_len], 0);
+    }
+    assert_eq!(map.path_exists_at(&target[..target_len]), !create);
+    map
+}
+
+fn long_key(len: usize, mask: u8) -> Vec<u8> {
+    let mut key = vec![b'-'; len];
+    for (bit, index) in [0, len / 4, len / 2, 3 * len / 4].into_iter().enumerate() {
+        key[index] = b'0' + ((mask >> (3 - bit)) & 1);
+    }
+    key
+}
+
+fn long_map(len: usize, create: bool) -> PathMap<u64> {
+    let mut map = PathMap::new();
+    for mask in 0..16 {
+        if !create || mask != 15 {
+            seed_val(&mut map, &long_key(len, mask), mask as u64);
+        }
+    }
+    assert_eq!(map.path_exists_at(long_key(len, 15)), !create);
+    map
+}
+
+#[divan::bench(sample_size = 64, args = [0usize, 1, 2, 3])]
+fn binary_set_val_at_short_replace(bencher: Bencher, key_len: usize) {
+    let key = short_key(7);
+    let mut map = short_map(key_len, false);
+    bencher.bench_local(|| {
+        black_box(&mut map).set_val_at(black_box(&key[..key_len]), black_box(1));
+    });
+}
+
+// The empty path is the root, so creating a new path starts at length one.
+#[divan::bench(sample_size = 16, args = [1usize, 2, 3])]
+fn binary_set_val_at_short_create(bencher: Bencher, key_len: usize) {
+    let key = short_key(7);
+    let out = bencher.with_inputs(|| short_map(key_len, true)).bench_local_values(|mut map| {
+        black_box(&mut map).set_val_at(black_box(&key[..key_len]), black_box(1));
+        map
+    });
+    divan::black_box_drop(out);
+}
+
+#[divan::bench(args = [160usize, 256])]
+fn binary_set_val_at_long_replace(bencher: Bencher, key_len: usize) {
+    let key = long_key(key_len, 15);
+    let mut map = long_map(key_len, false);
+    bencher.bench_local(|| {
+        black_box(&mut map).set_val_at(black_box(&key), black_box(1));
+    });
+}
+
+#[divan::bench(sample_size = 16, args = [160usize, 256])]
+fn binary_set_val_at_long_create(bencher: Bencher, key_len: usize) {
+    let key = long_key(key_len, 15);
+    let out = bencher.with_inputs(|| long_map(key_len, true)).bench_local_values(|mut map| {
+        black_box(&mut map).set_val_at(black_box(&key), black_box(1));
+        map
+    });
+    divan::black_box_drop(out);
+}
+
 #[divan::bench(args = [250, 500, 1000, 2000, 4000, 8000])]
 fn binary_get(bencher: Bencher, n: u64) {
 

src/experimental.rs

@@ -159,6 +159,7 @@ impl <V: TrieValue + 'static, A: Allocator> ZipperWriting<V, A> for NullZipper {
     fn get_val_or_set_mut(&mut self, default: V) -> &mut V { Box::leak(Box::new(default)) }
     fn get_val_or_set_mut_with<F>(&mut self, func: F) -> &mut V where F: FnOnce() -> V { Box::leak(Box::new(func())) }
     fn set_val(&mut self, _val: V) -> Option<V> { None }
+    fn set_val_at<K: AsRef<[u8]>>(&mut self, path: K, val: V) -> Option<V> { None }
     fn remove_val(&mut self, _prune: bool) -> Option<V> { None }
     fn zipper_head<'z>(&'z mut self) -> Self::ZipperHead<'z> { todo!() }
     fn graft<Z: ZipperInfallibleSubtries<V, A>>(&mut self, _read_zipper: &Z) {}

src/trie_map.rs

@@ -337,21 +337,19 @@ impl<V: Clone + Send + Sync + Unpin, A: Allocator> PathMap<V, A> {
         self.path_exists_at(k)
     }
 
-    /// Inserts `v` into the map at `path`.  Panics if `path` has a zero length
+    /// Inserts `v` into the map at `path`.
     ///
     /// Returns `Some(replaced_val)` if an existing value was replaced, otherwise returns `None` if
     /// the value was added to the map without replacing anything.
     pub fn set_val_at<K: AsRef<[u8]>>(&mut self, path: K, v: V) -> Option<V> {
         let path = path.as_ref();
-
-        //NOTE: Here is the old impl traversing without the zipper.  Kept here for benchmarking purposes
-        // However, the zipper version is basically identical performance, within the margin of error 
-        // traverse_to_leaf_static_result(&mut self.root, k,
-        // |node, remaining_key| node.node_set_val(remaining_key, v),
-        // |_new_leaf_node, _remaining_key| None)
-
-        let mut zipper = self.write_zipper_at_path(path);
-        zipper.set_val(v)
+        if path.is_empty() {
+            return core::mem::replace(self.root_val_mut(), Some(v));
+        }
+        let (old_val, _) = with_node_at_path_mut(self.get_or_init_root_mut(), path,
+            |node, remaining_key| node.node_set_val(remaining_key, v),
+            |_, _| (None, true));
+        old_val
     }
 
     /// Alias for [Self::set_val_at], so `PathMap` "feels" like other Rust collections

src/trie_node.rs

@@ -2659,6 +2659,26 @@ pub(crate) fn node_along_path_mut<'a, 'k, V: Clone + Send + Sync, A: Allocator>(
     (key, node)
 }
 
+/// Applies a node operation at a path, replacing the node if it needs to be upgraded.
+#[inline]
+pub(crate) fn with_node_at_path_mut<V, A, NodeF, RetryF, R>(root: &mut TrieNodeODRc<V, A>, path: &[u8], node_f: NodeF, retry_f: RetryF) -> R
+where
+    V: Clone + Send + Sync,
+    A: Allocator,
+    NodeF: FnOnce(&mut TaggedNodeRefMut<'_, V, A>, &[u8]) -> Result<R, TrieNodeODRc<V, A>>,
+    RetryF: FnOnce(&mut TaggedNodeRefMut<'_, V, A>, &[u8]) -> R,
+{
+    debug_assert!(!path.is_empty());
+    let (remaining_key, node) = node_along_path_mut(root, path, true);
+    match node_f(&mut node.make_mut(), remaining_key) {
+        Ok(result) => result,
+        Err(replacement_node) => {
+            *node = replacement_node;
+            retry_f(&mut node.make_mut(), remaining_key)
+        }
+    }
+}
+
 /// Ensures the node is a CellByteNode
 ///
 /// Returns `true` if the node was upgraded and `false` if it already was a CellByteNode