Add reusable zipper buffers

src/trie_map.rs

@@ -229,6 +229,12 @@ impl<V: Clone + Send + Sync + Unpin, A: Allocator> PathMap<V, A> {
         ReadZipperUntracked::new_with_node_and_path_internal_in(self.root().unwrap(), &[], 0, root_val, self.alloc.clone())
     }
 
+    /// Creates a new read-only [Zipper], starting at the root of a `PathMap`, that reuses the
+    /// allocations held by `buffer`.  See [ZipperBuffer]
+    pub fn read_zipper_with_buffer<'a>(&'a self, buffer: ZipperBuffer<V, A>) -> ReadZipperUntracked<'a, 'static, V, A> {
+        ReadZipperUntracked::new_with_buffer(self, buffer)
+    }
+
     /// Creates a new read-only [Zipper], with the specified path from the root of the map; This method is much more
     /// efficient than [read_zipper_at_path](Self::read_zipper_at_path), but means the resulting zipper is bound by
     /// the `'path` lifetime
@@ -260,6 +266,15 @@ impl<V: Clone + Send + Sync + Unpin, A: Allocator> PathMap<V, A> {
         WriteZipperUntracked::new_with_node_and_path_internal_in(root_node, Some(root_val), &[], 0, self.alloc.clone())
     }
 
+    /// Creates a new [write zipper](ZipperWriting) starting at the root of the `PathMap`, that reuses the
+    /// allocations held by `buffer`.  See [WriteZipperBuffer]
+    pub fn write_zipper_with_buffer(&mut self, buffer: WriteZipperBuffer<V, A>) -> WriteZipperUntracked<'_, 'static, V, A> {
+        self.ensure_root();
+        let root_node = self.root.get_mut().as_mut().unwrap();
+        let root_val = self.root_val.get_mut();
+        WriteZipperUntracked::new_with_node_and_buffer_internal_in(root_node, Some(root_val), &[], 0, self.alloc.clone(), buffer)
+    }
+
     /// Creates a new [write zipper](ZipperWriting) with the specified path from the root of the map
     pub fn write_zipper_at_path<'a, 'path>(&'a mut self, path: &'path[u8]) -> WriteZipperUntracked<'a, 'path, V, A> {
         self.ensure_root();

src/write_zipper.rs

@@ -424,6 +424,29 @@ impl<V: Clone + Send + Sync, Z, A: Allocator> WriteZipperPriv<V, A> for &mut Z w
     fn alloc(&self) -> A { (**self).alloc() }
 }
 
+// ***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---
+// WriteZipperBuffer
+// ***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---
+
+/// Buffers taken from a write zipper for reuse by another write zipper.
+///
+/// A write zipper's path buffer and node stack can grow substantially while it is being used.
+/// A `WriteZipperBuffer` keeps those allocations alive after the zipper is dropped, so they can be
+/// handed to [PathMap::write_zipper_with_buffer] or
+/// [ZipperCreation::write_zipper_at_exclusive_path_with_buffer].
+///
+/// A `WriteZipperBuffer` holds no cursor state: all its buffers are always empty, and only their
+/// capacity is carried over to the next zipper.
+pub struct WriteZipperBuffer<V: Clone + Send + Sync, A: Allocator = GlobalAlloc> {
+    prefix_buf: Vec<u8>,
+    prefix_idx: Vec<usize>,
+    stack: NodeStackVec<V, A>,
+}
+
+// SAFETY: The node stack is always empty, so the buffer doesn't refer to any node
+unsafe impl<V: Clone + Send + Sync, A: Allocator> Send for WriteZipperBuffer<V, A> {}
+unsafe impl<V: Clone + Send + Sync, A: Allocator> Sync for WriteZipperBuffer<V, A> {}
+
 // ***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---
 // WriteZipperTracked
 // ***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---
@@ -552,6 +575,21 @@ impl<'a, V: Clone + Send + Sync + Unpin, A: Allocator> WriteZipperTracked<'a, 's
         let core = WriteZipperCore::<'a, 'static, V, A>::new_with_node_and_cloned_path_internal_in(root_node, root_val, path, root_key_start, alloc);
         Self { z: core, _tracker: tracker }
     }
