library/alloc/src/collections/btree/node.rs

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