Auto merge of rust-lang#78631 - ssomers:btree-alias_for_underfull, r=Mark-Simulacrum

BTreeMap: fix pointer provenance rules in underfullness

Continuing on rust-lang#78480, and for readability, and possibly for performance: avoid aliasing when handling underfull nodes, and consolidate the code doing that. In particular:
- Avoid the rather explicit aliasing for internal nodes in `remove_kv_tracking`.
- Climb down to the root to handle underfull nodes using a reborrowed handle, rather than one copied with `ptr::read`, before resuming on the leaf level.
- Integrate the code tracking leaf edge position into the functions performing changes, rather than bolting it on.

r? `@Mark-Simulacrum`

Author: bors

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

@@ -89,20 +89,15 @@ impl<K, V> Root<K, V> {
         let mut cur_node = self.node_as_mut();
         while let Internal(internal) = cur_node.force() {
             // Check if right-most child is underfull.
-            let mut last_edge = internal.last_edge();
-            let right_child_len = last_edge.reborrow().descend().len();
+            let mut last_kv = internal.last_kv().consider_for_balancing();
+            let right_child_len = last_kv.right_child_len();
             if right_child_len < MIN_LEN {
                 // We need to steal.
-                let mut last_kv = match last_edge.left_kv() {
-                    Ok(left) => left,
-                    Err(_) => unreachable!(),
-                };
                 last_kv.bulk_steal_left(MIN_LEN - right_child_len);
-                last_edge = last_kv.right_edge();
             }
 
             // Go further down.
-            cur_node = last_edge.descend();
+            cur_node = last_kv.into_right_child();
         }
     }
 }

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

@@ -6,6 +6,7 @@ use core::ptr;
 /// relevant function.
 ///
 /// If a panic occurs in the `change` closure, the entire process will be aborted.
+#[allow(dead_code)] // keep as illustration and for future use
 #[inline]
 pub fn take_mut<T>(v: &mut T, change: impl FnOnce(T) -> T) {
     replace(v, |value| (change(value), ()))

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

@@ -362,20 +362,6 @@ impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::E
     }
 }
 
-impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
-    /// Moves the leaf edge handle to the next leaf edge.
-    ///
-    /// # Safety
-    /// There must be another KV in the direction travelled.
-    pub unsafe fn move_next_unchecked(&mut self) {
-        super::mem::take_mut(self, |leaf_edge| {
-            let kv = leaf_edge.next_kv();
-            let kv = unsafe { unwrap_unchecked(kv.ok()) };
-            kv.next_leaf_edge()
-        })
-    }
-}
-
 impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
     /// Moves the leaf edge handle to the next leaf edge and returns the key and value
     /// in between, deallocating any node left behind while leaving the corresponding

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