+    /// Same as [Self::new_with_node_and_cloned_path_internal_in], but reuses the allocations held by `buffer`
+    pub(crate) fn new_with_node_and_buffer_internal_in(root_node: &'a mut TrieNodeODRc<V, A>, root_val: Option<&'a mut Option<V>>, path: &[u8], root_key_start: usize, alloc: A, buffer: WriteZipperBuffer<V, A>, tracker: Option<ZipperTracker<TrackingWrite>>) -> Self {
+        let core = WriteZipperCore::<'a, 'static, V, A>::new_with_node_and_buffer_internal_in(root_node, root_val, path, root_key_start, alloc, buffer);
+        Self { z: core, _tracker: tracker }
+    }
+}
+
+impl<'a, 'path, V: Clone + Send + Sync + Unpin, A: Allocator> WriteZipperTracked<'a, 'path, V, A> {
+    /// Consumes the zipper and returns its reusable traversal buffers.
+    ///
+    /// NOTE: For a zipper created by a [ZipperHead], this does not prune the zipper's root prefix path the
+    /// way [ZipperCreation::cleanup_write_zipper] does
+    pub fn into_zipper_buffer(mut self) -> WriteZipperBuffer<V, A> {
+        self.z.take_zipper_buffer()
+    }
 }
 
 impl<'a, 'path, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ZipperWriting<V, A> for WriteZipperTracked<'a, 'path, V, A> {
@@ -724,6 +762,18 @@ impl <'a, 'path, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> WriteZipperU
     pub(crate) fn core(&mut self) -> &mut WriteZipperCore<'a, 'path, V, A> {
         &mut self.z
     }
+    /// Consumes the zipper and returns its reusable traversal buffers.
+    pub fn into_zipper_buffer(mut self) -> WriteZipperBuffer<V, A> {
+        self.z.take_zipper_buffer()
+    }
+}
+
+impl <'a, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> WriteZipperUntracked<'a, 'static, V, A> {
+    /// See [WriteZipperCore::new_with_node_and_buffer_internal_in]
+    pub(crate) fn new_with_node_and_buffer_internal_in(root_node: &'a mut TrieNodeODRc<V, A>, root_val: Option<&'a mut Option<V>>, path: &[u8], root_key_start: usize, alloc: A, buffer: WriteZipperBuffer<V, A>) -> Self {
+        let core = WriteZipperCore::<'a, 'static, V, A>::new_with_node_and_buffer_internal_in(root_node, root_val, path, root_key_start, alloc, buffer);
+        Self { z: core }
+    }
 }
 
 impl<'a, 'path, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ZipperWriting<V, A> for WriteZipperUntracked<'a, 'path, V, A> {
@@ -850,6 +900,10 @@ impl <V: Clone + Send + Sync + Unpin, A: Allocator> WriteZipperOwned<V, A> {
         let map = MaybeDangling::into_inner(self.map);
         *map
     }
+    /// Consumes the zipper and returns its reusable traversal buffers.
+    pub fn into_zipper_buffer(mut self) -> WriteZipperBuffer<V, A> {
+        self.z.take_zipper_buffer()
+    }
     /// Consumes the `WriteZipperOwned`, and returns a [ReadZipperOwned] in its place
     ///
     /// The returned read zipper will have the same root and focus as the the consumed write zipper.
@@ -1216,6 +1270,19 @@ impl<'a, 'path, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ZipperPathBuf
 }
 
 impl <'a, V: Clone + Send + Sync + Unpin, A: Allocator> WriteZipperCore<'a, 'static, V, A> {
