diff --git a/src/lib.rs b/src/lib.rs index d7f74c3..72bea4d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,7 +5,8 @@ strict_provenance_atomic_ptr, box_vec_non_null, maybe_uninit_slice, - debug_closure_helpers + debug_closure_helpers, + slice_ptr_get, ) )] #![cfg_attr(feature = "transposed-option", feature(try_trait_v2))] diff --git a/src/tree.rs b/src/tree.rs index e31dee5..0c3c78f 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -68,20 +68,13 @@ const CAPACITY: usize = 16; #[derive(Debug)] struct LeafNode { - parent: Option>>, + parent: Option>, parent_idx: MaybeUninit, value: Option, - leaf: bool, -} - -#[derive(Debug)] -#[repr(C)] -struct InternalNode { - data: LeafNode, - next: Option>>, len: u16, - keys: [MaybeUninit; CAPACITY], - edges: [MaybeUninit>; CAPACITY], + capacity: u16, + keys: NonNull>, + edges: NonNull>>, } type BoxedNode = NonNull>; @@ -137,9 +130,13 @@ impl LeafNode { unsafe fn init(this: *mut Self) { unsafe { (&raw mut (*this).parent).write(None); - (&raw mut (*this).leaf).write(true); - (&raw mut (*this).value).write(None); // parent_idx may be left uninitialized + (&raw mut (*this).value).write(None); + (&raw mut (*this).len).write(0); + (&raw mut (*this).capacity).write(0); + // keys and edges are dangling pointers + (&raw mut (*this).keys).write(NonNull::dangling()); + (&raw mut (*this).edges).write(NonNull::dangling()); } } @@ -152,44 +149,6 @@ impl LeafNode { } } -impl InternalNode { - unsafe fn init(this: *mut Self) { - unsafe { - LeafNode::init(&raw mut (*this).data); - (&raw mut (*this).data.leaf).write(false); - (&raw mut (*this).next).write(None); - (&raw mut (*this).len).write(0); - // keys and edges may be left uninitialized - } - } - - fn init_with_leaf(this: *mut Self, leaf: Box>) { - let mut uninit: Box>> = unsafe { core::mem::transmute(leaf) }; - unsafe { - core::mem::swap(uninit.assume_init_mut(), &mut (*this).data); - (&raw mut (*this).data.leaf).write(false); - (&raw mut (*this).next).write(None); - (&raw mut (*this).len).write(0); - } - } - - fn new_with_leaf(leaf: Box>) -> Box { - let mut this = Box::::new_uninit(); - unsafe { - Self::init_with_leaf(this.as_mut_ptr(), leaf); - this.assume_init() - } - } - - fn new() -> Box { - let mut this = Box::::new_uninit(); - unsafe { - Self::init(this.as_mut_ptr()); - this.assume_init() - } - } -} - impl NodeRef { fn new_leaf() -> Self { let node = LeafNode::new(); @@ -200,33 +159,15 @@ impl NodeRef { } } -impl NodeRef { - fn new_internal() -> Self { - let node = InternalNode::new(); - unsafe { NodeRef::from_new_internal(node) } - } - - /// # Safety - /// `height` must not be zero. - unsafe fn from_new_internal(internal: Box>) -> Self { - let node = NonNull::from(Box::leak(internal)).cast(); - let mut this = NodeRef { +impl NodeRef { + /// Unpack a node reference that was packed as `NodeRef::parent`. + fn new_internal(node: BoxedNode) -> Self { + let node = NodeRef { node, _marker: PhantomData, }; - this - } -} -impl NodeRef { - /// Unpack a node reference that was packed as `NodeRef::parent`. - fn from_internal(node: NonNull>) -> Self { - let node = NodeRef { - node: node.cast(), - _marker: PhantomData, - }; - - debug_assert!(!node.is_leaf()); + // debug_assert!(!node.is_leaf()); node } @@ -251,11 +192,13 @@ impl NodeRef // We need to use raw pointers to nodes because, if BorrowType is marker::ValMut, // there might be outstanding mutable references to values that we must not invalidate. - let leaf_ptr: *const _ = Self::as_leaf_ptr(&self); + let leaf_ptr = Self::as_leaf_ptr(&self).cast_const(); unsafe { (*leaf_ptr).parent } - .as_ref() .map(|parent| Handle { - node: NodeRef::from_internal(*parent), + node: NodeRef { + node: parent, + _marker: PhantomData, + }, idx: unsafe { usize::from((*leaf_ptr).parent_idx.assume_init()) }, _marker: PhantomData, }) @@ -295,12 +238,9 @@ impl // node pointer is dereferenced, we access the edges array with a // reference (Rust issue #73987) and invalidate any other references // to or inside the array, should any be around. - let parent_ptr = NodeRef::as_internal_ptr(&self.node); let node = unsafe { - (*parent_ptr) - .edges - .get_unchecked(self.idx) - .assume_init_read() + // in this case, no references to the node are created. + self.node.edge_area(self.idx).assume_init_read() }; NodeRef { node, @@ -340,39 +280,14 @@ impl<'a, K, V, NodeType, HandleType> Handle, K, V, NodeT } } -impl NodeRef { - /// Exposes the data of an internal node. - /// - /// Returns a raw ptr to avoid invalidating other references to this node. - fn as_internal_ptr(this: &Self) -> *mut InternalNode { - // SAFETY: the static node type is `Internal`. - this.node.as_ptr() as *mut InternalNode - } - - fn as_internal_non_null(this: &Self) -> NonNull> { - // SAFETY: the static node type is `Internal`. - unsafe { NonNull::new_unchecked(Self::as_internal_ptr(this)) } - } -} - -impl<'a, K, V> NodeRef, K, V, marker::Internal> { - /// Borrows exclusive access to the data of an internal node. - fn as_internal_mut(&mut self) -> &mut InternalNode { - let ptr = Self::as_internal_ptr(self); - unsafe { &mut *ptr } - } - +impl<'a, K, V, NodeType> NodeRef, K, V, NodeType> { pub(super) fn len_mut(&mut self) -> &mut u16 { // SAFETY: we have exclusive access to the entire node. - unsafe { &mut (*Self::as_internal_ptr(self)).len } + unsafe { &mut (*Self::as_leaf_ptr(self)).len } } - - pub(super) fn next_mut(&mut self) -> Option, K, V, marker::Internal>> { - let internal = self.as_internal_mut(); - internal.next.map(|next| NodeRef { - node: next.cast(), - _marker: PhantomData, - }) + pub(super) fn capacity_mut(&mut self) -> &mut u16 { + // SAFETY: we have exclusive access to the entire node. + unsafe { &mut (*Self::as_leaf_ptr(self)).len } } } @@ -447,7 +362,45 @@ impl<'a, K, V, Type> NodeRef, K, V, Type> { } } -impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Internal> { +impl<'a, K: 'a, V: 'a, BorrowType> NodeRef { + /// Borrows shared access to an element of the key storage area. + /// + /// # Safety + /// `index` is in bounds of 0..CAPACITY + unsafe fn key_area(&self, index: I) -> &Output + where + I: core::slice::SliceIndex<[MaybeUninit], Output = Output>, + { + // SAFETY: the caller will not be able to call further methods on self + // until the key slice reference is dropped, as we have unique access + // for the lifetime of the borrow. + unsafe { + let keys = (*self.node.as_ptr()).keys; + let capacity = (*self.node.as_ptr()).capacity as usize; + core::slice::from_raw_parts(keys.as_ptr(), capacity).get_unchecked(index) + } + } + + /// Borrows shared access to an element or slice of the node's value storage area. + /// + /// # Safety + /// `index` is in bounds of 0..node.capacity + unsafe fn edge_area(&self, index: I) -> &Output + where + I: core::slice::SliceIndex<[MaybeUninit>], Output = Output>, + { + // SAFETY: the caller will not be able to call further methods on self + // until the value slice reference is dropped, as we have unique access + // for the lifetime of the borrow. + unsafe { + let edges = (*self.node.as_ptr()).edges; + let capacity = (*self.node.as_ptr()).capacity as usize; + core::slice::from_raw_parts(edges.as_ptr(), capacity).get_unchecked(index) + } + } +} + +impl<'a, K: 'a, V: 'a, NodeType> NodeRef, K, V, NodeType> { /// Borrows exclusive access to an element of the key storage area. /// /// # Safety @@ -460,9 +413,8 @@ impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Internal> { // until the key slice reference is dropped, as we have unique access // for the lifetime of the borrow. unsafe { - self.as_internal_mut() - .keys - .as_mut_slice() + let leaf = self.as_leaf_mut(); + core::slice::from_raw_parts_mut(leaf.keys.as_ptr(), leaf.capacity as usize) .get_unchecked_mut(index) } } @@ -470,7 +422,7 @@ impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Internal> { /// Borrows exclusive access to an element or slice of the node's value storage area. /// /// # Safety - /// `index` is in bounds of 0..CAPACITY + /// `index` is in bounds of 0..node.capacity unsafe fn edge_area_mut(&mut self, index: I) -> &mut Output where I: core::slice::SliceIndex<[MaybeUninit>], Output = Output>, @@ -479,9 +431,8 @@ impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Internal> { // until the value slice reference is dropped, as we have unique access // for the lifetime of the borrow. unsafe { - self.as_internal_mut() - .edges - .as_mut_slice() + let leaf = self.as_leaf_mut(); + core::slice::from_raw_parts_mut(leaf.edges.as_ptr(), leaf.capacity as usize) .get_unchecked_mut(index) } } @@ -520,7 +471,7 @@ impl Handle NodeRef { +impl NodeRef { /// Finds the length of the node. This is the number of keys or values. /// The number of edges is `len() + 1`. /// Note that, despite being safe, calling this function can have the side effect @@ -528,13 +479,17 @@ impl NodeRef { pub(super) fn len(&self) -> usize { // Crucially, we only access the `len` field here. If BorrowType is marker::ValMut, // there might be outstanding mutable references to values that we must not invalidate. - unsafe { usize::from((*Self::as_internal_ptr(self)).len) } + unsafe { usize::from((*Self::as_leaf_ptr(self)).len) } + } + + pub(super) fn capacity(&self) -> usize { + unsafe { usize::from((*Self::as_leaf_ptr(self)).len) } } } impl NodeRef { pub(super) fn is_leaf(&self) -> bool { - unsafe { (*Self::as_leaf_ptr(self)).leaf } + unsafe { (*Self::as_leaf_ptr(self)).len == 0 } } /// Temporarily takes out another, immutable reference to the same node. @@ -555,6 +510,11 @@ impl NodeRef { this.node.as_ptr() } + fn as_leaf(&self) -> &LeafNode { + // SAFETY: the static node type is `Leaf`. + unsafe { &*Self::as_leaf_ptr(self) } + } + fn as_leaf_non_null(this: &Self) -> NonNull> { // SAFETY: the static node type is `Leaf`. unsafe { NonNull::new_unchecked(Self::as_leaf_ptr(this)) } @@ -562,38 +522,13 @@ impl NodeRef { } impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Internal> { - fn into_internal(self) -> &'a InternalNode { - // SAFETY: the static node type is `Internal`. - unsafe { &*Self::as_internal_ptr(&self) } - } - /// Borrows a view into the keys stored in the node. pub(super) fn keys(&self) -> &[K] { - let internal = self.into_internal(); - unsafe { - internal - .keys - .get_unchecked(..usize::from(internal.len)) - .assume_init_ref() - } + unsafe { self.key_area(..self.len()).assume_init_ref() } } pub(super) fn edges(&self) -> &[BoxedNode] { - let internal = self.into_internal(); - unsafe { - internal - .edges - .get_unchecked(..usize::from(internal.len + 1)) - .assume_init_ref() - } - } - - pub(super) fn next(&self) -> Option, K, V, marker::Internal>> { - let internal = self.into_internal(); - internal.next.map(|next| NodeRef { - node: next.cast(), - _marker: PhantomData, - }) + unsafe { self.edge_area(..self.len()).assume_init_ref() } } } @@ -746,43 +681,73 @@ impl NodeRef { } } -impl<'a, K, V> NodeRef, K, V, marker::Leaf> { - /// node must have parent/ be a leaf node. - pub(super) unsafe fn make_internal_node( - self, - ) -> NodeRef, K, V, marker::Internal> { - let leaf = self.node; - - let mut parent = self.ascend().ok().unwrap(); - - let mut internal = unsafe { - NodeRef::from_new_internal(InternalNode::new_with_leaf(Box::from_non_null(leaf))) +impl<'a, K, V, Type> NodeRef, K, V, Type> { + pub(super) fn grow_node(&mut self) { + // grow the node + let Some(new_capacity) = self.capacity_mut().checked_mul(2) else { + panic!("Node capacity overflow"); }; - parent.node.as_internal_mut().edges[parent.idx].write(NodeRef::as_leaf_non_null(&internal)); + // Ensure the new capacity is at least 16. + // capacity starts off at 0, so the first time it will grow to 16. + let new_capacity = new_capacity.max(16) as usize; - unsafe { internal.borrow_mut().dormant().awaken() } + // Allocate new keys and edges. + let new_keys = Box::into_non_null(Box::new_uninit_slice(new_capacity)); + let new_edges = Box::into_non_null(Box::new_uninit_slice(new_capacity)); + + let leaf = self.as_leaf_mut(); + + let old_keys = core::mem::replace(&mut leaf.keys, new_keys.as_non_null_ptr()); + let old_edges = core::mem::replace(&mut leaf.edges, new_edges.as_non_null_ptr()); + + // we don't want to copy and deallocate the old keys & edges if the + // capacity was 0, because then the old pointers are dangling. + if leaf.capacity > 0 { + unsafe { + let len = leaf.len as usize; + + core::ptr::copy_nonoverlapping(old_keys.as_ptr(), new_keys.as_mut_ptr(), len); + core::ptr::copy_nonoverlapping(old_edges.as_ptr(), new_edges.as_mut_ptr(), len); + + // dealloc old keys and edges. + // This doesn't drop because the keys and edges were moved. + _ = Box::from_non_null(old_keys); + _ = Box::from_non_null(old_edges); + } + } + + // SAFETY: new_capacity fits in a u16. + leaf.capacity = new_capacity as u16; } } impl NodeRef { - pub(super) fn reparent<'a>( + pub(super) fn reparent<'a, NodeType>( mut self, - mut parent: Handle, K, V, marker::Internal>, marker::Edge>, + mut parent: Handle, K, V, NodeType>, marker::Edge>, key: K, ) -> NodeRef, K, V, Type> where K: 'a, V: 'a, { - self.borrow_mut().as_leaf_mut().parent = Some(NodeRef::as_internal_non_null(&parent.node)); + self.borrow_mut().as_leaf_mut().parent = Some(NodeRef::as_leaf_non_null(&parent.node)); self.borrow_mut() .as_leaf_mut() .parent_idx .write(parent.idx as u16); - let new_len = parent.node.len() + 1; + let old_capacity = parent.node.capacity(); + let old_len = parent.node.len(); + let new_len = old_len + 1; + if new_len > old_capacity { + parent.node.grow_node(); + } + + // insert new key and child node. + // SAFETY: we just grew the allocations. unsafe { slice_insert(parent.node.key_area_mut(..new_len), parent.idx, key); slice_insert(parent.node.edge_area_mut(..new_len), parent.idx, self.node); @@ -794,19 +759,6 @@ impl NodeRef { } } -impl<'a, K, V, Type> Handle, K, V, marker::Leaf>, Type> { - /// node must have parent/ be a leaf node. - pub(super) unsafe fn make_internal_node( - self, - ) -> Handle, K, V, marker::Internal>, Type> { - Handle { - node: unsafe { self.node.make_internal_node() }, - idx: self.idx, - _marker: PhantomData, - } - } -} - impl Handle, marker::Edge> { /// Converts this edge handle into a value handle. pub(super) fn into_value( @@ -854,50 +806,8 @@ impl<'a, K: 'a, V: 'a, Type> Handle, K, V, Type>, marker return unsafe { Handle::new_value(self.node.as_leaf_or_internal()) }; }; - let idx = self.idx; - let mut internal = match self.force() { - ForceResult::Leaf(leaf) => { - // need to turn the leaf into an internal node. - - // SAFETY: root is allocated as an internal node. - // the only node that doesn't have a parent is the root node, - // all parent nodes are internal nodes. - let internal = unsafe { leaf.make_internal_node() }; - - internal - } - ForceResult::Internal(mut internal) => { - if idx >= CAPACITY { - // create a sibling node - let mut sibling = NodeRef::new_internal(); - let mut sibling = sibling.borrow_mut(); - // let sibling = Box::into_non_null(InternalNode::new()); - // set parent link - unsafe { - core::ptr::copy_nonoverlapping( - &raw mut internal.node.as_leaf_mut().parent, - &raw mut sibling.as_leaf_mut().parent, - 1, - ); - core::ptr::copy_nonoverlapping( - &raw mut internal.node.as_leaf_mut().parent_idx, - &raw mut sibling.as_leaf_mut().parent_idx, - 1, - ); - } - - internal.node.as_internal_mut().next = - Some(NodeRef::as_internal_non_null(&sibling)); - - unsafe { Handle::new_edge(sibling.dormant(), 0).awaken() } - } else { - internal - } - } - }; - let last = unsafe { - let child = NodeRef::new_leaf().reparent(internal, key); + let child = NodeRef::new_leaf().reparent(self, key); Handle::new_edge(child, 0) .insert_recursing(key_seq, val) @@ -1073,6 +983,8 @@ enum HandleOrTree<'a, BorrowType, K, V, NodeType, HandleType> { mod entry { use core::marker::PhantomData; + use crate::tree::LeafNode; + use super::{Handle, NodeRef, Tree, borrow::DormantMutRef, marker}; pub enum Entry<'a, Q: 'a, K: 'a, V: 'a> @@ -1172,7 +1084,11 @@ mod entry { HandleOrTree::Tree(mut tree) => { // SAFETY: there are no nodes in the tree yet let tree = unsafe { tree.reborrow() }; - let root = tree.root.insert(NodeRef::new_internal()); + let root = + tree.root + .insert(NodeRef::new_internal(alloc::boxed::Box::into_non_null( + LeafNode::new(), + ))); unsafe { Handle::new_edge(root.borrow_mut(), 0).insert_recursing(self.key, value) @@ -1200,6 +1116,7 @@ mod subtree { marker, }; + // BorrowType may be one of `Immut`, `Mut`. pub struct Subtree { pub(super) root: NodeRef, } @@ -1578,12 +1495,10 @@ mod tests { let mut leaf = NodeRef::<_, (), Test, _>::new_leaf(); leaf.borrow_mut().as_leaf_mut().value = Some(Test("test")); - let mut root = NodeRef::new_internal(); + let mut root = NodeRef::new_leaf(); - let leaf = leaf.reparent(unsafe { Handle::new_edge(root.borrow_mut(), 0) }, ()); + let mut leaf = leaf.reparent(unsafe { Handle::new_edge(root.borrow_mut(), 0) }, ()); - let mut internal = unsafe { leaf.make_internal_node() }; - - assert_eq!(internal.as_leaf_mut().value, Some(Test("test"))); + assert_eq!(leaf.as_leaf_mut().value, Some(Test("test"))); } }