Push new allocator trait through BTreeMap

The `BTreeMap` implementation was updated with support for custom
allocators, but many of the methods that take allocators do not require
that the provided allocator is the same allocator that is used for the
owning `BTreeMap`. This could lead to situations where nodes are created
with a different allocator, which would violate the safety conditions of
`deallocate` on drop. This was discovered because the changes to
`Box::leak` make it invalid to call without upgrading the allocator to
`A: PinSafeAllocator`.

This also updates `Box::pin_in` with the new `PinSafeAllocator` trait.

Author: djkoloski

library/alloc/src/boxed.rs

@@ -575,7 +575,7 @@ impl<T, A: Allocator> Box<T, A> {
     #[inline(always)]
     pub const fn pin_in(x: T, alloc: A) -> Pin<Self>
     where
-        A: 'static + ~const Allocator + ~const Destruct,
+        A: ~const Allocator + ~const PinSafeAllocator + ~const Destruct,
     {
         Self::into_pin(Self::new_in(x, alloc))
     }

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

@@ -15,7 +15,11 @@ impl<K, V> Root<K, V> {
     /// a `BTreeMap`, both iterators should produce keys in strictly ascending
     /// order, each greater than all keys in the tree, including any keys
     /// already in the tree upon entry.
-    pub fn append_from_sorted_iters<I, A: Allocator + Clone>(
+    ///
+    /// # Safety
+    ///
+    /// `alloc` must be the allocator for the owning `BTreeMap`.
+    pub unsafe fn append_from_sorted_iters<I, A: Allocator + Clone>(
         &mut self,
         left: I,
         right: I,
@@ -29,14 +33,25 @@ impl<K, V> Root<K, V> {
         let iter = MergeIter(MergeIterInner::new(left, right));
 
         // Meanwhile, we build a tree from the sorted sequence in linear time.
-        self.bulk_push(iter, length, alloc)
+
+        // SAFETY: The caller has guaranteed that `alloc` is the allocator for the owning
+        // `BTreeMap`.
+        unsafe { self.bulk_push(iter, length, alloc) }
     }
 
     /// Pushes all key-value pairs to the end of the tree, incrementing a
     /// `length` variable along the way. The latter makes it easier for the
     /// caller to avoid a leak when the iterator panicks.
-    pub fn bulk_push<I, A: Allocator + Clone>(&mut self, iter: I, length: &mut usize, alloc: A)
-    where
+    ///
+    /// # Safety
+    ///
+    /// `alloc` must be the allocator for the owning `BTreeMap`.
+    pub unsafe fn bulk_push<I, A: Allocator + Clone>(
+        &mut self,
+        iter: I,
+        length: &mut usize,
+        alloc: A,
+    ) where
         I: Iterator<Item = (K, V)>,
     {
         let mut cur_node = self.borrow_mut().last_leaf_edge().into_node();
@@ -64,17 +79,24 @@ impl<K, V> Root<K, V> {
                         }
                         Err(_) => {
                             // We are at the top, create a new root node and push there.
-                            open_node = self.push_internal_level(alloc.clone());
+
+                            // SAFETY: The caller has guaranteed that `alloc` is the allocator for
+                            // the owning `BTreeMap`.
+                            open_node = unsafe { self.push_internal_level(alloc.clone()) };
                             break;
                         }
                     }
                 }
 
                 // Push key-value pair and new right subtree.
                 let tree_height = open_node.height() - 1;
-                let mut right_tree = Root::new(alloc.clone());
+                // SAFETY: The caller has guaranteed that `alloc` is the allocator for the owning
+                // `BTreeMap`.
+                let mut right_tree = unsafe { Root::new(alloc.clone()) };
                 for _ in 0..tree_height {
-                    right_tree.push_internal_level(alloc.clone());
+                    // SAFETY: The caller has guaranteed that `alloc` is the allocator for the
+                    // owning `BTreeMap`.
+                    unsafe { right_tree.push_internal_level(alloc.clone()) };
                 }
                 open_node.push(key, value, right_tree);
 

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

@@ -216,7 +216,9 @@ impl<K: Clone, V: Clone, A: Allocator + Clone> Clone for BTreeMap<K, V, A> {
             match node.force() {
                 Leaf(leaf) => {
                     let mut out_tree = BTreeMap {
-                        root: Some(Root::new(alloc.clone())),
+                        // SAFETY: `alloc` is the allocator for both the original and the cloned
+                        // `BTreeMap`.
+                        root: unsafe { Some(Root::new(alloc.clone())) },
                         length: 0,
                         alloc: ManuallyDrop::new(alloc),
                         _marker: PhantomData,
@@ -247,7 +249,9 @@ impl<K: Clone, V: Clone, A: Allocator + Clone> Clone for BTreeMap<K, V, A> {
 
                     {
                         let out_root = out_tree.root.as_mut().unwrap();
-                        let mut out_node = out_root.push_internal_level(alloc.clone());
+                        // SAFETY: `alloc` is the allocator for both the original and the cloned
+                        // `BTreeMap`.
+                        let mut out_node = unsafe { out_root.push_internal_level(alloc.clone()) };
                         let mut in_edge = internal.first_edge();
                         while let Ok(kv) = in_edge.right_kv() {
                             let (k, v) = kv.into_kv();
@@ -269,7 +273,9 @@ impl<K: Clone, V: Clone, A: Allocator + Clone> Clone for BTreeMap<K, V, A> {
                             out_node.push(
                                 k,
                                 v,
-                                subroot.unwrap_or_else(|| Root::new(alloc.clone())),
+                                // SAFETY: `alloc` is the allocator for both the original and cloned
+                                // `BTreeMap`.
+                                subroot.unwrap_or_else(|| unsafe { Root::new(alloc.clone()) }),
                             );
                             out_tree.length += 1 + sublength;
                         }
@@ -323,8 +329,9 @@ where
 
     fn replace(&mut self, key: K) -> Option<K> {
         let (map, dormant_map) = DormantMutRef::new(self);
+        // SAFETY: `alloc` is the allocator for the `BTreeMap`.
         let root_node =
-            map.root.get_or_insert_with(|| Root::new((*map.alloc).clone())).borrow_mut();
+            map.root.get_or_insert_with(|| unsafe { Root::new((*map.alloc).clone()) }).borrow_mut();
         match root_node.search_tree::<K>(&key) {
             Found(mut kv) => Some(mem::replace(kv.key_mut(), key)),
             GoDown(handle) => {
@@ -1144,13 +1151,16 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
 
         let self_iter = mem::replace(self, Self::new_in((*self.alloc).clone())).into_iter();
         let other_iter = mem::replace(other, Self::new_in((*self.alloc).clone())).into_iter();
-        let root = self.root.get_or_insert_with(|| Root::new((*self.alloc).clone()));
-        root.append_from_sorted_iters(
-            self_iter,
-            other_iter,
-            &mut self.length,
-            (*self.alloc).clone(),
-        )
+        let root = self.root.get_or_insert_with(|| unsafe { Root::new((*self.alloc).clone()) });
+        // SAFETY: `self.alloc` is the allocator for the `BTreeMap`.
+        unsafe {
+            root.append_from_sorted_iters(
+                self_iter,
+                other_iter,
+                &mut self.length,
+                (*self.alloc).clone(),
+            )
+        }
     }
 
     /// Constructs a double-ended iterator over a sub-range of elements in the map.
@@ -1464,9 +1474,13 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
         K: Ord,
         I: IntoIterator<Item = (K, V)>,
     {
-        let mut root = Root::new(alloc.clone());
+        // SAFETY: `alloc` is the allocator for the returned `BTreeMap`.
+        let mut root = unsafe { Root::new(alloc.clone()) };
         let mut length = 0;
-        root.bulk_push(DedupSortedIter::new(iter.into_iter()), &mut length, alloc.clone());
+        // SAFETY: `alloc` is the allocator for the returned `BTreeMap`.
+        unsafe {
+            root.bulk_push(DedupSortedIter::new(iter.into_iter()), &mut length, alloc.clone());
+        }
         BTreeMap { root: Some(root), length, alloc: ManuallyDrop::new(alloc), _marker: PhantomData }
     }
 }

library/alloc/src/collections/btree/map/entry.rs

@@ -342,30 +342,39 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
             None => {
                 // SAFETY: There is no tree yet so no reference to it exists.
                 let map = unsafe { self.dormant_map.awaken() };
-                let mut root = NodeRef::new_leaf(self.alloc.clone());
+                // SAFETY: `self.alloc` is the allocator for the owning `BTreeMap`.
+                let mut root = unsafe { NodeRef::new_leaf(self.alloc.clone()) };
                 let val_ptr = root.borrow_mut().push(self.key, value) as *mut V;
                 map.root = Some(root.forget_type());
                 map.length = 1;
                 val_ptr
             }
-            Some(handle) => match handle.insert_recursing(self.key, value, self.alloc.clone()) {
-                (None, val_ptr) => {
-                    // SAFETY: We have consumed self.handle.
-                    let map = unsafe { self.dormant_map.awaken() };
-                    map.length += 1;
-                    val_ptr
+            Some(handle) => {
+                // SAFETY: `self.alloc` is the allocator for the owning `BTreeMap`.
+                let insert_result =
+                    unsafe { handle.insert_recursing(self.key, value, self.alloc.clone()) };
+                match insert_result {
+                    (None, val_ptr) => {
+                        // SAFETY: We have consumed self.handle.
+                        let map = unsafe { self.dormant_map.awaken() };
+                        map.length += 1;
+                        val_ptr
+                    }
+                    (Some(ins), val_ptr) => {
+                        drop(ins.left);
+                        // SAFETY: We have consumed self.handle and dropped the
+                        // remaining reference to the tree, ins.left.
+                        let map = unsafe { self.dormant_map.awaken() };
+                        let root = map.root.as_mut().unwrap(); // same as ins.left
+                        // SAFETY: `self.alloc` is the allocator for the owning `BTreeMap`.
+                        unsafe {
+                            root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right)
+                        };
+                        map.length += 1;
+                        val_ptr
+                    }
                 }
-                (Some(ins), val_ptr) => {
-                    drop(ins.left);
-                    // SAFETY: We have consumed self.handle and dropped the
-                    // remaining reference to the tree, ins.left.
-                    let map = unsafe { self.dormant_map.awaken() };
-                    let root = map.root.as_mut().unwrap(); // same as ins.left
-                    root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right);
-                    map.length += 1;
-                    val_ptr
-                }
-            },
+            }
         };
         // Now that we have finished growing the tree using borrowed references,
         // dereference the pointer to a part of it, that we picked up along the way.

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

@@ -116,7 +116,14 @@ impl<K, V> BTreeMap<K, V> {
     {
         let iter = mem::take(self).into_iter();
         if !iter.is_empty() {
-            self.root.insert(Root::new(*self.alloc)).bulk_push(iter, &mut self.length, *self.alloc);
+            // SAFETY: `self.alloc` is the allocator for this `BTreeMap`.
+            unsafe {
+                self.root.insert(Root::new(*self.alloc)).bulk_push(
+                    iter,
+                    &mut self.length,
+                    *self.alloc,
+                );
+            }
         }
     }
 }