diff --git a/src/trie_map.rs b/src/trie_map.rs index 37855a2d..5c0b1dfd 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 @@ -260,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..de20ef4e 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(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 } + } + /// 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.rs b/src/zipper.rs index d87d79ef..11933059 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::*; @@ -1939,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, @@ -1953,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(), } } @@ -2889,6 +3009,41 @@ 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 { 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 #[inline(always)] fn path_len(&self) -> usize { @@ -3514,6 +3669,131 @@ 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); + } + + /// 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::*; 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);