From 81a420c1fd83185611b0f781c1c8c6fa475a001b Mon Sep 17 00:00:00 2001 From: Igor Malovitsa Date: Sat, 24 Oct 2026 00:42:41 +0000 Subject: [PATCH 1/3] Add reusable zipper buffers A ZipperBuffer holds the path and ancestor-stack allocations of a read zipper, so a later zipper can reuse them instead of growing new ones. - ReadZipper{Untracked,Tracked,Owned}::into_zipper_buffer consumes a zipper and keeps its (emptied) buffers - PathMap::read_zipper_with_buffer creates a zipper at the map root on top of a buffer - ReadZipper{Untracked,Tracked}::clone_with_buffer is Clone, but writes the path and ancestor stack into the buffer's allocations The ancestor stack's element type borrows from the map, so the buffer stores that allocation with its element type erased; traversal code is unchanged. --- src/trie_map.rs | 6 ++ src/zipper.rs | 254 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 260 insertions(+) diff --git a/src/trie_map.rs b/src/trie_map.rs index 37855a2d..3cef3f9c 100644 --- a/src/trie_map.rs +++ b/src/trie_map.rs @@ -229,6 +229,12 @@ impl PathMap { 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) -> 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 diff --git a/src/zipper.rs b/src/zipper.rs index d87d79ef..ddd46e33 100644 --- a/src/zipper.rs +++ b/src/zipper.rs @@ -1554,6 +1554,24 @@ impl<'a, 'path, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ReadZipperTra let core = ReadZipperCore::new_with_node_and_cloned_path_in(root_node, owned_root, path, root_prefix_len, root_key_start, root_val, alloc); Self { z: core, tracker } } + + /// Consumes the zipper and returns its reusable traversal buffers. + pub fn into_zipper_buffer(mut self) -> ZipperBuffer { + let placeholder = ReadZipperCore::new_with_node_and_path_internal_in( + OwnedOrBorrowed::None, + &[], + 0, + None, + self.z.alloc.clone(), + ); + core::mem::replace(&mut self.z, placeholder).into_zipper_buffer() + } + + /// Same as [Clone::clone], but the new zipper stores its path and ancestor + /// stack in the allocations held by `buffer`. + pub fn clone_with_buffer(&self, buffer: ZipperBuffer) -> Self { + Self { z: self.z.clone_with_buffer(buffer), tracker: self.tracker.clone() } + } } //GOAT, the standard prototype of IntoIterator isn't compatible with ReadZipperTracked anymore because @@ -1667,6 +1685,25 @@ impl<'a, 'path, V: Clone + Send + Sync + Unpin + 'a, A: Allocator + 'a> ReadZipp pub(crate) fn new_forked_with_inner_zipper(core: ReadZipperCore<'a, 'path, V, A>) -> Self { ReadZipperUntracked{ z: core } } + + /// Creates a new zipper at the root of `map`, reusing the allocations held by `buffer`. + /// See [PathMap::read_zipper_with_buffer] + pub(crate) fn new_with_buffer(map: &'a PathMap, buffer: ZipperBuffer) -> ReadZipperUntracked<'a, 'static, V, A> { + let mut zipper = map.read_zipper(); + zipper.z.install_buffer(buffer); + zipper + } + + /// Consumes the zipper and returns its reusable traversal buffers. + pub fn into_zipper_buffer(self) -> ZipperBuffer { + self.z.into_zipper_buffer() + } + + /// Same as [Clone::clone], but the new zipper stores its path and ancestor + /// stack in the allocations held by `buffer`. + pub fn clone_with_buffer(&self, buffer: ZipperBuffer) -> Self { + Self { z: self.z.clone_with_buffer(buffer) } + } } impl<'a, 'path, V: Clone + Send + Sync + Unpin + 'a, A: Allocator + 'a> ReadZipperUntracked<'a, 'path, V, A> { @@ -1754,6 +1791,11 @@ impl ReadZipperOwned ZipperBuffer { + self.z.into_zipper_buffer() + } } @@ -1824,6 +1866,65 @@ pub(crate) const EXPECTED_DEPTH: usize = 16; /// Size in bytes to preallocate path storage in the zipper pub(crate) const EXPECTED_PATH_LEN: usize = 64; +/// Buffers taken from a read zipper for reuse by another read zipper. +/// +/// A zipper's path buffer and ancestor stack can grow substantially while it +/// is being used. A [`ZipperBuffer`] keeps those allocations alive after the +/// zipper is dropped, so they can be handed to +/// [PathMap::read_zipper_with_buffer] or [ReadZipperUntracked::clone_with_buffer]. +/// +/// A `ZipperBuffer` holds no cursor state: both buffers are always empty, and +/// only their capacity is carried over to the next zipper. +pub struct ZipperBuffer { + prefix_buf: Vec, + // The ancestor stack's element type borrows from the map, so its + // allocation is stored with the element type erased. + ancestors: ErasedVecAlloc, + phantom: core::marker::PhantomData (V, A)>, +} + +/// The allocation of an empty `Vec`, with the element type erased so it is no +/// longer bound by any lifetime in that type +struct ErasedVecAlloc { + ptr: core::ptr::NonNull, + cap: usize, + elem: core::alloc::Layout, +} + +// SAFETY: An `ErasedVecAlloc` owns an allocation with no live elements in it +unsafe impl Send for ErasedVecAlloc {} +unsafe impl Sync for ErasedVecAlloc {} + +impl ErasedVecAlloc { + fn from_vec(mut vec: Vec) -> Self { + vec.clear(); + let mut vec = core::mem::ManuallyDrop::new(vec); + let ptr = core::ptr::NonNull::new(vec.as_mut_ptr()).unwrap().cast(); + Self { ptr, cap: vec.capacity(), elem: core::alloc::Layout::new::() } + } + + /// Returns an empty `Vec` that owns the allocation. `T` must have the same layout + /// as the element type of the `Vec` the allocation came from. The zipper only + /// converts between ancestor stacks that differ in lifetime, so this always holds + fn into_vec(self) -> Vec { + assert_eq!(self.elem, core::alloc::Layout::new::()); + let this = core::mem::ManuallyDrop::new(self); + // SAFETY: The allocation came from a `Vec` with capacity `cap` whose element + // has the same size and alignment as `T`, and it holds no elements + unsafe { Vec::from_raw_parts(this.ptr.as_ptr().cast::(), 0, this.cap) } + } +} + +impl Drop for ErasedVecAlloc { + fn drop(&mut self) { + let size = self.elem.size() * self.cap; + if size > 0 { + // SAFETY: This is the layout `Vec` allocated with: `cap` elements of `elem` + unsafe { std::alloc::dealloc(self.ptr.as_ptr(), core::alloc::Layout::from_size_align_unchecked(size, self.elem.align())) } + } + } +} + pub(crate) mod read_zipper_core { use crate::PathMap; use crate::zipper::*; @@ -2889,6 +2990,61 @@ pub(crate) mod read_zipper_core { } } + /// Consumes the zipper and keeps its buffer allocations for reuse. + pub(crate) fn into_zipper_buffer(mut self) -> ZipperBuffer { + // Ancestor entries contain references into this zipper's map, so + // no entry may survive. Only the allocations are kept. + self.prefix_buf.clear(); + self.ancestors.clear(); + ZipperBuffer { + prefix_buf: self.prefix_buf, + ancestors: ErasedVecAlloc::from_vec(self.ancestors), + phantom: core::marker::PhantomData, + } + } + + /// Moves the allocations in `buffer` into a freshly created zipper, which must + /// not have moved yet (its buffers are still uninitialized) + pub(crate) fn install_buffer(&mut self, buffer: ZipperBuffer) { + debug_assert_eq!(self.prefix_buf.capacity(), 0); + debug_assert!(self.ancestors.is_empty()); + let ZipperBuffer { mut prefix_buf, ancestors, phantom: _ } = buffer; + // A non-zero capacity marks `prefix_buf` as initialized, in which case + // it must begin with the origin path. See `reserve_buffers` + if prefix_buf.capacity() > 0 { + prefix_buf.extend_from_slice(self.origin_path.as_slice()); + } + self.prefix_buf = prefix_buf; + self.ancestors = ancestors.into_vec(); + } + + /// Same as [Clone::clone], but writes the path and ancestor stack into the + /// allocations held by `buffer` + pub(crate) fn clone_with_buffer(&self, buffer: ZipperBuffer) -> Self { + let ZipperBuffer { mut prefix_buf, ancestors, phantom: _ } = buffer; + let mut ancestors: Vec<_> = ancestors.into_vec(); + if self.prefix_buf.capacity() > 0 { + prefix_buf.extend_from_slice(&self.prefix_buf); + ancestors.extend_from_slice(&self.ancestors); + } else if prefix_buf.capacity() > 0 { + // `self` hasn't moved yet, so its path is just the origin path + debug_assert!(self.ancestors.is_empty()); + prefix_buf.extend_from_slice(self.origin_path.as_slice()); + } + Self { + origin_path: self.origin_path.clone(), + root_key_start: self.root_key_start, + root_parent_key_start: self.root_parent_key_start, + root_val: self.root_val, + root_node: self.root_node.clone(), + focus_node: self.focus_node.clone(), + focus_iter_token: NODE_ITER_INVALID, + prefix_buf, + ancestors, + alloc: self.alloc.clone(), + } + } + /// Returns the length of the `self.path()`, saving a couple instructions, but is internal because it may panic #[inline(always)] fn path_len(&self) -> usize { @@ -3514,6 +3670,104 @@ pub(crate) mod read_zipper_core { assert_eq!(&rz.path(), b"-howdy"); assert_eq!(rz.origin_path(), b"hi-howdy"); } + + #[test] + fn read_zipper_with_buffer_reuses_allocations() { + let mut map = PathMap::::new(); + map.set_val_at(b"alpha", 1); + map.set_val_at(b"alphabet", 2); + + let mut zipper = map.read_zipper(); + zipper.reserve_buffers(4096, 512); + zipper.descend_to(b"alphabet"); + let prefix_ptr = zipper.z.prefix_buf.as_ptr(); + let ancestors_ptr = zipper.z.ancestors.as_ptr(); + let prefix_capacity = zipper.z.prefix_buf.capacity(); + let ancestors_capacity = zipper.z.ancestors.capacity(); + + let buffer = zipper.into_zipper_buffer(); + + let mut other_map = PathMap::::new(); + other_map.set_val_at(b"alphabet", 3); + let mut zipper = other_map.read_zipper_with_buffer(buffer); + + assert_eq!(zipper.path(), b""); + assert_eq!(zipper.origin_path(), b""); + assert_eq!(zipper.z.prefix_buf.as_ptr(), prefix_ptr); + assert_eq!(zipper.z.prefix_buf.capacity(), prefix_capacity); + assert_eq!(zipper.z.ancestors.as_ptr(), ancestors_ptr); + assert_eq!(zipper.z.ancestors.capacity(), ancestors_capacity); + + assert!(zipper.descend_to_check(b"alphabet")); + assert_eq!(zipper.val(), Some(&3)); + assert_eq!(zipper.z.prefix_buf.as_ptr(), prefix_ptr); + assert_eq!(zipper.z.ancestors.as_ptr(), ancestors_ptr); + } + + #[test] + fn read_zipper_with_unallocated_buffer() { + let mut map = PathMap::::new(); + map.set_val_at(b"alphabet", 1); + + let buffer = map.read_zipper().into_zipper_buffer(); + let mut zipper = map.read_zipper_with_buffer(buffer); + assert!(zipper.descend_to_check(b"alphabet")); + assert_eq!(zipper.val(), Some(&1)); + } + + #[test] + fn zipper_buffer_drop_and_send() { + fn assert_send_sync(_: &T) {} + let mut map = PathMap::::new(); + map.set_val_at(b"alphabet", 1); + + let mut zipper = map.read_zipper(); + zipper.reserve_buffers(256, 32); + zipper.descend_to(b"alphabet"); + let buffer = zipper.into_zipper_buffer(); + assert_send_sync(&buffer); + drop(buffer); + } + + #[test] + fn read_zipper_clone_with_buffer() { + let mut map = PathMap::::new(); + map.set_val_at(b"alphabet", 1); + map.set_val_at(b"alpine", 2); + + let mut donor = map.read_zipper(); + donor.reserve_buffers(4096, 512); + let prefix_ptr = donor.z.prefix_buf.as_ptr(); + let ancestors_ptr = donor.z.ancestors.as_ptr(); + + // Clone a zipper that has moved below a non-root origin + let mut zipper = map.read_zipper_at_borrowed_path(b"al"); + zipper.descend_to(b"pha"); + let mut cloned = zipper.clone_with_buffer(donor.into_zipper_buffer()); + assert_eq!(cloned.z.prefix_buf.as_ptr(), prefix_ptr); + assert_eq!(cloned.z.ancestors.as_ptr(), ancestors_ptr); + assert_eq!(cloned.path(), b"pha"); + assert_eq!(cloned.origin_path(), b"alpha"); + assert!(cloned.descend_to_check(b"bet")); + assert_eq!(cloned.val(), Some(&1)); + assert_eq!(cloned.ascend(6), 6); + assert!(cloned.descend_to_check(b"pine")); + assert_eq!(cloned.val(), Some(&2)); + assert_eq!(cloned.z.prefix_buf.as_ptr(), prefix_ptr); + + // Clone a zipper that has not moved yet, so its buffers are uninitialized + let mut donor = map.read_zipper(); + donor.reserve_buffers(4096, 512); + let prefix_ptr = donor.z.prefix_buf.as_ptr(); + let fresh = map.read_zipper_at_borrowed_path(b"alp"); + let mut cloned = fresh.clone_with_buffer(donor.into_zipper_buffer()); + assert_eq!(cloned.z.prefix_buf.as_ptr(), prefix_ptr); + assert_eq!(cloned.origin_path(), b"alp"); + assert!(cloned.descend_to_check(b"habet")); + assert_eq!(cloned.val(), Some(&1)); + assert_eq!(cloned.origin_path(), b"alphabet"); + assert_eq!(cloned.z.prefix_buf.as_ptr(), prefix_ptr); + } } use read_zipper_core::*; From 3671f88864f53197939733c15b6a7b560bb5f77e Mon Sep 17 00:00:00 2001 From: Igor Malovitsa Date: Sat, 24 Oct 2026 01:00:30 +0000 Subject: [PATCH 2/3] Add reusable write zipper buffers A WriteZipperBuffer holds the path, key-index and node-stack allocations of a write zipper, so a later write zipper can reuse them. - WriteZipper{Untracked,Tracked,Owned}::into_zipper_buffer consumes a zipper and keeps its (emptied) buffers - PathMap::write_zipper_with_buffer creates a write zipper at the map root on top of a buffer - ZipperCreation::write_zipper_at_exclusive_path_with_buffer(_unchecked) creates a ZipperHead write zipper on top of a buffer --- src/trie_map.rs | 9 ++ src/write_zipper.rs | 233 +++++++++++++++++++++++++++++++++++++++++++- src/zipper_head.rs | 39 ++++++++ 3 files changed, 280 insertions(+), 1 deletion(-) diff --git a/src/trie_map.rs b/src/trie_map.rs index 3cef3f9c..5c0b1dfd 100644 --- a/src/trie_map.rs +++ b/src/trie_map.rs @@ -266,6 +266,15 @@ impl PathMap { 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) -> 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(); diff --git a/src/write_zipper.rs b/src/write_zipper.rs index 624ab95d..23a48fbe 100644 --- a/src/write_zipper.rs +++ b/src/write_zipper.rs @@ -424,6 +424,29 @@ impl WriteZipperPriv 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 { + prefix_buf: Vec, + prefix_idx: Vec, + stack: NodeStackVec, +} + +// SAFETY: The node stack is always empty, so the buffer doesn't refer to any node +unsafe impl Send for WriteZipperBuffer {} +unsafe impl Sync for WriteZipperBuffer {} + // ***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***---***--- // 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, root_val: Option<&'a mut Option>, path: &[u8], root_key_start: usize, alloc: A, buffer: WriteZipperBuffer, tracker: Option>) -> 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 { + self.z.take_zipper_buffer() + } } impl<'a, 'path, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ZipperWriting 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 { + 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, root_val: Option<&'a mut Option>, path: &[u8], root_key_start: usize, alloc: A, buffer: WriteZipperBuffer) -> 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 for WriteZipperUntracked<'a, 'path, V, A> { @@ -850,6 +900,10 @@ impl WriteZipperOwned { 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 { + 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, root_val: Option<&'a mut Option>, path: &[u8], root_key_start: usize, alloc: A, buffer: WriteZipperBuffer) -> 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), + focus_stack, + alloc, + } + } pub(crate) fn new_with_node_and_cloned_path_in(root_node: &'a mut TrieNodeODRc, root_val: Option<&'a mut Option>, 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 { + 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, root_val: Option<&'a mut Option>, 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(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, prefix_idx: Vec) -> 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,7 +3043,7 @@ 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; @@ -2962,6 +3051,9 @@ mod mut_node_stack { 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 = SmallVec<[TaggedNodePtr; 1]>; + /// See [mutcursor::MutCursorRootedVec] for discussion about behavior pub struct MutNodeStack<'a, V: Clone + Send + Sync, A: Allocator> { root: Option>>, @@ -2975,6 +3067,23 @@ mod mut_node_stack { pub fn new(root: &'a mut TrieNodeODRc) -> 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 fn new_with_stack(root: &'a mut TrieNodeODRc, stack: NodeStackVec) -> 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> { + 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 { + self.stack.clear(); + core::mem::take(&mut self.stack) + } #[inline] pub fn top(&self) -> Option> { 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 { + 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(z: &WriteZipperCore<'_, '_, V, A>) -> (*const u8, *const usize, Option<*const TaggedNodePtr>) { + (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::::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::::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::::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) {} + 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::::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(before: &PathMap, after: &PathMap, path_probes: &[&[u8]]) diff --git a/src/zipper_head.rs b/src/zipper_head.rs index a379dc39..16830ede 100644 --- a/src/zipper_head.rs +++ b/src/zipper_head.rs @@ -55,6 +55,14 @@ pub trait ZipperCreation<'trie, V: Clone + Send + Sync, A: Allocator = GlobalAll /// break the rules, it's just an optimization that affects when to spend time enforcing them. unsafe fn write_zipper_at_exclusive_path_unchecked<'a, K: AsRef<[u8]>>(&'a self, path: K) -> WriteZipperTracked<'a, 'static, V, A> where 'trie: 'a; + /// Same as [write_zipper_at_exclusive_path](ZipperCreation::write_zipper_at_exclusive_path), but the + /// returned zipper reuses the allocations held by `buffer`. `buffer` is dropped if there is a conflict + fn write_zipper_at_exclusive_path_with_buffer<'a, K: AsRef<[u8]>>(&'a self, path: K, buffer: WriteZipperBuffer) -> Result, Conflict> where 'trie: 'a; + + /// Same as [write_zipper_at_exclusive_path_unchecked](ZipperCreation::write_zipper_at_exclusive_path_unchecked), + /// but the returned zipper reuses the allocations held by `buffer` + unsafe fn write_zipper_at_exclusive_path_with_buffer_unchecked<'a, K: AsRef<[u8]>>(&'a self, path: K, buffer: WriteZipperBuffer) -> WriteZipperTracked<'a, 'static, V, A> where 'trie: 'a; + //GOAT-TrackedOwnedZippers // /// Creates a [WriteZipperOwned] from the specified path by temporarily cutting the trie // /// @@ -297,6 +305,37 @@ impl<'trie, Z, V: 'trie + Clone + Send + Sync + Unpin, A: Allocator + 'trie> Zip WriteZipperTracked::new_with_node_and_cloned_path_internal_in(zipper_root_node, Some(zipper_root_val), path, path.len(), z.alloc.clone(), tracker) }) } + fn write_zipper_at_exclusive_path_with_buffer<'a, K: AsRef<[u8]>>(&'a self, path: K, buffer: WriteZipperBuffer) -> Result, Conflict> where 'trie: 'a { + let path = path.as_ref(); + let zipper_tracker = ZipperTracker::::new(self.tracker_paths().clone(), path)?; + self.with_inner_core_z(|z| { + let (zipper_root_node, zipper_root_val) = prepare_exclusive_write_path(z, path); + //SAFETY: See similar code in `read_zipper_at_borrowed_path` for more discussion + let zipper_root_node: &'trie mut TrieNodeODRc = unsafe{ &mut *(zipper_root_node as *mut _) }; + let zipper_root_val: &'trie mut Option = unsafe{ &mut *(zipper_root_val as *mut _) }; + + Ok(WriteZipperTracked::new_with_node_and_buffer_internal_in(zipper_root_node, Some(zipper_root_val), path, path.len(), z.alloc.clone(), buffer, Some(zipper_tracker))) + }) + } + unsafe fn write_zipper_at_exclusive_path_with_buffer_unchecked<'a, K: AsRef<[u8]>>(&'a self, path: K, buffer: WriteZipperBuffer) -> WriteZipperTracked<'a, 'static, V, A> where 'trie: 'a { + let path = path.as_ref(); + self.with_inner_core_z(|z| { + let (zipper_root_node, zipper_root_val) = prepare_exclusive_write_path(z, path); + //SAFETY: The user is asserting that the paths won't conflict. + // See similar code in `read_zipper_at_borrowed_path` for more discussion + let zipper_root_node: &'trie mut TrieNodeODRc = unsafe{ &mut *(zipper_root_node as *mut _) }; + let zipper_root_val: &'trie mut Option = unsafe{ &mut *(zipper_root_val as *mut _) }; + + //We still keep the trackers in debug mode + #[cfg(debug_assertions)] + let tracker = Some(ZipperTracker::::new(self.tracker_paths().clone(), path) + .unwrap_or_else(|conflict| panic!("Fatal error. WriteZipper at {path:?} {conflict}"))); + #[cfg(not(debug_assertions))] + let tracker = None; + + WriteZipperTracked::new_with_node_and_buffer_internal_in(zipper_root_node, Some(zipper_root_val), path, path.len(), z.alloc.clone(), buffer, tracker) + }) + } fn cleanup_write_zipper + ZipperAbsolutePath>(&self, mut z: ChildZ) { let origin_path = z.take_root_prefix_path(); drop(z); From 622f05192f53961548e4398512c7041180282834 Mon Sep 17 00:00:00 2001 From: Igor Malovitsa Date: Mon, 26 Oct 2026 14:09:38 +0000 Subject: [PATCH 3/3] Fix UAF in ZipperHead & clone_with_buffer --- src/write_zipper.rs | 2 +- src/zipper.rs | 76 ++++++++++++++++++++++++++++++--------------- 2 files changed, 52 insertions(+), 26 deletions(-) diff --git a/src/write_zipper.rs b/src/write_zipper.rs index 23a48fbe..de20ef4e 100644 --- a/src/write_zipper.rs +++ b/src/write_zipper.rs @@ -3069,7 +3069,7 @@ mod mut_node_stack { } /// Same as [Self::new], but uses `stack`, which must be empty, as the storage #[inline] - pub fn new_with_stack(root: &'a mut TrieNodeODRc, stack: NodeStackVec) -> Self { + pub(crate) fn new_with_stack(root: &'a mut TrieNodeODRc, stack: NodeStackVec) -> Self { debug_assert!(stack.is_empty()); Self { root: Some(NonNull::from(root)), stack, phantom: PhantomData } } diff --git a/src/zipper.rs b/src/zipper.rs index ddd46e33..11933059 100644 --- a/src/zipper.rs +++ b/src/zipper.rs @@ -2040,8 +2040,27 @@ pub(crate) mod read_zipper_core { impl Clone for ReadZipperCore<'_, '_, V, A> where V: Clone { fn clone(&self) -> Self { + self.clone_into_buffers(Vec::new(), Vec::new()) + } + } + + impl<'a, V: Clone + Send + Sync, A: Allocator> ReadZipperCore<'a, '_, V, A> { + /// Copies the zipper, storing the copy's path and ancestor stack in `prefix_buf` + /// and `ancestors`, which must be empty. Shared by `clone` and `clone_with_buffer` + fn clone_into_buffers(&self, mut prefix_buf: Vec, mut ancestors: Vec<(TaggedNodeRef<'a, V, A>, IterToken, usize)>) -> Self { + debug_assert!(prefix_buf.is_empty() && ancestors.is_empty()); + if self.prefix_buf.capacity() > 0 { + prefix_buf.extend_from_slice(&self.prefix_buf); + ancestors.extend_from_slice(&self.ancestors); + } else if prefix_buf.capacity() > 0 { + // A non-zero capacity marks `prefix_buf` as initialized, so it must begin with + // the origin path. `self` hasn't moved yet, so its path is just the origin path + debug_assert!(self.ancestors.is_empty()); + prefix_buf.extend_from_slice(self.origin_path.as_slice()); + } let root_node = self.root_node.clone(); - // The cloned private root has a new address. Its reference is valid while the clone lives. + // A private root value lives in `root_node`'s box, so the copy must point at its own + // box rather than at `self`'s. The box keeps its address while the copy lives. let root_val = match &root_node { OwnedOrBorrowed::Private(private) => Some(unsafe { &*(&private.1 as *const V) }), _ => self.root_val, @@ -2054,8 +2073,8 @@ pub(crate) mod read_zipper_core { root_node, focus_node: self.focus_node.clone(), focus_iter_token: NODE_ITER_INVALID, - prefix_buf: self.prefix_buf.clone(), - ancestors: self.ancestors.clone(), + prefix_buf, + ancestors, alloc: self.alloc.clone(), } } @@ -3021,28 +3040,8 @@ pub(crate) mod read_zipper_core { /// Same as [Clone::clone], but writes the path and ancestor stack into the /// allocations held by `buffer` pub(crate) fn clone_with_buffer(&self, buffer: ZipperBuffer) -> Self { - let ZipperBuffer { mut prefix_buf, ancestors, phantom: _ } = buffer; - let mut ancestors: Vec<_> = ancestors.into_vec(); - if self.prefix_buf.capacity() > 0 { - prefix_buf.extend_from_slice(&self.prefix_buf); - ancestors.extend_from_slice(&self.ancestors); - } else if prefix_buf.capacity() > 0 { - // `self` hasn't moved yet, so its path is just the origin path - debug_assert!(self.ancestors.is_empty()); - prefix_buf.extend_from_slice(self.origin_path.as_slice()); - } - Self { - origin_path: self.origin_path.clone(), - root_key_start: self.root_key_start, - root_parent_key_start: self.root_parent_key_start, - root_val: self.root_val, - root_node: self.root_node.clone(), - focus_node: self.focus_node.clone(), - focus_iter_token: NODE_ITER_INVALID, - prefix_buf, - ancestors, - alloc: self.alloc.clone(), - } + let ZipperBuffer { prefix_buf, ancestors, phantom: _ } = buffer; + self.clone_into_buffers(prefix_buf, ancestors.into_vec()) } /// Returns the length of the `self.path()`, saving a couple instructions, but is internal because it may panic @@ -3768,6 +3767,33 @@ pub(crate) mod read_zipper_core { assert_eq!(cloned.origin_path(), b"alphabet"); assert_eq!(cloned.z.prefix_buf.as_ptr(), prefix_ptr); } + + /// A zipper from a [ZipperHead] keeps a private copy of its root value. A buffered + /// clone must point at its own copy, so it stays valid after the original is dropped + #[test] + fn read_zipper_clone_with_buffer_private_root() { + let mut map = PathMap::::new(); + map.set_val_at(b"ab", 1); + map.set_val_at(b"abcd", 2); + let zh = map.zipper_head(); + + let (mut cloned, mut buffered) = { + let rz = zh.read_zipper_at_path(b"ab").unwrap(); + let buffer = PathMap::::new().read_zipper().into_zipper_buffer(); + (rz.clone(), rz.clone_with_buffer(buffer)) + }; + // Reuse the freed memory, so a dangling reference reads garbage outside miri too + let junk = Box::new([0xdead_usize; 2]); + assert_eq!(cloned.val(), Some(&1)); + assert_eq!(buffered.val(), Some(&1)); + assert!(buffered.descend_to_check(b"cd")); + assert_eq!(buffered.val(), Some(&2)); + assert_eq!(buffered.ascend(2), 2); + assert_eq!(buffered.val(), Some(&1)); + assert!(cloned.descend_to_check(b"cd")); + assert_eq!(cloned.val(), Some(&2)); + drop(junk); + } } use read_zipper_core::*;