From 137d4eaf12139af894da6c0b6fe74195d6b7b98c Mon Sep 17 00:00:00 2001
From: Sidney Cammeresi <sac@cheesecake.org>
Date: Sat, 20 Sep 2025 16:54:51 -0700
Subject: [PATCH] BTreeMap: Don't leak allocators when initializing nodes

Memory was allocated via `Box::leak` and thence intended to be tracked
and deallocated manually, but the allocator was also leaked, not
tracked, and never dropped.  Now it is dropped immediately.

According to my reading of the `Allocator` trait, if a copy of the
allocator remains live, then its allocations must remain live.  Since
the B-tree has a copy of the allocator that will only be dropped after
the nodes, it's safe to not store the allocator in each node (which
would be a much more intrusive change).
---
 library/alloc/src/collections/btree/map.rs  |  3 +++
 library/alloc/src/collections/btree/node.rs | 12 ++++++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index cebd25c6039cc..e3b53b2fe291a 100644
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -194,6 +194,9 @@ pub struct BTreeMap<
     root: Option<Root<K, V>>,
     length: usize,
     /// `ManuallyDrop` to control drop order (needs to be dropped after all the nodes).
+    // Although some of the accessory types store a copy of the allocator, the nodes do not.
+    // Because allocations will remain live as long as any copy (like this one) of the allocator
+    // is live, it's unnecessary to store the allocator in each node.
     pub(super) alloc: ManuallyDrop<A>,
     // For dropck; the `Box` avoids making the `Unpin` impl more strict than before
     _marker: PhantomData<crate::boxed::Box<(K, V), A>>,
diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
index b233e1740b7b7..77bd90f2c750c 100644
--- a/library/alloc/src/collections/btree/node.rs
+++ b/library/alloc/src/collections/btree/node.rs
@@ -224,7 +224,11 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
     }
 
     fn from_new_leaf<A: Allocator + Clone>(leaf: Box<LeafNode<K, V>, A>) -> Self {
-        NodeRef { height: 0, node: NonNull::from(Box::leak(leaf)), _marker: PhantomData }
+        // 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) };
+        NodeRef { height: 0, node, _marker: PhantomData }
     }
 }
 
@@ -242,7 +246,11 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
         height: usize,
     ) -> Self {
         debug_assert!(height > 0);
-        let node = NonNull::from(Box::leak(internal)).cast();
+        // 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 };
         this.borrow_mut().correct_all_childrens_parent_links();
         this