]> git.proxmox.com Git - rustc.git/blame - library/alloc/src/collections/btree/node.rs
New upstream version 1.57.0+dfsg1
[rustc.git] / library / alloc / src / collections / btree / node.rs
CommitLineData
9cc50fc6
SL
1// This is an attempt at an implementation following the ideal
2//
3// ```
4// struct BTreeMap<K, V> {
5// height: usize,
6// root: Option<Box<Node<K, V, height>>>
7// }
8//
9// struct Node<K, V, height: usize> {
10// keys: [K; 2 * B - 1],
11// vals: [V; 2 * B - 1],
29967ef6
XL
12// edges: [if height > 0 { Box<Node<K, V, height - 1>> } else { () }; 2 * B],
13// parent: Option<(NonNull<Node<K, V, height + 1>>, u16)>,
9cc50fc6
SL
14// len: u16,
15// }
16// ```
17//
7453a54e 18// Since Rust doesn't actually have dependent types and polymorphic recursion,
9cc50fc6 19// we make do with lots of unsafety.
1a4d82fc 20
a7813a04
XL
21// A major goal of this module is to avoid complexity by treating the tree as a generic (if
22// weirdly shaped) container and avoiding dealing with most of the B-Tree invariants. As such,
23// this module doesn't care whether the entries are sorted, which nodes can be underfull, or
24// even what underfull means. However, we do rely on a few invariants:
25//
26// - Trees must have uniform depth/height. This means that every path down to a leaf from a
27// given node has exactly the same length.
29967ef6
XL
28// - A node of length `n` has `n` keys, `n` values, and `n + 1` edges.
29// This implies that even an empty node has at least one edge.
fc512014
XL
30// For a leaf node, "having an edge" only means we can identify a position in the node,
31// since leaf edges are empty and need no data representation. In an internal node,
32// an edge both identifies a position and contains a pointer to a child node.
a7813a04 33
85aaf69f 34use core::marker::PhantomData;
0bf4aa26 35use core::mem::{self, MaybeUninit};
fc512014
XL
36use core::ptr::{self, NonNull};
37use core::slice::SliceIndex;
9cc50fc6 38
fc512014 39use crate::alloc::{Allocator, Global, Layout};
9fa01778 40use crate::boxed::Box;
9cc50fc6
SL
41
42const B: usize = 6;
43pub const CAPACITY: usize = 2 * B - 1;
29967ef6 44pub const MIN_LEN_AFTER_SPLIT: usize = B - 1;
3dfed10e
XL
45const KV_IDX_CENTER: usize = B - 1;
46const EDGE_IDX_LEFT_OF_CENTER: usize = B - 1;
47const EDGE_IDX_RIGHT_OF_CENTER: usize = B;
9cc50fc6 48
1b1a35ee 49/// The underlying representation of leaf nodes and part of the representation of internal nodes.
9cc50fc6 50struct LeafNode<K, V> {
1b1a35ee
XL
51 /// We want to be covariant in `K` and `V`.
52 parent: Option<NonNull<InternalNode<K, V>>>,
a7813a04
XL
53
54 /// This node's index into the parent node's `edges` array.
55 /// `*node.parent.edges[node.parent_idx]` should be the same thing as `node`.
a1dfa0c6 56 /// This is only guaranteed to be initialized when `parent` is non-null.
0bf4aa26 57 parent_idx: MaybeUninit<u16>,
a7813a04
XL
58
59 /// The number of keys and values this node stores.
9cc50fc6 60 len: u16,
94b46f34
XL
61
62 /// The arrays storing the actual data of the node. Only the first `len` elements of each
63 /// array are initialized and valid.
9fa01778
XL
64 keys: [MaybeUninit<K>; CAPACITY],
65 vals: [MaybeUninit<V>; CAPACITY],
1a4d82fc
JJ
66}
67
9cc50fc6 68impl<K, V> LeafNode<K, V> {
6a06907d
XL
69 /// Initializes a new `LeafNode` in-place.
70 unsafe fn init(this: *mut Self) {
71 // As a general policy, we leave fields uninitialized if they can be, as this should
72 // be both slightly faster and easier to track in Valgrind.
73 unsafe {
74 // parent_idx, keys, and vals are all MaybeUninit
75 ptr::addr_of_mut!((*this).parent).write(None);
76 ptr::addr_of_mut!((*this).len).write(0);
77 }
78 }
79
80 /// Creates a new boxed `LeafNode`.
81 fn new() -> Box<Self> {
82 unsafe {
83 let mut leaf = Box::new_uninit();
84 LeafNode::init(leaf.as_mut_ptr());
85 leaf.assume_init()
b039eaaf 86 }
9cc50fc6 87 }
0731742a 88}
94b46f34 89
a7813a04
XL
90/// The underlying representation of internal nodes. As with `LeafNode`s, these should be hidden
91/// behind `BoxedNode`s to prevent dropping uninitialized keys and values. Any pointer to an
cdc7bbd5 92/// `InternalNode` can be directly cast to a pointer to the underlying `LeafNode` portion of the
a7813a04
XL
93/// node, allowing code to act on leaf and internal nodes generically without having to even check
94/// which of the two a pointer is pointing at. This property is enabled by the use of `repr(C)`.
9cc50fc6 95#[repr(C)]
1b1a35ee 96// gdb_providers.py uses this type name for introspection.
9cc50fc6
SL
97struct InternalNode<K, V> {
98 data: LeafNode<K, V>,
a7813a04
XL
99
100 /// The pointers to the children of this node. `len + 1` of these are considered
5869c6ff
XL
101 /// initialized and valid, except that near the end, while the tree is held
102 /// through borrow type `Dying`, some of these pointers are dangling.
9fa01778 103 edges: [MaybeUninit<BoxedNode<K, V>>; 2 * B],
1a4d82fc
JJ
104}
105
9cc50fc6 106impl<K, V> InternalNode<K, V> {
6a06907d 107 /// Creates a new boxed `InternalNode`.
a7813a04 108 ///
6a06907d
XL
109 /// # Safety
110 /// An invariant of internal nodes is that they have at least one
111 /// initialized and valid edge. This function does not set up
a7813a04 112 /// such an edge.
6a06907d
XL
113 unsafe fn new() -> Box<Self> {
114 unsafe {
115 let mut node = Box::<Self>::new_uninit();
116 // We only need to initialize the data; the edges are MaybeUninit.
117 LeafNode::init(ptr::addr_of_mut!((*node.as_mut_ptr()).data));
118 node.assume_init()
119 }
1a4d82fc
JJ
120 }
121}
122
74b04a01 123/// A managed, non-null pointer to a node. This is either an owned pointer to
ba9703b0
XL
124/// `LeafNode<K, V>` or an owned pointer to `InternalNode<K, V>`.
125///
126/// However, `BoxedNode` contains no information as to which of the two types
74b04a01 127/// of nodes it actually contains, and, partially due to this lack of information,
fc512014
XL
128/// is not a separate type and has no destructor.
129type BoxedNode<K, V> = NonNull<LeafNode<K, V>>;
1a4d82fc 130
9cc50fc6
SL
131// N.B. `NodeRef` is always covariant in `K` and `V`, even when the `BorrowType`
132// is `Mut`. This is technically wrong, but cannot result in any unsafety due to
133// internal use of `NodeRef` because we stay completely generic over `K` and `V`.
134// However, whenever a public type wraps `NodeRef`, make sure that it has the
135// correct variance.
29967ef6 136///
9cc50fc6 137/// A reference to a node.
1a4d82fc 138///
3b2f2976 139/// This type has a number of parameters that controls how it acts:
29967ef6
XL
140/// - `BorrowType`: A dummy type that describes the kind of borrow and carries a lifetime.
141/// - When this is `Immut<'a>`, the `NodeRef` acts roughly like `&'a Node`.
142/// - When this is `ValMut<'a>`, the `NodeRef` acts roughly like `&'a Node`
143/// with respect to keys and tree structure, but also allows many
144/// mutable references to values throughout the tree to coexist.
145/// - When this is `Mut<'a>`, the `NodeRef` acts roughly like `&'a mut Node`,
146/// although insert methods allow a mutable pointer to a value to coexist.
147/// - When this is `Owned`, the `NodeRef` acts roughly like `Box<Node>`,
148/// but does not have a destructor, and must be cleaned up manually.
5869c6ff
XL
149/// - When this is `Dying`, the `NodeRef` still acts roughly like `Box<Node>`,
150/// but has methods to destroy the tree bit by bit, and ordinary methods,
151/// while not marked as unsafe to call, can invoke UB if called incorrectly.
fc512014 152/// Since any `NodeRef` allows navigating through the tree, `BorrowType`
5869c6ff 153/// effectively applies to the entire tree, not just to the node itself.
29967ef6 154/// - `K` and `V`: These are the types of keys and values stored in the nodes.
9cc50fc6
SL
155/// - `Type`: This can be `Leaf`, `Internal`, or `LeafOrInternal`. When this is
156/// `Leaf`, the `NodeRef` points to a leaf node, when this is `Internal` the
157/// `NodeRef` points to an internal node, and when this is `LeafOrInternal` the
158/// `NodeRef` could be pointing to either type of node.
29967ef6
XL
159/// `Type` is named `NodeType` when used outside `NodeRef`.
160///
161/// Both `BorrowType` and `NodeType` restrict what methods we implement, to
162/// exploit static type safety. There are limitations in the way we can apply
163/// such restrictions:
164/// - For each type parameter, we can only define a method either generically
165/// or for one particular type. For example, we cannot define a method like
5869c6ff
XL
166/// `into_kv` generically for all `BorrowType`, or once for all types that
167/// carry a lifetime, because we want it to return `&'a` references.
29967ef6
XL
168/// Therefore, we define it only for the least powerful type `Immut<'a>`.
169/// - We cannot get implicit coercion from say `Mut<'a>` to `Immut<'a>`.
136023e0 170/// Therefore, we have to explicitly call `reborrow` on a more powerful
5869c6ff 171/// `NodeRef` in order to reach a method like `into_kv`.
fc512014
XL
172///
173/// All methods on `NodeRef` that return some kind of reference, either:
174/// - Take `self` by value, and return the lifetime carried by `BorrowType`.
175/// Sometimes, to invoke such a method, we need to call `reborrow_mut`.
176/// - Take `self` by reference, and (implicitly) return that reference's
177/// lifetime, instead of the lifetime carried by `BorrowType`. That way,
178/// the borrow checker guarantees that the `NodeRef` remains borrowed as long
179/// as the returned reference is used.
180/// The methods supporting insert bend this rule by returning a raw pointer,
181/// i.e., a reference without any lifetime.
9cc50fc6 182pub struct NodeRef<BorrowType, K, V, Type> {
fc512014
XL
183 /// The number of levels that the node and the level of leaves are apart, a
184 /// constant of the node that cannot be entirely described by `Type`, and that
185 /// the node itself does not store. We only need to store the height of the root
186 /// node, and derive every other node's height from it.
187 /// Must be zero if `Type` is `Leaf` and non-zero if `Type` is `Internal`.
9cc50fc6 188 height: usize,
29967ef6
XL
189 /// The pointer to the leaf or internal node. The definition of `InternalNode`
190 /// ensures that the pointer is valid either way.
0531ce1d 191 node: NonNull<LeafNode<K, V>>,
dfeec247 192 _marker: PhantomData<(BorrowType, Type)>,
1a4d82fc
JJ
193}
194
cdc7bbd5
XL
195/// The root node of an owned tree.
196///
197/// Note that this does not have a destructor, and must be cleaned up manually.
198pub type Root<K, V> = NodeRef<marker::Owned, K, V, marker::LeafOrInternal>;
199
dfeec247 200impl<'a, K: 'a, V: 'a, Type> Copy for NodeRef<marker::Immut<'a>, K, V, Type> {}
9cc50fc6
SL
201impl<'a, K: 'a, V: 'a, Type> Clone for NodeRef<marker::Immut<'a>, K, V, Type> {
202 fn clone(&self) -> Self {
203 *self
204 }
1a4d82fc
JJ
205}
206
dfeec247 207unsafe impl<BorrowType, K: Sync, V: Sync, Type> Sync for NodeRef<BorrowType, K, V, Type> {}
9cc50fc6 208
dfeec247
XL
209unsafe impl<'a, K: Sync + 'a, V: Sync + 'a, Type> Send for NodeRef<marker::Immut<'a>, K, V, Type> {}
210unsafe impl<'a, K: Send + 'a, V: Send + 'a, Type> Send for NodeRef<marker::Mut<'a>, K, V, Type> {}
1b1a35ee 211unsafe impl<'a, K: Send + 'a, V: Send + 'a, Type> Send for NodeRef<marker::ValMut<'a>, K, V, Type> {}
dfeec247 212unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Owned, K, V, Type> {}
5869c6ff 213unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Dying, K, V, Type> {}
9cc50fc6 214
cdc7bbd5
XL
215impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
216 fn new_leaf() -> Self {
217 Self::from_new_leaf(LeafNode::new())
218 }
219
220 fn from_new_leaf(leaf: Box<LeafNode<K, V>>) -> Self {
221 NodeRef { height: 0, node: NonNull::from(Box::leak(leaf)), _marker: PhantomData }
222 }
223}
224
225impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
226 fn new_internal(child: Root<K, V>) -> Self {
227 let mut new_node = unsafe { InternalNode::new() };
228 new_node.edges[0].write(child.node);
229 unsafe { NodeRef::from_new_internal(new_node, child.height + 1) }
230 }
231
232 /// # Safety
233 /// `height` must not be zero.
234 unsafe fn from_new_internal(internal: Box<InternalNode<K, V>>, height: usize) -> Self {
235 debug_assert!(height > 0);
236 let node = NonNull::from(Box::leak(internal)).cast();
237 let mut this = NodeRef { height, node, _marker: PhantomData };
238 this.borrow_mut().correct_all_childrens_parent_links();
239 this
240 }
241}
242
29967ef6
XL
243impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
244 /// Unpack a node reference that was packed as `NodeRef::parent`.
245 fn from_internal(node: NonNull<InternalNode<K, V>>, height: usize) -> Self {
246 debug_assert!(height > 0);
247 NodeRef { height, node: node.cast(), _marker: PhantomData }
248 }
249}
250
9cc50fc6 251impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
29967ef6 252 /// Exposes the data of an internal node.
1b1a35ee 253 ///
29967ef6
XL
254 /// Returns a raw ptr to avoid invalidating other references to this node.
255 fn as_internal_ptr(this: &Self) -> *mut InternalNode<K, V> {
256 // SAFETY: the static node type is `Internal`.
257 this.node.as_ptr() as *mut InternalNode<K, V>
1a4d82fc 258 }
1a4d82fc
JJ
259}
260
1b1a35ee 261impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
fc512014
XL
262 /// Borrows exclusive access to the data of an internal node.
263 fn as_internal_mut(&mut self) -> &mut InternalNode<K, V> {
264 let ptr = Self::as_internal_ptr(self);
29967ef6 265 unsafe { &mut *ptr }
1a4d82fc 266 }
9cc50fc6
SL
267}
268
9cc50fc6 269impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
29967ef6
XL
270 /// Finds the length of the node. This is the number of keys or values.
271 /// The number of edges is `len() + 1`.
74b04a01
XL
272 /// Note that, despite being safe, calling this function can have the side effect
273 /// of invalidating mutable references that unsafe code has created.
85aaf69f 274 pub fn len(&self) -> usize {
1b1a35ee
XL
275 // Crucially, we only access the `len` field here. If BorrowType is marker::ValMut,
276 // there might be outstanding mutable references to values that we must not invalidate.
29967ef6 277 unsafe { usize::from((*Self::as_leaf_ptr(self)).len) }
1a4d82fc
JJ
278 }
279
fc512014
XL
280 /// Returns the number of levels that the node and leaves are apart. Zero
281 /// height means the node is a leaf itself. If you picture trees with the
282 /// root on top, the number says at which elevation the node appears.
283 /// If you picture trees with leaves on top, the number says how high
284 /// the tree extends above the node.
a7813a04
XL
285 pub fn height(&self) -> usize {
286 self.height
287 }
288
a7813a04 289 /// Temporarily takes out another, immutable reference to the same node.
29967ef6 290 pub fn reborrow(&self) -> NodeRef<marker::Immut<'_>, K, V, Type> {
1b1a35ee 291 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
1a4d82fc
JJ
292 }
293
1b1a35ee 294 /// Exposes the leaf portion of any leaf or internal node.
1b1a35ee 295 ///
29967ef6
XL
296 /// Returns a raw ptr to avoid invalidating other references to this node.
297 fn as_leaf_ptr(this: &Self) -> *mut LeafNode<K, V> {
ba9703b0
XL
298 // The node must be valid for at least the LeafNode portion.
299 // This is not a reference in the NodeRef type because we don't know if
300 // it should be unique or shared.
29967ef6 301 this.node.as_ptr()
94b46f34 302 }
29967ef6 303}
94b46f34 304
5869c6ff 305impl<BorrowType: marker::BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
a7813a04
XL
306 /// Finds the parent of the current node. Returns `Ok(handle)` if the current
307 /// node actually has a parent, where `handle` points to the edge of the parent
308 /// that points to the current node. Returns `Err(self)` if the current node has
309 /// no parent, giving back the original `NodeRef`.
310 ///
fc512014
XL
311 /// The method name assumes you picture trees with the root node on top.
312 ///
a7813a04
XL
313 /// `edge.descend().ascend().unwrap()` and `node.ascend().unwrap().descend()` should
314 /// both, upon success, do nothing.
dfeec247
XL
315 pub fn ascend(
316 self,
317 ) -> Result<Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge>, Self> {
5869c6ff 318 assert!(BorrowType::PERMITS_TRAVERSAL);
1b1a35ee
XL
319 // We need to use raw pointers to nodes because, if BorrowType is marker::ValMut,
320 // there might be outstanding mutable references to values that we must not invalidate.
29967ef6 321 let leaf_ptr: *const _ = Self::as_leaf_ptr(&self);
1b1a35ee
XL
322 unsafe { (*leaf_ptr).parent }
323 .as_ref()
324 .map(|parent| Handle {
29967ef6 325 node: NodeRef::from_internal(*parent, self.height + 1),
1b1a35ee 326 idx: unsafe { usize::from((*leaf_ptr).parent_idx.assume_init()) },
dfeec247 327 _marker: PhantomData,
9cc50fc6 328 })
1b1a35ee 329 .ok_or(self)
1a4d82fc 330 }
1a4d82fc 331
9cc50fc6 332 pub fn first_edge(self) -> Handle<Self, marker::Edge> {
74b04a01 333 unsafe { Handle::new_edge(self, 0) }
1a4d82fc
JJ
334 }
335
9cc50fc6
SL
336 pub fn last_edge(self) -> Handle<Self, marker::Edge> {
337 let len = self.len();
74b04a01 338 unsafe { Handle::new_edge(self, len) }
1a4d82fc 339 }
3157f602
XL
340
341 /// Note that `self` must be nonempty.
342 pub fn first_kv(self) -> Handle<Self, marker::KV> {
74b04a01
XL
343 let len = self.len();
344 assert!(len > 0);
345 unsafe { Handle::new_kv(self, 0) }
3157f602
XL
346 }
347
348 /// Note that `self` must be nonempty.
349 pub fn last_kv(self) -> Handle<Self, marker::KV> {
350 let len = self.len();
74b04a01
XL
351 assert!(len > 0);
352 unsafe { Handle::new_kv(self, len - 1) }
1a4d82fc
JJ
353 }
354}
355
cdc7bbd5
XL
356impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
357 /// Could be a public implementation of PartialEq, but only used in this module.
358 fn eq(&self, other: &Self) -> bool {
359 let Self { node, height, _marker } = self;
360 if node.eq(&other.node) {
361 debug_assert_eq!(*height, other.height);
362 true
363 } else {
364 false
365 }
366 }
367}
368
1b1a35ee 369impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, K, V, Type> {
29967ef6 370 /// Exposes the leaf portion of any leaf or internal node in an immutable tree.
fc512014
XL
371 fn into_leaf(self) -> &'a LeafNode<K, V> {
372 let ptr = Self::as_leaf_ptr(&self);
29967ef6
XL
373 // SAFETY: there can be no mutable references into this tree borrowed as `Immut`.
374 unsafe { &*ptr }
1b1a35ee 375 }
5869c6ff
XL
376
377 /// Borrows a view into the keys stored in the node.
378 pub fn keys(&self) -> &[K] {
379 let leaf = self.into_leaf();
380 unsafe {
381 MaybeUninit::slice_assume_init_ref(leaf.keys.get_unchecked(..usize::from(leaf.len)))
382 }
383 }
1b1a35ee
XL
384}
385
5869c6ff 386impl<K, V> NodeRef<marker::Dying, K, V, marker::LeafOrInternal> {
a7813a04 387 /// Similar to `ascend`, gets a reference to a node's parent node, but also
5869c6ff 388 /// deallocates the current node in the process. This is unsafe because the
a7813a04 389 /// current node will still be accessible despite being deallocated.
dfeec247
XL
390 pub unsafe fn deallocate_and_ascend(
391 self,
5869c6ff 392 ) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::Internal>, marker::Edge>> {
74b04a01 393 let height = self.height;
83c7162d 394 let node = self.node;
9cc50fc6 395 let ret = self.ascend().ok();
f035d41b 396 unsafe {
fc512014 397 Global.deallocate(
f035d41b
XL
398 node.cast(),
399 if height > 0 {
400 Layout::new::<InternalNode<K, V>>()
401 } else {
402 Layout::new::<LeafNode<K, V>>()
403 },
404 );
405 }
9cc50fc6 406 ret
1a4d82fc
JJ
407 }
408}
409
6a06907d 410impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
cdc7bbd5 411 /// Temporarily takes out another mutable reference to the same node. Beware, as
94222f64 412 /// this method is very dangerous, doubly so since it might not immediately appear
a7813a04
XL
413 /// dangerous.
414 ///
1b1a35ee
XL
415 /// Because mutable pointers can roam anywhere around the tree, the returned
416 /// pointer can easily be used to make the original pointer dangling, out of
417 /// bounds, or invalid under stacked borrow rules.
418 // FIXME(@gereeter) consider adding yet another type parameter to `NodeRef`
419 // that restricts the use of navigation methods on reborrowed pointers,
420 // preventing this unsafety.
9fa01778 421 unsafe fn reborrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, Type> {
1b1a35ee 422 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
1a4d82fc
JJ
423 }
424
17df50a5 425 /// Borrows exclusive access to the leaf portion of a leaf or internal node.
fc512014
XL
426 fn as_leaf_mut(&mut self) -> &mut LeafNode<K, V> {
427 let ptr = Self::as_leaf_ptr(self);
428 // SAFETY: we have exclusive access to the entire node.
429 unsafe { &mut *ptr }
430 }
431
17df50a5 432 /// Offers exclusive access to the leaf portion of a leaf or internal node.
fc512014
XL
433 fn into_leaf_mut(mut self) -> &'a mut LeafNode<K, V> {
434 let ptr = Self::as_leaf_ptr(&mut self);
29967ef6
XL
435 // SAFETY: we have exclusive access to the entire node.
436 unsafe { &mut *ptr }
1b1a35ee 437 }
29967ef6 438}
1b1a35ee 439
17df50a5
XL
440impl<K, V, Type> NodeRef<marker::Dying, K, V, Type> {
441 /// Borrows exclusive access to the leaf portion of a dying leaf or internal node.
442 fn as_leaf_dying(&mut self) -> &mut LeafNode<K, V> {
443 let ptr = Self::as_leaf_ptr(self);
444 // SAFETY: we have exclusive access to the entire node.
445 unsafe { &mut *ptr }
446 }
447}
448
29967ef6 449impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
fc512014 450 /// Borrows exclusive access to an element of the key storage area.
1b1a35ee
XL
451 ///
452 /// # Safety
fc512014 453 /// `index` is in bounds of 0..CAPACITY
5869c6ff 454 unsafe fn key_area_mut<I, Output: ?Sized>(&mut self, index: I) -> &mut Output
fc512014
XL
455 where
456 I: SliceIndex<[MaybeUninit<K>], Output = Output>,
457 {
458 // SAFETY: the caller will not be able to call further methods on self
459 // until the key slice reference is dropped, as we have unique access
460 // for the lifetime of the borrow.
461 unsafe { self.as_leaf_mut().keys.as_mut_slice().get_unchecked_mut(index) }
1a4d82fc
JJ
462 }
463
fc512014 464 /// Borrows exclusive access to an element or slice of the node's value storage area.
1b1a35ee
XL
465 ///
466 /// # Safety
fc512014 467 /// `index` is in bounds of 0..CAPACITY
5869c6ff 468 unsafe fn val_area_mut<I, Output: ?Sized>(&mut self, index: I) -> &mut Output
fc512014
XL
469 where
470 I: SliceIndex<[MaybeUninit<V>], Output = Output>,
471 {
472 // SAFETY: the caller will not be able to call further methods on self
473 // until the value slice reference is dropped, as we have unique access
474 // for the lifetime of the borrow.
475 unsafe { self.as_leaf_mut().vals.as_mut_slice().get_unchecked_mut(index) }
1b1a35ee
XL
476 }
477}
478
29967ef6 479impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
fc512014 480 /// Borrows exclusive access to an element or slice of the node's storage area for edge contents.
29967ef6
XL
481 ///
482 /// # Safety
fc512014 483 /// `index` is in bounds of 0..CAPACITY + 1
5869c6ff 484 unsafe fn edge_area_mut<I, Output: ?Sized>(&mut self, index: I) -> &mut Output
fc512014
XL
485 where
486 I: SliceIndex<[MaybeUninit<BoxedNode<K, V>>], Output = Output>,
487 {
488 // SAFETY: the caller will not be able to call further methods on self
489 // until the edge slice reference is dropped, as we have unique access
490 // for the lifetime of the borrow.
491 unsafe { self.as_internal_mut().edges.as_mut_slice().get_unchecked_mut(index) }
1a4d82fc
JJ
492 }
493}
494
1b1a35ee
XL
495impl<'a, K, V, Type> NodeRef<marker::ValMut<'a>, K, V, Type> {
496 /// # Safety
29967ef6 497 /// - The node has more than `idx` initialized elements.
29967ef6 498 unsafe fn into_key_val_mut_at(mut self, idx: usize) -> (&'a K, &'a mut V) {
1b1a35ee
XL
499 // We only create a reference to the one element we are interested in,
500 // to avoid aliasing with outstanding references to other elements,
501 // in particular, those returned to the caller in earlier iterations.
29967ef6 502 let leaf = Self::as_leaf_ptr(&mut self);
5869c6ff
XL
503 let keys = unsafe { ptr::addr_of!((*leaf).keys) };
504 let vals = unsafe { ptr::addr_of_mut!((*leaf).vals) };
1b1a35ee 505 // We must coerce to unsized array pointers because of Rust issue #74679.
29967ef6
XL
506 let keys: *const [_] = keys;
507 let vals: *mut [_] = vals;
1b1a35ee
XL
508 let key = unsafe { (&*keys.get_unchecked(idx)).assume_init_ref() };
509 let val = unsafe { (&mut *vals.get_unchecked_mut(idx)).assume_init_mut() };
510 (key, val)
94b46f34 511 }
1a4d82fc
JJ
512}
513
29967ef6 514impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
fc512014
XL
515 /// Borrows exclusive access to the length of the node.
516 pub fn len_mut(&mut self) -> &mut u16 {
517 &mut self.as_leaf_mut().len
29967ef6
XL
518 }
519}
520
cdc7bbd5
XL
521impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
522 /// # Safety
523 /// Every item returned by `range` is a valid edge index for the node.
524 unsafe fn correct_childrens_parent_links<R: Iterator<Item = usize>>(&mut self, range: R) {
525 for i in range {
526 debug_assert!(i <= self.len());
527 unsafe { Handle::new_edge(self.reborrow_mut(), i) }.correct_parent_link();
528 }
529 }
530
531 fn correct_all_childrens_parent_links(&mut self) {
532 let len = self.len();
533 unsafe { self.correct_childrens_parent_links(0..=len) };
534 }
535}
536
29967ef6 537impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
fc512014 538 /// Sets the node's link to its parent edge,
29967ef6
XL
539 /// without invalidating other references to the node.
540 fn set_parent_link(&mut self, parent: NonNull<InternalNode<K, V>>, parent_idx: usize) {
541 let leaf = Self::as_leaf_ptr(self);
542 unsafe { (*leaf).parent = Some(parent) };
543 unsafe { (*leaf).parent_idx.write(parent_idx as u16) };
544 }
fc512014 545}
29967ef6 546
fc512014
XL
547impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
548 /// Clears the root's link to its parent edge.
29967ef6 549 fn clear_parent_link(&mut self) {
fc512014
XL
550 let mut root_node = self.borrow_mut();
551 let leaf = root_node.as_leaf_mut();
29967ef6
XL
552 leaf.parent = None;
553 }
554}
555
cdc7bbd5
XL
556impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
557 /// Returns a new owned tree, with its own root node that is initially empty.
558 pub fn new() -> Self {
559 NodeRef::new_leaf().forget_type()
560 }
561
562 /// Adds a new internal node with a single edge pointing to the previous root node,
563 /// make that new node the root node, and return it. This increases the height by 1
564 /// and is the opposite of `pop_internal_level`.
565 pub fn push_internal_level(&mut self) -> NodeRef<marker::Mut<'_>, K, V, marker::Internal> {
566 super::mem::take_mut(self, |old_root| NodeRef::new_internal(old_root).forget_type());
567
568 // `self.borrow_mut()`, except that we just forgot we're internal now:
569 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
570 }
571
572 /// Removes the internal root node, using its first child as the new root node.
573 /// As it is intended only to be called when the root node has only one child,
574 /// no cleanup is done on any of the keys, values and other children.
575 /// This decreases the height by 1 and is the opposite of `push_internal_level`.
576 ///
c295e0f8 577 /// Requires exclusive access to the `NodeRef` object but not to the root node;
cdc7bbd5
XL
578 /// it will not invalidate other handles or references to the root node.
579 ///
580 /// Panics if there is no internal level, i.e., if the root node is a leaf.
581 pub fn pop_internal_level(&mut self) {
582 assert!(self.height > 0);
583
584 let top = self.node;
585
586 // SAFETY: we asserted to be internal.
587 let internal_self = unsafe { self.borrow_mut().cast_to_internal_unchecked() };
588 // SAFETY: we borrowed `self` exclusively and its borrow type is exclusive.
589 let internal_node = unsafe { &mut *NodeRef::as_internal_ptr(&internal_self) };
590 // SAFETY: the first edge is always initialized.
591 self.node = unsafe { internal_node.edges[0].assume_init_read() };
592 self.height -= 1;
593 self.clear_parent_link();
594
595 unsafe {
596 Global.deallocate(top.cast(), Layout::new::<InternalNode<K, V>>());
597 }
598 }
599}
600
601impl<K, V, Type> NodeRef<marker::Owned, K, V, Type> {
602 /// Mutably borrows the owned root node. Unlike `reborrow_mut`, this is safe
603 /// because the return value cannot be used to destroy the root, and there
604 /// cannot be other references to the tree.
605 pub fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, Type> {
606 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
607 }
608
609 /// Slightly mutably borrows the owned root node.
610 pub fn borrow_valmut(&mut self) -> NodeRef<marker::ValMut<'_>, K, V, Type> {
611 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
612 }
613
614 /// Irreversibly transitions to a reference that permits traversal and offers
615 /// destructive methods and little else.
616 pub fn into_dying(self) -> NodeRef<marker::Dying, K, V, Type> {
617 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
618 }
619}
620
1b1a35ee 621impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> {
fc512014 622 /// Adds a key-value pair to the end of the node.
9cc50fc6 623 pub fn push(&mut self, key: K, val: V) {
fc512014 624 let len = self.len_mut();
1b1a35ee
XL
625 let idx = usize::from(*len);
626 assert!(idx < CAPACITY);
627 *len += 1;
1a4d82fc 628 unsafe {
5869c6ff
XL
629 self.key_area_mut(idx).write(key);
630 self.val_area_mut(idx).write(val);
0731742a 631 }
1a4d82fc
JJ
632 }
633}
634
1b1a35ee 635impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
fc512014 636 /// Adds a key-value pair, and an edge to go to the right of that pair,
1b1a35ee 637 /// to the end of the node.
9cc50fc6 638 pub fn push(&mut self, key: K, val: V, edge: Root<K, V>) {
74b04a01 639 assert!(edge.height == self.height - 1);
1a4d82fc 640
fc512014 641 let len = self.len_mut();
1b1a35ee
XL
642 let idx = usize::from(*len);
643 assert!(idx < CAPACITY);
644 *len += 1;
9cc50fc6 645 unsafe {
5869c6ff
XL
646 self.key_area_mut(idx).write(key);
647 self.val_area_mut(idx).write(val);
648 self.edge_area_mut(idx + 1).write(edge.node);
9cc50fc6 649 Handle::new_edge(self.reborrow_mut(), idx + 1).correct_parent_link();
1a4d82fc
JJ
650 }
651 }
1a4d82fc
JJ
652}
653
cdc7bbd5
XL
654impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Leaf> {
655 /// Removes any static information asserting that this node is a `Leaf` node.
656 pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
657 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
658 }
659}
660
661impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
662 /// Removes any static information asserting that this node is an `Internal` node.
663 pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
664 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
665 }
666}
667
9cc50fc6 668impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
a7813a04 669 /// Checks whether a node is an `Internal` node or a `Leaf` node.
dfeec247
XL
670 pub fn force(
671 self,
672 ) -> ForceResult<
9cc50fc6 673 NodeRef<BorrowType, K, V, marker::Leaf>,
dfeec247 674 NodeRef<BorrowType, K, V, marker::Internal>,
9cc50fc6
SL
675 > {
676 if self.height == 0 {
677 ForceResult::Leaf(NodeRef {
678 height: self.height,
679 node: self.node,
dfeec247 680 _marker: PhantomData,
9cc50fc6
SL
681 })
682 } else {
683 ForceResult::Internal(NodeRef {
684 height: self.height,
685 node: self.node,
dfeec247 686 _marker: PhantomData,
9cc50fc6 687 })
1a4d82fc
JJ
688 }
689 }
9cc50fc6 690}
1a4d82fc 691
cdc7bbd5
XL
692impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
693 /// Unsafely asserts to the compiler the static information that this node is a `Leaf`.
694 unsafe fn cast_to_leaf_unchecked(self) -> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> {
695 debug_assert!(self.height == 0);
696 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
697 }
698
699 /// Unsafely asserts to the compiler the static information that this node is an `Internal`.
700 unsafe fn cast_to_internal_unchecked(self) -> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
701 debug_assert!(self.height > 0);
702 NodeRef { height: self.height, node: self.node, _marker: PhantomData }
703 }
704}
705
fc512014
XL
706/// A reference to a specific key-value pair or edge within a node. The `Node` parameter
707/// must be a `NodeRef`, while the `Type` can either be `KV` (signifying a handle on a key-value
a7813a04
XL
708/// pair) or `Edge` (signifying a handle on an edge).
709///
710/// Note that even `Leaf` nodes can have `Edge` handles. Instead of representing a pointer to
fc512014 711/// a child node, these represent the spaces where child pointers would go between the key-value
a7813a04
XL
712/// pairs. For example, in a node with length 2, there would be 3 possible edge locations - one
713/// to the left of the node, one between the two pairs, and one at the right of the node.
9cc50fc6
SL
714pub struct Handle<Node, Type> {
715 node: Node,
716 idx: usize,
dfeec247 717 _marker: PhantomData<Type>,
9cc50fc6 718}
1a4d82fc 719
dfeec247 720impl<Node: Copy, Type> Copy for Handle<Node, Type> {}
a7813a04
XL
721// We don't need the full generality of `#[derive(Clone)]`, as the only time `Node` will be
722// `Clone`able is when it is an immutable reference and therefore `Copy`.
9cc50fc6
SL
723impl<Node: Copy, Type> Clone for Handle<Node, Type> {
724 fn clone(&self) -> Self {
725 *self
1a4d82fc 726 }
9cc50fc6 727}
1a4d82fc 728
9cc50fc6 729impl<Node, Type> Handle<Node, Type> {
fc512014 730 /// Retrieves the node that contains the edge or key-value pair this handle points to.
9cc50fc6
SL
731 pub fn into_node(self) -> Node {
732 self.node
1a4d82fc 733 }
ba9703b0
XL
734
735 /// Returns the position of this handle in the node.
736 pub fn idx(&self) -> usize {
737 self.idx
738 }
9cc50fc6 739}
1a4d82fc 740
9cc50fc6 741impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV> {
fc512014 742 /// Creates a new handle to a key-value pair in `node`.
74b04a01
XL
743 /// Unsafe because the caller must ensure that `idx < node.len()`.
744 pub unsafe fn new_kv(node: NodeRef<BorrowType, K, V, NodeType>, idx: usize) -> Self {
9cc50fc6 745 debug_assert!(idx < node.len());
1a4d82fc 746
dfeec247 747 Handle { node, idx, _marker: PhantomData }
1a4d82fc 748 }
1a4d82fc 749
9cc50fc6 750 pub fn left_edge(self) -> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
74b04a01 751 unsafe { Handle::new_edge(self.node, self.idx) }
1a4d82fc
JJ
752 }
753
9cc50fc6 754 pub fn right_edge(self) -> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
74b04a01 755 unsafe { Handle::new_edge(self.node, self.idx + 1) }
1a4d82fc 756 }
9cc50fc6 757}
1a4d82fc 758
9cc50fc6 759impl<BorrowType, K, V, NodeType, HandleType> PartialEq
dfeec247
XL
760 for Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType>
761{
9cc50fc6 762 fn eq(&self, other: &Self) -> bool {
fc512014 763 let Self { node, idx, _marker } = self;
29967ef6 764 node.eq(&other.node) && *idx == other.idx
9cc50fc6
SL
765 }
766}
1a4d82fc 767
9cc50fc6 768impl<BorrowType, K, V, NodeType, HandleType>
dfeec247
XL
769 Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType>
770{
cdc7bbd5 771 /// Temporarily takes out another immutable handle on the same location.
dfeec247 772 pub fn reborrow(&self) -> Handle<NodeRef<marker::Immut<'_>, K, V, NodeType>, HandleType> {
9cc50fc6 773 // We can't use Handle::new_kv or Handle::new_edge because we don't know our type
dfeec247 774 Handle { node: self.node.reborrow(), idx: self.idx, _marker: PhantomData }
1a4d82fc 775 }
9cc50fc6 776}
1a4d82fc 777
6a06907d 778impl<'a, K, V, NodeType, HandleType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, HandleType> {
cdc7bbd5 779 /// Temporarily takes out another mutable handle on the same location. Beware, as
94222f64 780 /// this method is very dangerous, doubly so since it might not immediately appear
a7813a04
XL
781 /// dangerous.
782 ///
1b1a35ee 783 /// For details, see `NodeRef::reborrow_mut`.
dfeec247
XL
784 pub unsafe fn reborrow_mut(
785 &mut self,
786 ) -> Handle<NodeRef<marker::Mut<'_>, K, V, NodeType>, HandleType> {
9cc50fc6 787 // We can't use Handle::new_kv or Handle::new_edge because we don't know our type
f035d41b 788 Handle { node: unsafe { self.node.reborrow_mut() }, idx: self.idx, _marker: PhantomData }
1a4d82fc 789 }
9cc50fc6 790}
1a4d82fc 791
dfeec247 792impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
74b04a01
XL
793 /// Creates a new handle to an edge in `node`.
794 /// Unsafe because the caller must ensure that `idx <= node.len()`.
795 pub unsafe fn new_edge(node: NodeRef<BorrowType, K, V, NodeType>, idx: usize) -> Self {
9cc50fc6 796 debug_assert!(idx <= node.len());
1a4d82fc 797
dfeec247 798 Handle { node, idx, _marker: PhantomData }
9cc50fc6 799 }
1a4d82fc 800
dfeec247 801 pub fn left_kv(self) -> Result<Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV>, Self> {
74b04a01
XL
802 if self.idx > 0 {
803 Ok(unsafe { Handle::new_kv(self.node, self.idx - 1) })
804 } else {
805 Err(self)
806 }
1a4d82fc
JJ
807 }
808
dfeec247 809 pub fn right_kv(self) -> Result<Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV>, Self> {
74b04a01
XL
810 if self.idx < self.node.len() {
811 Ok(unsafe { Handle::new_kv(self.node, self.idx) })
812 } else {
813 Err(self)
814 }
1a4d82fc
JJ
815 }
816}
817
fc512014
XL
818pub enum LeftOrRight<T> {
819 Left(T),
820 Right(T),
3dfed10e
XL
821}
822
823/// Given an edge index where we want to insert into a node filled to capacity,
824/// computes a sensible KV index of a split point and where to perform the insertion.
825/// The goal of the split point is for its key and value to end up in a parent node;
826/// the keys, values and edges to the left of the split point become the left child;
827/// the keys, values and edges to the right of the split point become the right child.
fc512014 828fn splitpoint(edge_idx: usize) -> (usize, LeftOrRight<usize>) {
3dfed10e
XL
829 debug_assert!(edge_idx <= CAPACITY);
830 // Rust issue #74834 tries to explain these symmetric rules.
831 match edge_idx {
fc512014
XL
832 0..EDGE_IDX_LEFT_OF_CENTER => (KV_IDX_CENTER - 1, LeftOrRight::Left(edge_idx)),
833 EDGE_IDX_LEFT_OF_CENTER => (KV_IDX_CENTER, LeftOrRight::Left(edge_idx)),
834 EDGE_IDX_RIGHT_OF_CENTER => (KV_IDX_CENTER, LeftOrRight::Right(0)),
835 _ => (KV_IDX_CENTER + 1, LeftOrRight::Right(edge_idx - (KV_IDX_CENTER + 1 + 1))),
3dfed10e
XL
836 }
837}
838
29967ef6 839impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
fc512014 840 /// Inserts a new key-value pair between the key-value pairs to the right and left of
a7813a04
XL
841 /// this edge. This method assumes that there is enough space in the node for the new
842 /// pair to fit.
29967ef6
XL
843 ///
844 /// The returned pointer points to the inserted value.
845 fn insert_fit(&mut self, key: K, val: V) -> *mut V {
9cc50fc6 846 debug_assert!(self.node.len() < CAPACITY);
fc512014 847 let new_len = self.node.len() + 1;
1a4d82fc
JJ
848
849 unsafe {
5869c6ff
XL
850 slice_insert(self.node.key_area_mut(..new_len), self.idx, key);
851 slice_insert(self.node.val_area_mut(..new_len), self.idx, val);
fc512014 852 *self.node.len_mut() = new_len as u16;
1a4d82fc 853
5869c6ff 854 self.node.val_area_mut(self.idx).assume_init_mut()
1a4d82fc
JJ
855 }
856 }
3dfed10e
XL
857}
858
1b1a35ee 859impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
fc512014 860 /// Inserts a new key-value pair between the key-value pairs to the right and left of
a7813a04
XL
861 /// this edge. This method splits the node if there isn't enough room.
862 ///
863 /// The returned pointer points to the inserted value.
3dfed10e 864 fn insert(mut self, key: K, val: V) -> (InsertResult<'a, K, V, marker::Leaf>, *mut V) {
9cc50fc6 865 if self.node.len() < CAPACITY {
1b1a35ee 866 let val_ptr = self.insert_fit(key, val);
74b04a01 867 let kv = unsafe { Handle::new_kv(self.node, self.idx) };
1b1a35ee 868 (InsertResult::Fit(kv), val_ptr)
9cc50fc6 869 } else {
3dfed10e
XL
870 let (middle_kv_idx, insertion) = splitpoint(self.idx);
871 let middle = unsafe { Handle::new_kv(self.node, middle_kv_idx) };
fc512014 872 let mut result = middle.split();
1b1a35ee 873 let mut insertion_edge = match insertion {
fc512014
XL
874 LeftOrRight::Left(insert_idx) => unsafe {
875 Handle::new_edge(result.left.reborrow_mut(), insert_idx)
3dfed10e 876 },
fc512014
XL
877 LeftOrRight::Right(insert_idx) => unsafe {
878 Handle::new_edge(result.right.borrow_mut(), insert_idx)
3dfed10e 879 },
9cc50fc6 880 };
1b1a35ee 881 let val_ptr = insertion_edge.insert_fit(key, val);
fc512014 882 (InsertResult::Split(result), val_ptr)
1a4d82fc
JJ
883 }
884 }
885}
886
9cc50fc6 887impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge> {
fc512014
XL
888 /// Fixes the parent pointer and index in the child node that this edge
889 /// links to. This is useful when the ordering of edges has been changed,
29967ef6
XL
890 fn correct_parent_link(self) {
891 // Create backpointer without invalidating other references to the node.
892 let ptr = unsafe { NonNull::new_unchecked(NodeRef::as_internal_ptr(&self.node)) };
893 let idx = self.idx;
9cc50fc6 894 let mut child = self.descend();
29967ef6 895 child.set_parent_link(ptr, idx);
9cc50fc6 896 }
1b1a35ee 897}
1a4d82fc 898
1b1a35ee 899impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge> {
fc512014
XL
900 /// Inserts a new key-value pair and an edge that will go to the right of that new pair
901 /// between this edge and the key-value pair to the right of this edge. This method assumes
a7813a04 902 /// that there is enough space in the node for the new pair to fit.
9cc50fc6 903 fn insert_fit(&mut self, key: K, val: V, edge: Root<K, V>) {
29967ef6 904 debug_assert!(self.node.len() < CAPACITY);
9cc50fc6 905 debug_assert!(edge.height == self.node.height - 1);
fc512014 906 let new_len = self.node.len() + 1;
1a4d82fc 907
9cc50fc6 908 unsafe {
5869c6ff
XL
909 slice_insert(self.node.key_area_mut(..new_len), self.idx, key);
910 slice_insert(self.node.val_area_mut(..new_len), self.idx, val);
911 slice_insert(self.node.edge_area_mut(..new_len + 1), self.idx + 1, edge.node);
fc512014 912 *self.node.len_mut() = new_len as u16;
9cc50fc6 913
fc512014 914 self.node.correct_childrens_parent_links(self.idx + 1..new_len + 1);
9cc50fc6
SL
915 }
916 }
1a4d82fc 917
fc512014
XL
918 /// Inserts a new key-value pair and an edge that will go to the right of that new pair
919 /// between this edge and the key-value pair to the right of this edge. This method splits
a7813a04 920 /// the node if there isn't enough room.
3dfed10e 921 fn insert(
dfeec247
XL
922 mut self,
923 key: K,
924 val: V,
925 edge: Root<K, V>,
926 ) -> InsertResult<'a, K, V, marker::Internal> {
74b04a01 927 assert!(edge.height == self.node.height - 1);
1a4d82fc 928
9cc50fc6
SL
929 if self.node.len() < CAPACITY {
930 self.insert_fit(key, val, edge);
74b04a01
XL
931 let kv = unsafe { Handle::new_kv(self.node, self.idx) };
932 InsertResult::Fit(kv)
9cc50fc6 933 } else {
3dfed10e
XL
934 let (middle_kv_idx, insertion) = splitpoint(self.idx);
935 let middle = unsafe { Handle::new_kv(self.node, middle_kv_idx) };
fc512014 936 let mut result = middle.split();
29967ef6 937 let mut insertion_edge = match insertion {
fc512014
XL
938 LeftOrRight::Left(insert_idx) => unsafe {
939 Handle::new_edge(result.left.reborrow_mut(), insert_idx)
3dfed10e 940 },
fc512014
XL
941 LeftOrRight::Right(insert_idx) => unsafe {
942 Handle::new_edge(result.right.borrow_mut(), insert_idx)
3dfed10e 943 },
29967ef6
XL
944 };
945 insertion_edge.insert_fit(key, val, edge);
fc512014 946 InsertResult::Split(result)
3dfed10e
XL
947 }
948 }
949}
950
1b1a35ee 951impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
fc512014 952 /// Inserts a new key-value pair between the key-value pairs to the right and left of
3dfed10e
XL
953 /// this edge. This method splits the node if there isn't enough room, and tries to
954 /// insert the split off portion into the parent node recursively, until the root is reached.
955 ///
956 /// If the returned result is a `Fit`, its handle's node can be this edge's node or an ancestor.
957 /// If the returned result is a `Split`, the `left` field will be the root node.
958 /// The returned pointer points to the inserted value.
959 pub fn insert_recursing(
960 self,
961 key: K,
962 value: V,
963 ) -> (InsertResult<'a, K, V, marker::LeafOrInternal>, *mut V) {
964 let (mut split, val_ptr) = match self.insert(key, value) {
965 (InsertResult::Fit(handle), ptr) => {
966 return (InsertResult::Fit(handle.forget_node_type()), ptr);
967 }
fc512014 968 (InsertResult::Split(split), val_ptr) => (split.forget_node_type(), val_ptr),
3dfed10e
XL
969 };
970
971 loop {
972 split = match split.left.ascend() {
fc512014 973 Ok(parent) => match parent.insert(split.kv.0, split.kv.1, split.right) {
3dfed10e
XL
974 InsertResult::Fit(handle) => {
975 return (InsertResult::Fit(handle.forget_node_type()), val_ptr);
976 }
fc512014 977 InsertResult::Split(split) => split.forget_node_type(),
3dfed10e
XL
978 },
979 Err(root) => {
980 return (InsertResult::Split(SplitResult { left: root, ..split }), val_ptr);
981 }
982 };
9cc50fc6
SL
983 }
984 }
1a4d82fc
JJ
985}
986
5869c6ff
XL
987impl<BorrowType: marker::BorrowType, K, V>
988 Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge>
989{
a7813a04
XL
990 /// Finds the node pointed to by this edge.
991 ///
fc512014
XL
992 /// The method name assumes you picture trees with the root node on top.
993 ///
a7813a04
XL
994 /// `edge.descend().ascend().unwrap()` and `node.ascend().unwrap().descend()` should
995 /// both, upon success, do nothing.
9cc50fc6 996 pub fn descend(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
5869c6ff 997 assert!(BorrowType::PERMITS_TRAVERSAL);
1b1a35ee
XL
998 // We need to use raw pointers to nodes because, if BorrowType is
999 // marker::ValMut, there might be outstanding mutable references to
1000 // values that we must not invalidate. There's no worry accessing the
1001 // height field because that value is copied. Beware that, once the
1002 // node pointer is dereferenced, we access the edges array with a
1003 // reference (Rust issue #73987) and invalidate any other references
1004 // to or inside the array, should any be around.
29967ef6 1005 let parent_ptr = NodeRef::as_internal_ptr(&self.node);
fc512014
XL
1006 let node = unsafe { (*parent_ptr).edges.get_unchecked(self.idx).assume_init_read() };
1007 NodeRef { node, height: self.node.height - 1, _marker: PhantomData }
9cc50fc6 1008 }
1a4d82fc
JJ
1009}
1010
dfeec247 1011impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Immut<'a>, K, V, NodeType>, marker::KV> {
9cc50fc6 1012 pub fn into_kv(self) -> (&'a K, &'a V) {
5869c6ff
XL
1013 debug_assert!(self.idx < self.node.len());
1014 let leaf = self.node.into_leaf();
1015 let k = unsafe { leaf.keys.get_unchecked(self.idx).assume_init_ref() };
1016 let v = unsafe { leaf.vals.get_unchecked(self.idx).assume_init_ref() };
1017 (k, v)
1a4d82fc 1018 }
9cc50fc6 1019}
1a4d82fc 1020
dfeec247 1021impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV> {
fc512014 1022 pub fn key_mut(&mut self) -> &mut K {
5869c6ff 1023 unsafe { self.node.key_area_mut(self.idx).assume_init_mut() }
3dfed10e
XL
1024 }
1025
1026 pub fn into_val_mut(self) -> &'a mut V {
5869c6ff 1027 debug_assert!(self.idx < self.node.len());
fc512014
XL
1028 let leaf = self.node.into_leaf_mut();
1029 unsafe { leaf.vals.get_unchecked_mut(self.idx).assume_init_mut() }
3dfed10e 1030 }
1b1a35ee 1031}
3dfed10e 1032
1b1a35ee
XL
1033impl<'a, K, V, NodeType> Handle<NodeRef<marker::ValMut<'a>, K, V, NodeType>, marker::KV> {
1034 pub fn into_kv_valmut(self) -> (&'a K, &'a mut V) {
1035 unsafe { self.node.into_key_val_mut_at(self.idx) }
1a4d82fc 1036 }
9cc50fc6 1037}
1a4d82fc 1038
1b1a35ee 1039impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV> {
9cc50fc6 1040 pub fn kv_mut(&mut self) -> (&mut K, &mut V) {
5869c6ff 1041 debug_assert!(self.idx < self.node.len());
29967ef6 1042 // We cannot call separate key and value methods, because calling the second one
1b1a35ee 1043 // invalidates the reference returned by the first.
29967ef6 1044 unsafe {
fc512014 1045 let leaf = self.node.as_leaf_mut();
29967ef6
XL
1046 let key = leaf.keys.get_unchecked_mut(self.idx).assume_init_mut();
1047 let val = leaf.vals.get_unchecked_mut(self.idx).assume_init_mut();
1048 (key, val)
1049 }
1a4d82fc 1050 }
1a4d82fc 1051
17df50a5 1052 /// Replaces the key and value that the KV handle refers to.
fc512014
XL
1053 pub fn replace_kv(&mut self, k: K, v: V) -> (K, V) {
1054 let (key, val) = self.kv_mut();
1055 (mem::replace(key, k), mem::replace(val, v))
29967ef6 1056 }
fc512014 1057}
29967ef6 1058
17df50a5
XL
1059impl<K, V, NodeType> Handle<NodeRef<marker::Dying, K, V, NodeType>, marker::KV> {
1060 /// Extracts the key and value that the KV handle refers to.
94222f64
XL
1061 /// # Safety
1062 /// The node that the handle refers to must not yet have been deallocated.
1063 pub unsafe fn into_key_val(mut self) -> (K, V) {
17df50a5
XL
1064 debug_assert!(self.idx < self.node.len());
1065 let leaf = self.node.as_leaf_dying();
1066 unsafe {
1067 let key = leaf.keys.get_unchecked_mut(self.idx).assume_init_read();
1068 let val = leaf.vals.get_unchecked_mut(self.idx).assume_init_read();
1069 (key, val)
1070 }
1071 }
1072
1073 /// Drops the key and value that the KV handle refers to.
94222f64
XL
1074 /// # Safety
1075 /// The node that the handle refers to must not yet have been deallocated.
17df50a5 1076 #[inline]
94222f64 1077 pub unsafe fn drop_key_val(mut self) {
17df50a5
XL
1078 debug_assert!(self.idx < self.node.len());
1079 let leaf = self.node.as_leaf_dying();
1080 unsafe {
1081 leaf.keys.get_unchecked_mut(self.idx).assume_init_drop();
1082 leaf.vals.get_unchecked_mut(self.idx).assume_init_drop();
1083 }
1084 }
1085}
1086
fc512014 1087impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV> {
3dfed10e
XL
1088 /// Helps implementations of `split` for a particular `NodeType`,
1089 /// by taking care of leaf data.
29967ef6 1090 fn split_leaf_data(&mut self, new_node: &mut LeafNode<K, V>) -> (K, V) {
fc512014 1091 debug_assert!(self.idx < self.node.len());
5869c6ff
XL
1092 let old_len = self.node.len();
1093 let new_len = old_len - self.idx - 1;
29967ef6 1094 new_node.len = new_len as u16;
9cc50fc6 1095 unsafe {
5869c6ff
XL
1096 let k = self.node.key_area_mut(self.idx).assume_init_read();
1097 let v = self.node.val_area_mut(self.idx).assume_init_read();
9cc50fc6 1098
5869c6ff
XL
1099 move_to_slice(
1100 self.node.key_area_mut(self.idx + 1..old_len),
1101 &mut new_node.keys[..new_len],
9cc50fc6 1102 );
5869c6ff
XL
1103 move_to_slice(
1104 self.node.val_area_mut(self.idx + 1..old_len),
1105 &mut new_node.vals[..new_len],
9cc50fc6
SL
1106 );
1107
fc512014 1108 *self.node.len_mut() = self.idx as u16;
29967ef6 1109 (k, v)
3dfed10e
XL
1110 }
1111 }
1112}
1113
1b1a35ee 1114impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
3dfed10e
XL
1115 /// Splits the underlying node into three parts:
1116 ///
fc512014 1117 /// - The node is truncated to only contain the key-value pairs to the left of
3dfed10e 1118 /// this handle.
1b1a35ee 1119 /// - The key and value pointed to by this handle are extracted.
fc512014 1120 /// - All the key-value pairs to the right of this handle are put into a newly
3dfed10e 1121 /// allocated node.
fc512014 1122 pub fn split(mut self) -> SplitResult<'a, K, V, marker::Leaf> {
6a06907d 1123 let mut new_node = LeafNode::new();
3dfed10e 1124
6a06907d 1125 let kv = self.split_leaf_data(&mut new_node);
9cc50fc6 1126
6a06907d
XL
1127 let right = NodeRef::from_new_leaf(new_node);
1128 SplitResult { left: self.node, kv, right }
9cc50fc6 1129 }
1a4d82fc 1130
fc512014
XL
1131 /// Removes the key-value pair pointed to by this handle and returns it, along with the edge
1132 /// that the key-value pair collapsed into.
dfeec247
XL
1133 pub fn remove(
1134 mut self,
3dfed10e 1135 ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
fc512014 1136 let old_len = self.node.len();
9cc50fc6 1137 unsafe {
5869c6ff
XL
1138 let k = slice_remove(self.node.key_area_mut(..old_len), self.idx);
1139 let v = slice_remove(self.node.val_area_mut(..old_len), self.idx);
fc512014 1140 *self.node.len_mut() = (old_len - 1) as u16;
3dfed10e 1141 ((k, v), self.left_edge())
9cc50fc6 1142 }
1a4d82fc
JJ
1143 }
1144}
1145
1b1a35ee 1146impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV> {
a7813a04
XL
1147 /// Splits the underlying node into three parts:
1148 ///
fc512014 1149 /// - The node is truncated to only contain the edges and key-value pairs to the
29967ef6 1150 /// left of this handle.
1b1a35ee 1151 /// - The key and value pointed to by this handle are extracted.
fc512014 1152 /// - All the edges and key-value pairs to the right of this handle are put into
a7813a04 1153 /// a newly allocated node.
fc512014 1154 pub fn split(mut self) -> SplitResult<'a, K, V, marker::Internal> {
5869c6ff 1155 let old_len = self.node.len();
9cc50fc6 1156 unsafe {
6a06907d 1157 let mut new_node = InternalNode::new();
fc512014
XL
1158 let kv = self.split_leaf_data(&mut new_node.data);
1159 let new_len = usize::from(new_node.data.len);
5869c6ff
XL
1160 move_to_slice(
1161 self.node.edge_area_mut(self.idx + 1..old_len + 1),
1162 &mut new_node.edges[..new_len + 1],
9cc50fc6
SL
1163 );
1164
29967ef6 1165 let height = self.node.height;
5869c6ff 1166 let right = NodeRef::from_new_internal(new_node, height);
1a4d82fc 1167
fc512014 1168 SplitResult { left: self.node, kv, right }
9cc50fc6
SL
1169 }
1170 }
fc512014 1171}
1a4d82fc 1172
fc512014
XL
1173/// Represents a session for evaluating and performing a balancing operation
1174/// around an internal key-value pair.
1175pub struct BalancingContext<'a, K, V> {
1176 parent: Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV>,
1177 left_child: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
1178 right_child: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
1179}
1180
1181impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV> {
1182 pub fn consider_for_balancing(self) -> BalancingContext<'a, K, V> {
9cc50fc6
SL
1183 let self1 = unsafe { ptr::read(&self) };
1184 let self2 = unsafe { ptr::read(&self) };
fc512014
XL
1185 BalancingContext {
1186 parent: self,
1187 left_child: self1.left_edge().descend(),
1188 right_child: self2.right_edge().descend(),
1189 }
1190 }
1191}
1192
1193impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
1194 /// Chooses a balancing context involving the node as a child, thus between
1195 /// the KV immediately to the left or to the right in the parent node.
1196 /// Returns an `Err` if there is no parent.
1197 /// Panics if the parent is empty.
1198 ///
1199 /// Prefers the left side, to be optimal if the given node is somehow
1200 /// underfull, meaning here only that it has fewer elements than its left
1201 /// sibling and than its right sibling, if they exist. In that case,
1202 /// merging with the left sibling is faster, since we only need to move
1203 /// the node's N elements, instead of shifting them to the right and moving
1204 /// more than N elements in front. Stealing from the left sibling is also
1205 /// typically faster, since we only need to shift the node's N elements to
1206 /// the right, instead of shifting at least N of the sibling's elements to
1207 /// the left.
1208 pub fn choose_parent_kv(self) -> Result<LeftOrRight<BalancingContext<'a, K, V>>, Self> {
1209 match unsafe { ptr::read(&self) }.ascend() {
1210 Ok(parent_edge) => match parent_edge.left_kv() {
1211 Ok(left_parent_kv) => Ok(LeftOrRight::Left(BalancingContext {
1212 parent: unsafe { ptr::read(&left_parent_kv) },
1213 left_child: left_parent_kv.left_edge().descend(),
1214 right_child: self,
1215 })),
1216 Err(parent_edge) => match parent_edge.right_kv() {
1217 Ok(right_parent_kv) => Ok(LeftOrRight::Right(BalancingContext {
1218 parent: unsafe { ptr::read(&right_parent_kv) },
1219 left_child: self,
1220 right_child: right_parent_kv.right_edge().descend(),
1221 })),
1222 Err(_) => unreachable!("empty internal node"),
1223 },
1224 },
1225 Err(root) => Err(root),
1226 }
1227 }
1228}
1229
1230impl<'a, K, V> BalancingContext<'a, K, V> {
1231 pub fn left_child_len(&self) -> usize {
1232 self.left_child.len()
1233 }
1234
1235 pub fn right_child_len(&self) -> usize {
1236 self.right_child.len()
1237 }
1238
1239 pub fn into_left_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
1240 self.left_child
1241 }
1242
1243 pub fn into_right_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
1244 self.right_child
1245 }
1246
5869c6ff
XL
1247 /// Returns whether merging is possible, i.e., whether there is enough room
1248 /// in a node to combine the central KV with both adjacent child nodes.
fc512014
XL
1249 pub fn can_merge(&self) -> bool {
1250 self.left_child.len() + 1 + self.right_child.len() <= CAPACITY
1251 }
1252}
1253
1254impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
5869c6ff
XL
1255 /// Performs a merge and lets a closure decide what to return.
1256 fn do_merge<
1257 F: FnOnce(
1258 NodeRef<marker::Mut<'a>, K, V, marker::Internal>,
1259 NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
1260 ) -> R,
1261 R,
1262 >(
fc512014 1263 self,
5869c6ff
XL
1264 result: F,
1265 ) -> R {
fc512014
XL
1266 let Handle { node: mut parent_node, idx: parent_idx, _marker } = self.parent;
1267 let old_parent_len = parent_node.len();
1268 let mut left_node = self.left_child;
1269 let old_left_len = left_node.len();
5869c6ff 1270 let mut right_node = self.right_child;
9cc50fc6 1271 let right_len = right_node.len();
fc512014 1272 let new_left_len = old_left_len + 1 + right_len;
9cc50fc6 1273
fc512014 1274 assert!(new_left_len <= CAPACITY);
1a4d82fc 1275
9cc50fc6 1276 unsafe {
fc512014 1277 *left_node.len_mut() = new_left_len as u16;
29967ef6 1278
5869c6ff
XL
1279 let parent_key = slice_remove(parent_node.key_area_mut(..old_parent_len), parent_idx);
1280 left_node.key_area_mut(old_left_len).write(parent_key);
1281 move_to_slice(
1282 right_node.key_area_mut(..right_len),
1283 left_node.key_area_mut(old_left_len + 1..new_left_len),
dfeec247 1284 );
29967ef6 1285
5869c6ff
XL
1286 let parent_val = slice_remove(parent_node.val_area_mut(..old_parent_len), parent_idx);
1287 left_node.val_area_mut(old_left_len).write(parent_val);
1288 move_to_slice(
1289 right_node.val_area_mut(..right_len),
1290 left_node.val_area_mut(old_left_len + 1..new_left_len),
9cc50fc6
SL
1291 );
1292
5869c6ff 1293 slice_remove(&mut parent_node.edge_area_mut(..old_parent_len + 1), parent_idx + 1);
fc512014
XL
1294 parent_node.correct_childrens_parent_links(parent_idx + 1..old_parent_len);
1295 *parent_node.len_mut() -= 1;
7453a54e 1296
fc512014 1297 if parent_node.height > 1 {
3dfed10e
XL
1298 // SAFETY: the height of the nodes being merged is one below the height
1299 // of the node of this edge, thus above zero, so they are internal.
fc512014 1300 let mut left_node = left_node.reborrow_mut().cast_to_internal_unchecked();
5869c6ff
XL
1301 let mut right_node = right_node.cast_to_internal_unchecked();
1302 move_to_slice(
1303 right_node.edge_area_mut(..right_len + 1),
1304 left_node.edge_area_mut(old_left_len + 1..new_left_len + 1),
9cc50fc6
SL
1305 );
1306
fc512014 1307 left_node.correct_childrens_parent_links(old_left_len + 1..new_left_len + 1);
1a4d82fc 1308
fc512014 1309 Global.deallocate(right_node.node.cast(), Layout::new::<InternalNode<K, V>>());
9cc50fc6 1310 } else {
fc512014 1311 Global.deallocate(right_node.node.cast(), Layout::new::<LeafNode<K, V>>());
3dfed10e 1312 }
9cc50fc6 1313 }
5869c6ff
XL
1314 result(parent_node, left_node)
1315 }
1316
1317 /// Merges the parent's key-value pair and both adjacent child nodes into
1318 /// the left child node and returns the shrunk parent node.
1319 ///
1320 /// Panics unless we `.can_merge()`.
1321 pub fn merge_tracking_parent(self) -> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
1322 self.do_merge(|parent, _child| parent)
1323 }
1324
1325 /// Merges the parent's key-value pair and both adjacent child nodes into
1326 /// the left child node and returns that child node.
1327 ///
1328 /// Panics unless we `.can_merge()`.
1329 pub fn merge_tracking_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
1330 self.do_merge(|_parent, child| child)
1331 }
1332
1333 /// Merges the parent's key-value pair and both adjacent child nodes into
1334 /// the left child node and returns the edge handle in that child node
1335 /// where the tracked child edge ended up,
1336 ///
1337 /// Panics unless we `.can_merge()`.
1338 pub fn merge_tracking_child_edge(
1339 self,
1340 track_edge_idx: LeftOrRight<usize>,
1341 ) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Edge> {
1342 let old_left_len = self.left_child.len();
1343 let right_len = self.right_child.len();
1344 assert!(match track_edge_idx {
1345 LeftOrRight::Left(idx) => idx <= old_left_len,
1346 LeftOrRight::Right(idx) => idx <= right_len,
1347 });
1348 let child = self.merge_tracking_child();
1349 let new_idx = match track_edge_idx {
1350 LeftOrRight::Left(idx) => idx,
1351 LeftOrRight::Right(idx) => old_left_len + 1 + idx,
1352 };
1353 unsafe { Handle::new_edge(child, new_idx) }
85aaf69f 1354 }
a7813a04 1355
fc512014
XL
1356 /// Removes a key-value pair from the left child and places it in the key-value storage
1357 /// of the parent, while pushing the old parent key-value pair into the right child.
1358 /// Returns a handle to the edge in the right child corresponding to where the original
1359 /// edge specified by `track_right_edge_idx` ended up.
1360 pub fn steal_left(
1361 mut self,
1362 track_right_edge_idx: usize,
1363 ) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Edge> {
5869c6ff
XL
1364 self.bulk_steal_left(1);
1365 unsafe { Handle::new_edge(self.right_child, 1 + track_right_edge_idx) }
a7813a04
XL
1366 }
1367
fc512014
XL
1368 /// Removes a key-value pair from the right child and places it in the key-value storage
1369 /// of the parent, while pushing the old parent key-value pair onto the left child.
1370 /// Returns a handle to the edge in the left child specified by `track_left_edge_idx`,
1371 /// which didn't move.
1372 pub fn steal_right(
1373 mut self,
1374 track_left_edge_idx: usize,
1375 ) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Edge> {
5869c6ff
XL
1376 self.bulk_steal_right(1);
1377 unsafe { Handle::new_edge(self.left_child, track_left_edge_idx) }
a7813a04
XL
1378 }
1379
1380 /// This does stealing similar to `steal_left` but steals multiple elements at once.
3157f602 1381 pub fn bulk_steal_left(&mut self, count: usize) {
fc512014 1382 assert!(count > 0);
a7813a04 1383 unsafe {
fc512014
XL
1384 let left_node = &mut self.left_child;
1385 let old_left_len = left_node.len();
1386 let right_node = &mut self.right_child;
1387 let old_right_len = right_node.len();
a7813a04
XL
1388
1389 // Make sure that we may steal safely.
fc512014
XL
1390 assert!(old_right_len + count <= CAPACITY);
1391 assert!(old_left_len >= count);
3157f602 1392
fc512014
XL
1393 let new_left_len = old_left_len - count;
1394 let new_right_len = old_right_len + count;
1395 *left_node.len_mut() = new_left_len as u16;
1396 *right_node.len_mut() = new_right_len as u16;
3157f602 1397
fc512014 1398 // Move leaf data.
3157f602 1399 {
3157f602 1400 // Make room for stolen elements in the right child.
5869c6ff
XL
1401 slice_shr(right_node.key_area_mut(..new_right_len), count);
1402 slice_shr(right_node.val_area_mut(..new_right_len), count);
3157f602
XL
1403
1404 // Move elements from the left child to the right one.
5869c6ff
XL
1405 move_to_slice(
1406 left_node.key_area_mut(new_left_len + 1..old_left_len),
1407 right_node.key_area_mut(..count - 1),
1408 );
1409 move_to_slice(
1410 left_node.val_area_mut(new_left_len + 1..old_left_len),
1411 right_node.val_area_mut(..count - 1),
1412 );
3157f602
XL
1413
1414 // Move the left-most stolen pair to the parent.
5869c6ff
XL
1415 let k = left_node.key_area_mut(new_left_len).assume_init_read();
1416 let v = left_node.val_area_mut(new_left_len).assume_init_read();
1417 let (k, v) = self.parent.replace_kv(k, v);
1418
1419 // Move parent's key-value pair to the right child.
1420 right_node.key_area_mut(count - 1).write(k);
1421 right_node.val_area_mut(count - 1).write(v);
a7813a04
XL
1422 }
1423
fc512014 1424 match (left_node.reborrow_mut().force(), right_node.reborrow_mut().force()) {
5869c6ff 1425 (ForceResult::Internal(mut left), ForceResult::Internal(mut right)) => {
3157f602 1426 // Make room for stolen edges.
5869c6ff 1427 slice_shr(right.edge_area_mut(..new_right_len + 1), count);
3157f602 1428
fc512014 1429 // Steal edges.
5869c6ff
XL
1430 move_to_slice(
1431 left.edge_area_mut(new_left_len + 1..old_left_len + 1),
1432 right.edge_area_mut(..count),
1433 );
1434
1435 right.correct_childrens_parent_links(0..new_right_len + 1);
dfeec247
XL
1436 }
1437 (ForceResult::Leaf(_), ForceResult::Leaf(_)) => {}
1b1a35ee 1438 _ => unreachable!(),
a7813a04 1439 }
3157f602
XL
1440 }
1441 }
a7813a04 1442
3157f602
XL
1443 /// The symmetric clone of `bulk_steal_left`.
1444 pub fn bulk_steal_right(&mut self, count: usize) {
fc512014 1445 assert!(count > 0);
3157f602 1446 unsafe {
fc512014
XL
1447 let left_node = &mut self.left_child;
1448 let old_left_len = left_node.len();
1449 let right_node = &mut self.right_child;
1450 let old_right_len = right_node.len();
3157f602
XL
1451
1452 // Make sure that we may steal safely.
fc512014
XL
1453 assert!(old_left_len + count <= CAPACITY);
1454 assert!(old_right_len >= count);
3157f602 1455
fc512014
XL
1456 let new_left_len = old_left_len + count;
1457 let new_right_len = old_right_len - count;
1458 *left_node.len_mut() = new_left_len as u16;
1459 *right_node.len_mut() = new_right_len as u16;
3157f602 1460
fc512014 1461 // Move leaf data.
3157f602 1462 {
5869c6ff
XL
1463 // Move the right-most stolen pair to the parent.
1464 let k = right_node.key_area_mut(count - 1).assume_init_read();
1465 let v = right_node.val_area_mut(count - 1).assume_init_read();
1466 let (k, v) = self.parent.replace_kv(k, v);
3157f602 1467
fc512014 1468 // Move parent's key-value pair to the left child.
5869c6ff
XL
1469 left_node.key_area_mut(old_left_len).write(k);
1470 left_node.val_area_mut(old_left_len).write(v);
3157f602
XL
1471
1472 // Move elements from the right child to the left one.
5869c6ff
XL
1473 move_to_slice(
1474 right_node.key_area_mut(..count - 1),
1475 left_node.key_area_mut(old_left_len + 1..new_left_len),
1476 );
1477 move_to_slice(
1478 right_node.val_area_mut(..count - 1),
1479 left_node.val_area_mut(old_left_len + 1..new_left_len),
1480 );
3157f602 1481
fc512014 1482 // Fill gap where stolen elements used to be.
5869c6ff
XL
1483 slice_shl(right_node.key_area_mut(..old_right_len), count);
1484 slice_shl(right_node.val_area_mut(..old_right_len), count);
3157f602
XL
1485 }
1486
fc512014 1487 match (left_node.reborrow_mut().force(), right_node.reborrow_mut().force()) {
5869c6ff 1488 (ForceResult::Internal(mut left), ForceResult::Internal(mut right)) => {
fc512014 1489 // Steal edges.
5869c6ff
XL
1490 move_to_slice(
1491 right.edge_area_mut(..count),
1492 left.edge_area_mut(old_left_len + 1..new_left_len + 1),
1493 );
3157f602 1494
fc512014 1495 // Fill gap where stolen edges used to be.
5869c6ff
XL
1496 slice_shl(right.edge_area_mut(..old_right_len + 1), count);
1497
1498 left.correct_childrens_parent_links(old_left_len + 1..new_left_len + 1);
1499 right.correct_childrens_parent_links(0..new_right_len + 1);
dfeec247
XL
1500 }
1501 (ForceResult::Leaf(_), ForceResult::Leaf(_)) => {}
1b1a35ee 1502 _ => unreachable!(),
a7813a04
XL
1503 }
1504 }
1505 }
85aaf69f
SL
1506}
1507
74b04a01
XL
1508impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
1509 pub fn forget_node_type(
1510 self,
1511 ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
1512 unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
1513 }
1514}
1515
1516impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
1517 pub fn forget_node_type(
1518 self,
1519 ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
1520 unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
1521 }
1522}
1523
1524impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::KV> {
1525 pub fn forget_node_type(
1526 self,
1527 ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
1528 unsafe { Handle::new_kv(self.node.forget_type(), self.idx) }
1529 }
1530}
1531
3dfed10e
XL
1532impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV> {
1533 pub fn forget_node_type(
1534 self,
1535 ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
1536 unsafe { Handle::new_kv(self.node.forget_type(), self.idx) }
1537 }
1538}
1539
6a06907d 1540impl<BorrowType, K, V, Type> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, Type> {
9fa01778 1541 /// Checks whether the underlying node is an `Internal` node or a `Leaf` node.
dfeec247
XL
1542 pub fn force(
1543 self,
1544 ) -> ForceResult<
6a06907d
XL
1545 Handle<NodeRef<BorrowType, K, V, marker::Leaf>, Type>,
1546 Handle<NodeRef<BorrowType, K, V, marker::Internal>, Type>,
9cc50fc6
SL
1547 > {
1548 match self.node.force() {
dfeec247
XL
1549 ForceResult::Leaf(node) => {
1550 ForceResult::Leaf(Handle { node, idx: self.idx, _marker: PhantomData })
1551 }
1552 ForceResult::Internal(node) => {
1553 ForceResult::Internal(Handle { node, idx: self.idx, _marker: PhantomData })
1554 }
9cc50fc6 1555 }
85aaf69f
SL
1556 }
1557}
1a4d82fc 1558
cdc7bbd5
XL
1559impl<'a, K, V, Type> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, Type> {
1560 /// Unsafely asserts to the compiler the static information that the handle's node is a `Leaf`.
1561 pub unsafe fn cast_to_leaf_unchecked(
1562 self,
1563 ) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, Type> {
1564 let node = unsafe { self.node.cast_to_leaf_unchecked() };
1565 Handle { node, idx: self.idx, _marker: PhantomData }
1566 }
1567}
1568
3157f602
XL
1569impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Edge> {
1570 /// Move the suffix after `self` from one node to another one. `right` must be empty.
1571 /// The first edge of `right` remains unchanged.
dfeec247
XL
1572 pub fn move_suffix(
1573 &mut self,
1574 right: &mut NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
1575 ) {
3157f602 1576 unsafe {
fc512014 1577 let new_left_len = self.idx;
3157f602 1578 let mut left_node = self.reborrow_mut().into_node();
5869c6ff 1579 let old_left_len = left_node.len();
3157f602 1580
5869c6ff 1581 let new_right_len = old_left_len - new_left_len;
3157f602
XL
1582 let mut right_node = right.reborrow_mut();
1583
74b04a01
XL
1584 assert!(right_node.len() == 0);
1585 assert!(left_node.height == right_node.height);
3157f602 1586
fc512014 1587 if new_right_len > 0 {
fc512014
XL
1588 *left_node.len_mut() = new_left_len as u16;
1589 *right_node.len_mut() = new_right_len as u16;
3157f602 1590
5869c6ff
XL
1591 move_to_slice(
1592 left_node.key_area_mut(new_left_len..old_left_len),
1593 right_node.key_area_mut(..new_right_len),
1594 );
1595 move_to_slice(
1596 left_node.val_area_mut(new_left_len..old_left_len),
1597 right_node.val_area_mut(..new_right_len),
1598 );
74b04a01 1599 match (left_node.force(), right_node.force()) {
5869c6ff
XL
1600 (ForceResult::Internal(mut left), ForceResult::Internal(mut right)) => {
1601 move_to_slice(
1602 left.edge_area_mut(new_left_len + 1..old_left_len + 1),
1603 right.edge_area_mut(1..new_right_len + 1),
1604 );
1605 right.correct_childrens_parent_links(1..new_right_len + 1);
74b04a01
XL
1606 }
1607 (ForceResult::Leaf(_), ForceResult::Leaf(_)) => {}
1b1a35ee 1608 _ => unreachable!(),
dfeec247 1609 }
3157f602
XL
1610 }
1611 }
1612 }
1613}
1614
9cc50fc6
SL
1615pub enum ForceResult<Leaf, Internal> {
1616 Leaf(Leaf),
dfeec247 1617 Internal(Internal),
9cc50fc6 1618}
85aaf69f 1619
3dfed10e 1620/// Result of insertion, when a node needed to expand beyond its capacity.
fc512014
XL
1621pub struct SplitResult<'a, K, V, NodeType> {
1622 // Altered node in existing tree with elements and edges that belong to the left of `kv`.
1623 pub left: NodeRef<marker::Mut<'a>, K, V, NodeType>,
3dfed10e 1624 // Some key and value split off, to be inserted elsewhere.
fc512014
XL
1625 pub kv: (K, V),
1626 // Owned, unattached, new node with elements and edges that belong to the right of `kv`.
1627 pub right: NodeRef<marker::Owned, K, V, NodeType>,
1628}
1629
1630impl<'a, K, V> SplitResult<'a, K, V, marker::Leaf> {
1631 pub fn forget_node_type(self) -> SplitResult<'a, K, V, marker::LeafOrInternal> {
1632 SplitResult { left: self.left.forget_type(), kv: self.kv, right: self.right.forget_type() }
1633 }
1634}
1635
1636impl<'a, K, V> SplitResult<'a, K, V, marker::Internal> {
1637 pub fn forget_node_type(self) -> SplitResult<'a, K, V, marker::LeafOrInternal> {
1638 SplitResult { left: self.left.forget_type(), kv: self.kv, right: self.right.forget_type() }
1639 }
3dfed10e
XL
1640}
1641
fc512014
XL
1642pub enum InsertResult<'a, K, V, NodeType> {
1643 Fit(Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV>),
1644 Split(SplitResult<'a, K, V, NodeType>),
1a4d82fc
JJ
1645}
1646
9cc50fc6
SL
1647pub mod marker {
1648 use core::marker::PhantomData;
1a4d82fc 1649
dfeec247
XL
1650 pub enum Leaf {}
1651 pub enum Internal {}
1652 pub enum LeafOrInternal {}
85aaf69f 1653
dfeec247 1654 pub enum Owned {}
5869c6ff 1655 pub enum Dying {}
9cc50fc6
SL
1656 pub struct Immut<'a>(PhantomData<&'a ()>);
1657 pub struct Mut<'a>(PhantomData<&'a mut ()>);
1b1a35ee 1658 pub struct ValMut<'a>(PhantomData<&'a mut ()>);
85aaf69f 1659
5869c6ff
XL
1660 pub trait BorrowType {
1661 // Whether node references of this borrow type allow traversing
1662 // to other nodes in the tree.
1663 const PERMITS_TRAVERSAL: bool = true;
1664 }
1665 impl BorrowType for Owned {
c295e0f8 1666 // Traversal isn't needed, it happens using the result of `borrow_mut`.
5869c6ff
XL
1667 // By disabling traversal, and only creating new references to roots,
1668 // we know that every reference of the `Owned` type is to a root node.
1669 const PERMITS_TRAVERSAL: bool = false;
1670 }
1671 impl BorrowType for Dying {}
1672 impl<'a> BorrowType for Immut<'a> {}
1673 impl<'a> BorrowType for Mut<'a> {}
1674 impl<'a> BorrowType for ValMut<'a> {}
1675
dfeec247
XL
1676 pub enum KV {}
1677 pub enum Edge {}
1a4d82fc 1678}
85aaf69f 1679
29967ef6
XL
1680/// Inserts a value into a slice of initialized elements followed by one uninitialized element.
1681///
1682/// # Safety
1683/// The slice has more than `idx` elements.
1684unsafe fn slice_insert<T>(slice: &mut [MaybeUninit<T>], idx: usize, val: T) {
f035d41b 1685 unsafe {
29967ef6
XL
1686 let len = slice.len();
1687 debug_assert!(len > idx);
1688 let slice_ptr = slice.as_mut_ptr();
1689 if len > idx + 1 {
1690 ptr::copy(slice_ptr.add(idx), slice_ptr.add(idx + 1), len - idx - 1);
1691 }
1692 (*slice_ptr.add(idx)).write(val);
f035d41b 1693 }
9cc50fc6
SL
1694}
1695
29967ef6
XL
1696/// Removes and returns a value from a slice of all initialized elements, leaving behind one
1697/// trailing uninitialized element.
1698///
1699/// # Safety
1700/// The slice has more than `idx` elements.
1701unsafe fn slice_remove<T>(slice: &mut [MaybeUninit<T>], idx: usize) -> T {
f035d41b 1702 unsafe {
29967ef6
XL
1703 let len = slice.len();
1704 debug_assert!(idx < len);
1705 let slice_ptr = slice.as_mut_ptr();
1706 let ret = (*slice_ptr.add(idx)).assume_init_read();
1707 ptr::copy(slice_ptr.add(idx + 1), slice_ptr.add(idx), len - idx - 1);
f035d41b
XL
1708 ret
1709 }
9cc50fc6 1710}
3dfed10e 1711
5869c6ff
XL
1712/// Shifts the elements in a slice `distance` positions to the left.
1713///
1714/// # Safety
1715/// The slice has at least `distance` elements.
1716unsafe fn slice_shl<T>(slice: &mut [MaybeUninit<T>], distance: usize) {
1717 unsafe {
1718 let slice_ptr = slice.as_mut_ptr();
1719 ptr::copy(slice_ptr.add(distance), slice_ptr, slice.len() - distance);
1720 }
1721}
1722
1723/// Shifts the elements in a slice `distance` positions to the right.
1724///
1725/// # Safety
1726/// The slice has at least `distance` elements.
1727unsafe fn slice_shr<T>(slice: &mut [MaybeUninit<T>], distance: usize) {
1728 unsafe {
1729 let slice_ptr = slice.as_mut_ptr();
1730 ptr::copy(slice_ptr, slice_ptr.add(distance), slice.len() - distance);
1731 }
1732}
1733
1734/// Moves all values from a slice of initialized elements to a slice
1735/// of uninitialized elements, leaving behind `src` as all uninitialized.
1736/// Works like `dst.copy_from_slice(src)` but does not require `T` to be `Copy`.
1737fn move_to_slice<T>(src: &mut [MaybeUninit<T>], dst: &mut [MaybeUninit<T>]) {
1738 assert!(src.len() == dst.len());
1739 unsafe {
1740 ptr::copy_nonoverlapping(src.as_ptr(), dst.as_mut_ptr(), src.len());
1741 }
1742}
1743
3dfed10e
XL
1744#[cfg(test)]
1745mod tests;