+    /// Same as [Self::new_with_node_and_cloned_path_internal_in], but stores the path and node stack in
+    /// the allocations held by `buffer`
+    pub(crate) fn new_with_node_and_buffer_internal_in(root_node: &'a mut TrieNodeODRc<V, A>, root_val: Option<&'a mut Option<V>>, path: &[u8], root_key_start: usize, alloc: A, buffer: WriteZipperBuffer<V, A>) -> Self {
+        let WriteZipperBuffer { prefix_buf, prefix_idx, stack } = buffer;
+        let focus_stack = MutNodeStack::new_with_stack(root_node, stack);
+        debug_assert!((path.len()-root_key_start == 0) != (root_val.is_none())); //We must have either a node_path or a root_val, but never both
+        Self {
+            key: KeyFields::new_cloned_path_in_buffer(path, root_key_start, prefix_buf, prefix_idx),
+            root_val: root_val.map(|val| val as *mut Option<V>),
+            focus_stack,
+            alloc,
+        }
+    }
     pub(crate) fn new_with_node_and_cloned_path_in(root_node: &'a mut TrieNodeODRc<V, A>, root_val: Option<&'a mut Option<V>>, path: &[u8], root_prefix_len: usize, root_key_start: usize, alloc: A) -> Self {
         let (key, node) = node_along_path_mut(root_node, &path[root_key_start..], true);
 
@@ -1236,6 +1303,16 @@ impl <'a, V: Clone + Send + Sync + Unpin, A: Allocator> WriteZipperCore<'a, 'sta
 }
 
 impl <'a, 'path, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> WriteZipperCore<'a, 'path, V, A> {
+    /// Takes the buffer allocations out of the zipper for reuse.  The zipper is left without its path,
+    /// so it must be dropped right after this call
+    pub(crate) fn take_zipper_buffer(&mut self) -> WriteZipperBuffer<V, A> {
+        let mut prefix_buf = core::mem::take(&mut self.key.prefix_buf);
+        prefix_buf.clear();
+        let mut prefix_idx = core::mem::take(&mut self.key.prefix_idx);
+        prefix_idx.clear();
+        let stack = self.focus_stack.take_stack();
+        WriteZipperBuffer { prefix_buf, prefix_idx, stack }
+    }
     /// Creates a new zipper, with a path relative to a node
     pub(crate) fn new_with_node_and_path_in(root_node: &'a mut TrieNodeODRc<V, A>, root_val: Option<&'a mut Option<V>>, path: &'path [u8], root_prefix_len: usize, root_key_start: usize, alloc: A) -> Self {
         let (key, node) = node_along_path_mut(root_node, &path[root_key_start..], true);
@@ -2851,6 +2928,18 @@ fn make_parents_in<V: Clone + Send + Sync, A: Allocator>(path: &[u8], child_node
 }
 
 impl KeyFields<'static> {
+    /// Same as [Self::new_cloned_path], but copies `path` into `prefix_buf`, which must be empty
+    #[inline]
+    fn new_cloned_path_in_buffer(path: &[u8], root_key_start: usize, mut prefix_buf: Vec<u8>, prefix_idx: Vec<usize>) -> Self {
+        debug_assert!(prefix_buf.is_empty() && prefix_idx.is_empty());
+        prefix_buf.extend_from_slice(path);
+        Self {
+            origin_path: SliceOrLen::new_owned(path.len()),
+            root_key_start,
+            prefix_buf,
+            prefix_idx,
+        }
+    }
     #[inline]
     fn new_cloned_path(path: &[u8], root_key_start: usize) -> Self {
         let prefix_buf = path.to_vec();
@@ -2954,14 +3043,17 @@ impl<'k> KeyFields<'k> {
 // MutNodeStack: A replacement for MutCursor, follows the same pattern but stores TaggedNodeRefMut internally
 // ***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---
 
-use mut_node_stack::MutNodeStack;
+use mut_node_stack::{MutNodeStack, NodeStackVec};
 mod mut_node_stack {
     use smallvec::{SmallVec, smallvec};
     use core::ptr::NonNull;
     use core::marker::PhantomData;
     use crate::alloc::Allocator;
     use super::{TaggedNodeRef, TaggedNodeRefMut, TaggedNodePtr, TrieNodeODRc};
 
+    /// The storage behind a [MutNodeStack], which can be moved between stacks for reuse
+    pub type NodeStackVec<V, A> = SmallVec<[TaggedNodePtr<V, A>; 1]>;
+
     /// See [mutcursor::MutCursorRootedVec] for discussion about behavior
     pub struct MutNodeStack<'a, V: Clone + Send + Sync, A: Allocator> {
         root: Option<NonNull<TrieNodeODRc<V, A>>>,
@@ -2975,6 +3067,23 @@ mod mut_node_stack {
         pub fn new(root: &'a mut TrieNodeODRc<V, A>) -> Self {
             Self { root: Some(NonNull::from(root)), stack: smallvec![], phantom: PhantomData }
         }
+        /// Same as [Self::new], but uses `stack`, which must be empty, as the storage
+        #[inline]
+        pub(crate) fn new_with_stack(root: &'a mut TrieNodeODRc<V, A>, stack: NodeStackVec<V, A>) -> Self {
+            debug_assert!(stack.is_empty());
+            Self { root: Some(NonNull::from(root)), stack, phantom: PhantomData }
+        }
+        /// Returns the heap allocation backing the stack, if it has spilled out of its inline slot
+        #[cfg(test)]
+        pub fn heap_ptr(&self) -> Option<*const TaggedNodePtr<V, A>> {
+            self.stack.spilled().then(|| self.stack.as_ptr())
+        }
+        /// Moves to the root, and takes the (empty) storage out of the stack
+        #[inline]
+        pub fn take_stack(&mut self) -> NodeStackVec<V, A> {
+            self.stack.clear();
+            core::mem::take(&mut self.stack)
+        }
         #[inline]
         pub fn top(&self) -> Option<TaggedNodeRef<'_, V, A>> {
             match self.stack.last() {
@@ -3093,6 +3202,128 @@ mod tests {
     use crate::trie_node::*;
     use crate::alloc::GlobalAlloc;
 
+    /// Paths that branch 4 ways at each of 6 levels, so traversals cross several nodes
+    fn branchy_map() -> PathMap<usize> {
+        let mut map = PathMap::new();
+        let mut path = [0u8; 6];
+        for i in 0..4usize.pow(6) {
+            for (level, byte) in path.iter_mut().enumerate() {
+                *byte = b'a' + ((i >> (2 * level)) & 3) as u8;
+            }
+            map.set_val_at(&path, i);
+        }
+        map
+    }
+
+    fn buffer_ptrs<V: Clone + Send + Sync + Unpin, A: crate::alloc::Allocator>(z: &WriteZipperCore<'_, '_, V, A>) -> (*const u8, *const usize, Option<*const TaggedNodePtr<V, A>>) {
+        (z.key.prefix_buf.as_ptr(), z.key.prefix_idx.as_ptr(), z.focus_stack.heap_ptr())
+    }
+
+    #[test]
+    fn write_zipper_with_buffer_reuses_allocations() {
+        let mut map = branchy_map();
+        let mut wz = map.write_zipper();
+        wz.reserve_buffers(4096, 512);
+        wz.descend_to(b"dcbadc");
+        assert_eq!(wz.val(), Some(&branchy_val(b"dcbadc")));
+        let ptrs = buffer_ptrs(&wz.z);
+        assert!(ptrs.2.is_some(), "node stack should have spilled to the heap");
+        let buffer = wz.into_zipper_buffer();
+
+        let mut wz = map.write_zipper_with_buffer(buffer);
+        assert_eq!(wz.path(), b"");
+        assert_eq!(wz.origin_path(), b"");
+        wz.descend_to(b"abcdab");
+        assert_eq!(buffer_ptrs(&wz.z), ptrs);
+        assert_eq!(wz.set_val(7), Some(branchy_val(b"abcdab")));
+        assert_eq!(buffer_ptrs(&wz.z), ptrs);
+
+        // Move the buffer to a zipper on another map
+        let buffer = wz.into_zipper_buffer();
+        assert_eq!(map.get_val_at(b"abcdab"), Some(&7));
+        let mut other = PathMap::<usize>::new();
+        let mut wz = other.write_zipper_with_buffer(buffer);
+        wz.descend_to(b"xyz");
+        wz.set_val(1);
+        assert_eq!(wz.z.key.prefix_buf.as_ptr(), ptrs.0);
+        drop(wz);
+        assert_eq!(other.get_val_at(b"xyz"), Some(&1));
+        assert_eq!(other.val_count(), 1);
+    }
+
+    /// The value [branchy_map] stores at `path`
+    fn branchy_val(path: &[u8]) -> usize {
+        path.iter().enumerate().map(|(level, byte)| ((byte - b'a') as usize) << (2 * level)).sum()
+    }
+
+    #[test]
+    fn write_zipper_with_unallocated_buffer() {
+        let mut map = PathMap::<usize>::new();
+        let buffer = map.write_zipper().into_zipper_buffer();
+        let mut wz = map.write_zipper_with_buffer(buffer);
+        wz.descend_to(b"alphabet");
+        wz.set_val(1);
+        drop(wz);
+        assert_eq!(map.get_val_at(b"alphabet"), Some(&1));
+
+        let buffer = WriteZipperOwned::<usize>::new().into_zipper_buffer();
+        let mut wz = map.write_zipper_with_buffer(buffer);
+        wz.descend_to(b"alpine");
+        wz.set_val(2);
+        drop(wz);
+        assert_eq!(map.get_val_at(b"alpine"), Some(&2));
+    }
+
+    #[test]
+    fn write_zipper_buffer_drop_and_send() {
+        fn assert_send_sync<T: Send + Sync>(_: &T) {}
+        let mut map = branchy_map();
+        let mut wz = map.write_zipper();
+        wz.reserve_buffers(256, 32);
+        wz.descend_to(b"abcdab");
+        let buffer = wz.into_zipper_buffer();
+        assert_send_sync(&buffer);
+        drop(buffer);
+    }
+
+    #[test]
+    fn zipper_head_write_zipper_with_buffer() {
+        let mut map = branchy_map();
+        let zh = map.zipper_head();
+
+        let mut wz = zh.write_zipper_at_exclusive_path(b"ab").unwrap();
+        wz.reserve_buffers(4096, 512);
+        wz.descend_to(b"cdab");
+        let (prefix_ptr, idx_ptr, _) = buffer_ptrs(&wz.z);
+        let buffer = wz.into_zipper_buffer();
+
+        let mut wz = zh.write_zipper_at_exclusive_path_with_buffer(b"ba", buffer).unwrap();
+        assert_eq!(wz.path(), b"");
+        assert_eq!(wz.origin_path(), b"ba");
+        assert_eq!(wz.val_count(), 4usize.pow(4));
+        wz.descend_to(b"cdab");
+        assert_eq!(wz.z.key.prefix_buf.as_ptr(), prefix_ptr);
+        assert_eq!(wz.z.key.prefix_idx.as_ptr(), idx_ptr);
+        wz.set_val(100);
+
+        // A conflicting zipper is refused while `wz` is alive
+        let spare = PathMap::<usize>::new().write_zipper().into_zipper_buffer();
+        assert!(zh.write_zipper_at_exclusive_path_with_buffer(b"bacd", spare).is_err());
+
+        let buffer = wz.into_zipper_buffer();
+        let mut wz = unsafe{ zh.write_zipper_at_exclusive_path_with_buffer_unchecked(b"zz", buffer) };
+        assert_eq!(wz.origin_path(), b"zz");
+        wz.descend_to(b"top");
+        wz.set_val(200);
+        assert_eq!(wz.z.key.prefix_buf.as_ptr(), prefix_ptr);
+        zh.cleanup_write_zipper(wz);
+        drop(zh);
+
+        assert_eq!(map.get_val_at(b"bacdab"), Some(&100));
+        assert_eq!(map.get_val_at(b"zztop"), Some(&200));
+        assert_eq!(map.get_val_at(b"abcdab"), Some(&branchy_val(b"abcdab")));
+    }
+
     /// `PathMap::iter` exposes values but not dangling paths.  Check every value and, when a
     /// test cares about dangling structure, its explicitly supplied path probes as well.
     fn assert_map_unchanged<V>(before: &PathMap<V>, after: &PathMap<V>, path_probes: &[&[u8]])