@@ -54,11 +54,11 @@ fn test_splitpoint() {
         let mut left_len = middle_kv_idx;
         let mut right_len = CAPACITY - middle_kv_idx - 1;
         match insertion {
-            InsertionPlace::Left(edge_idx) => {
+            LeftOrRight::Left(edge_idx) => {
                 assert!(edge_idx <= left_len);
                 left_len += 1;
             }
-            InsertionPlace::Right(edge_idx) => {
+            LeftOrRight::Right(edge_idx) => {
                 assert!(edge_idx <= right_len);
                 right_len += 1;
             }

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

@@ -1,133 +1,153 @@
 use super::map::MIN_LEN;
-use super::node::{marker, ForceResult, Handle, NodeRef};
+use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef};
 use super::unwrap_unchecked;
 use core::mem;
-use core::ptr;
 
 impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
-    /// Removes a key/value-pair from the map, and returns that pair, as well as
-    /// the leaf edge corresponding to that former pair.
+    /// Removes a key/value-pair from the tree, and returns that pair, as well as
+    /// the leaf edge corresponding to that former pair. It's possible this empties
+    /// a root node that is internal, which the caller should pop from the map
+    /// holding the tree. The caller should also decrement the map's length.
     pub fn remove_kv_tracking<F: FnOnce()>(
         self,
         handle_emptied_internal_root: F,
     ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
-        let (old_kv, mut pos, was_internal) = match self.force() {
-            ForceResult::Leaf(leaf) => {
-                let (old_kv, pos) = leaf.remove();
-                (old_kv, pos, false)
-            }
-            ForceResult::Internal(mut internal) => {
-                // Replace the location freed in the internal node with an
-                // adjacent KV, and remove that adjacent KV from its leaf.
-                // Always choose the adjacent KV on the left side because
-                // it is typically faster to pop an element from the end
-                // of the KV arrays without needing to shift other elements.
-
-                let key_loc = internal.kv_mut().0 as *mut K;
-                let val_loc = internal.kv_mut().1 as *mut V;
-
-                let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok();
-                let to_remove = unsafe { unwrap_unchecked(to_remove) };
-
-                let (kv, pos) = to_remove.remove();
-
-                let old_key = unsafe { mem::replace(&mut *key_loc, kv.0) };
-                let old_val = unsafe { mem::replace(&mut *val_loc, kv.1) };
-
-                ((old_key, old_val), pos, true)
-            }
-        };
-
-        // Handle underflow
-        let mut cur_node = unsafe { ptr::read(&pos).into_node().forget_type() };
-        let mut at_leaf = true;
-        while cur_node.len() < MIN_LEN {
-            match handle_underfull_node(cur_node) {
-                UnderflowResult::AtRoot => break,
-                UnderflowResult::Merged(edge, merged_with_left, offset) => {
-                    // If we merged with our right sibling then our tracked
-                    // position has not changed. However if we merged with our
-                    // left sibling then our tracked position is now dangling.
-                    if at_leaf && merged_with_left {
-                        let idx = pos.idx() + offset;
-                        let node = match unsafe { ptr::read(&edge).descend().force() } {
-                            ForceResult::Leaf(leaf) => leaf,
-                            ForceResult::Internal(_) => unreachable!(),
-                        };
-                        pos = unsafe { Handle::new_edge(node, idx) };
-                    }
+        match self.force() {
+            Leaf(node) => node.remove_leaf_kv(handle_emptied_internal_root),
+            Internal(node) => node.remove_internal_kv(handle_emptied_internal_root),
+        }
+    }
+}
 
-                    let parent = edge.into_node();
-                    if parent.len() == 0 {
-                        // The parent that was just emptied must be the root,
-                        // because nodes on a lower level would not have been
-                        // left with a single child.
-                        handle_emptied_internal_root();
-                        break;
+impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
+    fn remove_leaf_kv<F: FnOnce()>(
+        self,
+        handle_emptied_internal_root: F,
+    ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
+        let (old_kv, mut pos) = self.remove();
+        let len = pos.reborrow().into_node().len();
+        if len < MIN_LEN {
+            let idx = pos.idx();
+            // We have to temporarily forget the child type, because there is no
+            // distinct node type for the immediate parents of a leaf.
+            let new_pos = match pos.into_node().forget_type().choose_parent_kv() {
+                Ok(Left(left_parent_kv)) => {
+                    debug_assert!(left_parent_kv.right_child_len() == MIN_LEN - 1);
+                    if left_parent_kv.can_merge() {
+                        left_parent_kv.merge(Some(Right(idx)))
                     } else {
-                        cur_node = parent.forget_type();
-                        at_leaf = false;
+                        debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
+                        left_parent_kv.steal_left(idx)
                     }
                 }
-                UnderflowResult::Stole(stole_from_left) => {
-                    // Adjust the tracked position if we stole from a left sibling
-                    if stole_from_left && at_leaf {
-                        // SAFETY: This is safe since we just added an element to our node.
-                        unsafe {
-                            pos.move_next_unchecked();
-                        }
+                Ok(Right(right_parent_kv)) => {
+                    debug_assert!(right_parent_kv.left_child_len() == MIN_LEN - 1);
+                    if right_parent_kv.can_merge() {
+                        right_parent_kv.merge(Some(Left(idx)))
+                    } else {
+                        debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
+                        right_parent_kv.steal_right(idx)
                     }
-                    break;
                 }
-            }
-        }
+                Err(pos) => unsafe { Handle::new_edge(pos, idx) },
+            };
+            // SAFETY: `new_pos` is the leaf we started from or a sibling.
+            pos = unsafe { new_pos.cast_to_leaf_unchecked() };
 
-        // If we deleted from an internal node then we need to compensate for
-        // the earlier swap and adjust the tracked position to point to the
-        // next element.
-        if was_internal {
-            pos = unsafe { unwrap_unchecked(pos.next_kv().ok()).next_leaf_edge() };
+            // Only if we merged, the parent (if any) has shrunk, but skipping
+            // the following step does not pay off in benchmarks.
+            //
+            // SAFETY: We won't destroy or rearrange the leaf where `pos` is at
+            // by handling its parent recursively; at worst we will destroy or
+            // rearrange the parent through the grandparent, thus change the
+            // leaf's parent pointer.
+            if let Ok(parent) = unsafe { pos.reborrow_mut() }.into_node().ascend() {
+                parent.into_node().handle_shrunk_node_recursively(handle_emptied_internal_root);
+            }
         }
-
         (old_kv, pos)
     }
 }
 
-enum UnderflowResult<'a, K, V> {
-    AtRoot,
-    Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
-    Stole(bool),
-}
+impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV> {
+    fn remove_internal_kv<F: FnOnce()>(
+        self,
+        handle_emptied_internal_root: F,
+    ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
+        // Remove an adjacent KV from its leaf and then put it back in place of
+        // the element we were asked to remove. Prefer the left adjacent KV,
+        // for the reasons listed in `choose_parent_kv`.
+        let left_leaf_kv = self.left_edge().descend().last_leaf_edge().left_kv();
+        let left_leaf_kv = unsafe { unwrap_unchecked(left_leaf_kv.ok()) };
+        let (left_kv, left_hole) = left_leaf_kv.remove_leaf_kv(handle_emptied_internal_root);
 
-fn handle_underfull_node<'a, K: 'a, V: 'a>(
-    node: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
-) -> UnderflowResult<'_, K, V> {
-    let parent = match node.ascend() {
-        Ok(parent) => parent,
-        Err(_) => return UnderflowResult::AtRoot,
-    };
+        // The internal node may have been stolen from or merged. Go back right
+        // to find where the original KV ended up.
+        let mut internal = unsafe { unwrap_unchecked(left_hole.next_kv().ok()) };
+        let old_key = mem::replace(internal.kv_mut().0, left_kv.0);
+        let old_val = mem::replace(internal.kv_mut().1, left_kv.1);
+        let pos = internal.next_leaf_edge();
+        ((old_key, old_val), pos)
+    }
+}
 
-    // Prefer the left KV if it exists. Merging with the left side is faster,
-    // since merging happens towards the left and `node` has fewer elements.
-    // Stealing from the left side is faster, since we can pop from the end of
-    // the KV arrays.
-    let (is_left, mut handle) = match parent.left_kv() {
-        Ok(left) => (true, left),
-        Err(parent) => {
-            let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
-            (false, right)
+impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
+    /// Stocks up a possibly underfull internal node, recursively.
+    /// Climbs up until it reaches an ancestor that has elements to spare or the root.
+    fn handle_shrunk_node_recursively<F: FnOnce()>(mut self, handle_emptied_internal_root: F) {
+        loop {
+            self = match self.len() {
+                0 => {
+                    // An empty node must be the root, because length is only
+                    // reduced by one, and non-root underfull nodes are stocked up,
+                    // so non-root nodes never have fewer than MIN_LEN - 1 elements.
+                    debug_assert!(self.ascend().is_err());
+                    handle_emptied_internal_root();
+                    return;
+                }
+                1..MIN_LEN => {
+                    if let Some(parent) = self.handle_underfull_node_locally() {
+                        parent
+                    } else {
+                        return;
+                    }
+                }
+                _ => return,
+            }
         }
-    };
+    }
 
-    if handle.can_merge() {
-        let offset = if is_left { handle.reborrow().left_edge().descend().len() + 1 } else { 0 };
-        UnderflowResult::Merged(handle.merge(), is_left, offset)
-    } else {
-        if is_left {
-            handle.steal_left();
-        } else {
-            handle.steal_right();
+    /// Stocks up an underfull internal node, possibly at the cost of shrinking
+    /// its parent instead, which is then returned.
+    fn handle_underfull_node_locally(
+        self,
+    ) -> Option<NodeRef<marker::Mut<'a>, K, V, marker::Internal>> {
+        match self.forget_type().choose_parent_kv() {
+            Ok(Left(left_parent_kv)) => {
+                debug_assert!(left_parent_kv.right_child_len() == MIN_LEN - 1);
+                if left_parent_kv.can_merge() {
+                    let pos = left_parent_kv.merge(None);
+                    let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) };
+                    Some(parent_edge.into_node())
+                } else {
+                    debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
+                    left_parent_kv.steal_left(0);
+                    None
+                }
+            }
+            Ok(Right(right_parent_kv)) => {
+                debug_assert!(right_parent_kv.left_child_len() == MIN_LEN - 1);
+                if right_parent_kv.can_merge() {
+                    let pos = right_parent_kv.merge(None);
+                    let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) };
+                    Some(parent_edge.into_node())
+                } else {
+                    debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
+                    right_parent_kv.steal_right(0);
+                    None
+                }
+            }
+            Err(_) => None,
         }
-        UnderflowResult::Stole(is_left)
     }
 }

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

@@ -60,17 +60,17 @@ impl<K, V> Root<K, V> {
             let mut cur_node = self.node_as_mut();
 
             while let Internal(node) = cur_node.force() {
-                let mut last_kv = node.last_kv();
+                let mut last_kv = node.last_kv().consider_for_balancing();
 
                 if last_kv.can_merge() {
-                    cur_node = last_kv.merge().descend();
+                    cur_node = last_kv.merge(None).into_node();
                 } else {
-                    let right_len = last_kv.reborrow().right_edge().descend().len();
+                    let right_len = last_kv.right_child_len();
                     // `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
                     if right_len < MIN_LEN + 1 {
                         last_kv.bulk_steal_left(MIN_LEN + 1 - right_len);
                     }
-                    cur_node = last_kv.right_edge().descend();
+                    cur_node = last_kv.into_right_child();
                 }
             }
         }
@@ -86,17 +86,17 @@ impl<K, V> Root<K, V> {
             let mut cur_node = self.node_as_mut();
 
             while let Internal(node) = cur_node.force() {
-                let mut first_kv = node.first_kv();
+                let mut first_kv = node.first_kv().consider_for_balancing();
 
                 if first_kv.can_merge() {
-                    cur_node = first_kv.merge().descend();
+                    cur_node = first_kv.merge(None).into_node();
                 } else {
-                    let left_len = first_kv.reborrow().left_edge().descend().len();
+                    let left_len = first_kv.left_child_len();
                     // `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
                     if left_len < MIN_LEN + 1 {
                         first_kv.bulk_steal_right(MIN_LEN + 1 - left_len);
                     }
-                    cur_node = first_kv.left_edge().descend();
+                    cur_node = first_kv.into_left_child();
                 }
             }
         }