about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--library/alloc/src/collections/btree/map.rs3
-rw-r--r--library/alloc/src/collections/btree/node.rs12
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 cebd25c6039..e3b53b2fe29 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 b233e1740b7..77bd90f2c75 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