From 7c95768dbd8a9e376c3d8724a28f327d556686e1 Mon Sep 17 00:00:00 2001
From: Marijn Schouten <mhkbst@gmail.com>
Date: Fri, 17 Oct 2025 15:22:54 +0000
Subject: [PATCH] btree: some cleanup with less unsafe

---
 library/alloc/src/collections/btree/node.rs | 33 +++++++++------------
 1 file changed, 14 insertions(+), 19 deletions(-)

diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
index a87259e7c58f2..84dd4b7e49def 100644
--- a/library/alloc/src/collections/btree/node.rs
+++ b/library/alloc/src/collections/btree/node.rs
@@ -33,6 +33,7 @@
 
 use core::marker::PhantomData;
 use core::mem::{self, MaybeUninit};
+use core::num::NonZero;
 use core::ptr::{self, NonNull};
 use core::slice::SliceIndex;
 
@@ -143,7 +144,7 @@ type BoxedNode<K, V> = NonNull<LeafNode<K, V>>;
 ///
 /// A reference to a node.
 ///
-/// This type has a number of parameters that controls how it acts:
+/// This type has a number of parameters that control how it acts:
 /// - `BorrowType`: A dummy type that describes the kind of borrow and carries a lifetime.
 ///    - When this is `Immut<'a>`, the `NodeRef` acts roughly like `&'a Node`.
 ///    - When this is `ValMut<'a>`, the `NodeRef` acts roughly like `&'a Node`
@@ -226,33 +227,27 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
 
     fn from_new_leaf<A: Allocator + Clone>(leaf: Box<LeafNode<K, V>, A>) -> Self {
         // The allocator must be dropped, not leaked.  See also `BTreeMap::alloc`.
-        let (leaf, _alloc) = Box::into_raw_with_allocator(leaf);
-        // SAFETY: the node was just allocated.
-        let node = unsafe { NonNull::new_unchecked(leaf) };
+        let (node, _alloc) = Box::into_non_null_with_allocator(leaf);
         NodeRef { height: 0, node, _marker: PhantomData }
     }
 }
 
 impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
+    /// Creates a new internal (height > 0) `NodeRef`
     fn new_internal<A: Allocator + Clone>(child: Root<K, V>, alloc: A) -> Self {
         let mut new_node = unsafe { InternalNode::new(alloc) };
         new_node.edges[0].write(child.node);
-        unsafe { NodeRef::from_new_internal(new_node, child.height + 1) }
+        NodeRef::from_new_internal(new_node, NonZero::new(child.height + 1).unwrap())
     }
 
-    /// # Safety
-    /// `height` must not be zero.
-    unsafe fn from_new_internal<A: Allocator + Clone>(
+    /// Creates a new internal (height > 0) `NodeRef` from an existing internal node
+    fn from_new_internal<A: Allocator + Clone>(
         internal: Box<InternalNode<K, V>, A>,
-        height: usize,
+        height: NonZero<usize>,
     ) -> Self {
-        debug_assert!(height > 0);
         // The allocator must be dropped, not leaked.  See also `BTreeMap::alloc`.
-        let (internal, _alloc) = Box::into_raw_with_allocator(internal);
-        // SAFETY: the node was just allocated.
-        let internal = unsafe { NonNull::new_unchecked(internal) };
-        let node = internal.cast();
-        let mut this = NodeRef { height, node, _marker: PhantomData };
+        let (node, _alloc) = Box::into_non_null_with_allocator(internal);
+        let mut this = NodeRef { height: height.into(), node: node.cast(), _marker: PhantomData };
         this.borrow_mut().correct_all_childrens_parent_links();
         this
     }
@@ -625,9 +620,8 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
         let top = self.node;
 
         // SAFETY: we asserted to be internal.
-        let internal_self = unsafe { self.borrow_mut().cast_to_internal_unchecked() };
-        // SAFETY: we borrowed `self` exclusively and its borrow type is exclusive.
-        let internal_node = unsafe { &mut *NodeRef::as_internal_ptr(&internal_self) };
+        let mut internal_self = unsafe { self.borrow_mut().cast_to_internal_unchecked() };
+        let internal_node = internal_self.as_internal_mut();
         // SAFETY: the first edge is always initialized.
         self.node = unsafe { internal_node.edges[0].assume_init_read() };
         self.height -= 1;
@@ -1305,7 +1299,8 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>,
                 &mut new_node.edges[..new_len + 1],
             );
 
-            let height = self.node.height;
+            // SAFETY: self is `marker::Internal`, so `self.node.height` is positive
+            let height = NonZero::new_unchecked(self.node.height);
             let right = NodeRef::from_new_internal(new_node, height);
 
             SplitResult { left: self.node, kv, right }