Skip to content

Commit 483c8a3

Browse files
djkoloskidjkoloski
committed
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.
1 parent fe3d3fc commit 483c8a3

File tree

12 files changed

+203
-75
lines changed

12 files changed

+203
-75
lines changed

‎library/alloc/src/boxed.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -575,7 +575,7 @@ impl<T, A: Allocator> Box<T, A> {
575575
#[inline(always)]
576576
pub const fn pin_in(x: T, alloc: A) -> Pin<Self>
577577
where
578-
A: 'static + ~const Allocator + ~const Destruct,
578+
A: ~const Allocator + ~const PinSafeAllocator + ~const Destruct,
579579
{
580580
Self::into_pin(Self::new_in(x, alloc))
581581
}

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

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,11 @@ impl<K, V> Root<K, V> {
1515
/// a `BTreeMap`, both iterators should produce keys in strictly ascending
1616
/// order, each greater than all keys in the tree, including any keys
1717
/// already in the tree upon entry.
18-
pub fn append_from_sorted_iters<I, A: Allocator + Clone>(
18+
///
19+
/// # Safety
20+
///
21+
/// `alloc` must be the allocator for the owning `BTreeMap`.
22+
pub unsafe fn append_from_sorted_iters<I, A: Allocator + Clone>(
1923
&mut self,
2024
left: I,
2125
right: I,
@@ -29,14 +33,25 @@ impl<K, V> Root<K, V> {
2933
let iter = MergeIter(MergeIterInner::new(left, right));
3034

3135
// Meanwhile, we build a tree from the sorted sequence in linear time.
32-
self.bulk_push(iter, length, alloc)
36+
37+
// SAFETY: The caller has guaranteed that `alloc` is the allocator for the owning
38+
// `BTreeMap`.
39+
unsafe { self.bulk_push(iter, length, alloc) }
3340
}
3441

3542
/// Pushes all key-value pairs to the end of the tree, incrementing a
3643
/// `length` variable along the way. The latter makes it easier for the
3744
/// caller to avoid a leak when the iterator panicks.
38-
pub fn bulk_push<I, A: Allocator + Clone>(&mut self, iter: I, length: &mut usize, alloc: A)
39-
where
45+
///
46+
/// # Safety
47+
///
48+
/// `alloc` must be the allocator for the owning `BTreeMap`.
49+
pub unsafe fn bulk_push<I, A: Allocator + Clone>(
50+
&mut self,
51+
iter: I,
52+
length: &mut usize,
53+
alloc: A,
54+
) where
4055
I: Iterator<Item = (K, V)>,
4156
{
4257
let mut cur_node = self.borrow_mut().last_leaf_edge().into_node();
@@ -64,17 +79,24 @@ impl<K, V> Root<K, V> {
6479
}
6580
Err(_) => {
6681
// We are at the top, create a new root node and push there.
67-
open_node = self.push_internal_level(alloc.clone());
82+
83+
// SAFETY: The caller has guaranteed that `alloc` is the allocator for
84+
// the owning `BTreeMap`.
85+
open_node = unsafe { self.push_internal_level(alloc.clone()) };
6886
break;
6987
}
7088
}
7189
}
7290

7391
// Push key-value pair and new right subtree.
7492
let tree_height = open_node.height() - 1;
75-
let mut right_tree = Root::new(alloc.clone());
93+
// SAFETY: The caller has guaranteed that `alloc` is the allocator for the owning
94+
// `BTreeMap`.
95+
let mut right_tree = unsafe { Root::new(alloc.clone()) };
7696
for _ in 0..tree_height {
77-
right_tree.push_internal_level(alloc.clone());
97+
// SAFETY: The caller has guaranteed that `alloc` is the allocator for the
98+
// owning `BTreeMap`.
99+
unsafe { right_tree.push_internal_level(alloc.clone()) };
78100
}
79101
open_node.push(key, value, right_tree);
80102

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

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,9 @@ impl<K: Clone, V: Clone, A: Allocator + Clone> Clone for BTreeMap<K, V, A> {
216216
match node.force() {
217217
Leaf(leaf) => {
218218
let mut out_tree = BTreeMap {
219-
root: Some(Root::new(alloc.clone())),
219+
// SAFETY: `alloc` is the allocator for both the original and the cloned
220+
// `BTreeMap`.
221+
root: unsafe { Some(Root::new(alloc.clone())) },
220222
length: 0,
221223
alloc: ManuallyDrop::new(alloc),
222224
_marker: PhantomData,
@@ -247,7 +249,9 @@ impl<K: Clone, V: Clone, A: Allocator + Clone> Clone for BTreeMap<K, V, A> {
247249

248250
{
249251
let out_root = out_tree.root.as_mut().unwrap();
250-
let mut out_node = out_root.push_internal_level(alloc.clone());
252+
// SAFETY: `alloc` is the allocator for both the original and the cloned
253+
// `BTreeMap`.
254+
let mut out_node = unsafe { out_root.push_internal_level(alloc.clone()) };
251255
let mut in_edge = internal.first_edge();
252256
while let Ok(kv) = in_edge.right_kv() {
253257
let (k, v) = kv.into_kv();
@@ -269,7 +273,9 @@ impl<K: Clone, V: Clone, A: Allocator + Clone> Clone for BTreeMap<K, V, A> {
269273
out_node.push(
270274
k,
271275
v,
272-
subroot.unwrap_or_else(|| Root::new(alloc.clone())),
276+
// SAFETY: `alloc` is the allocator for both the original and cloned
277+
// `BTreeMap`.
278+
subroot.unwrap_or_else(|| unsafe { Root::new(alloc.clone()) }),
273279
);
274280
out_tree.length += 1 + sublength;
275281
}
@@ -323,8 +329,9 @@ where
323329

324330
fn replace(&mut self, key: K) -> Option<K> {
325331
let (map, dormant_map) = DormantMutRef::new(self);
332+
// SAFETY: `alloc` is the allocator for the `BTreeMap`.
326333
let root_node =
327-
map.root.get_or_insert_with(|| Root::new((*map.alloc).clone())).borrow_mut();
334+
map.root.get_or_insert_with(|| unsafe { Root::new((*map.alloc).clone()) }).borrow_mut();
328335
match root_node.search_tree::<K>(&key) {
329336
Found(mut kv) => Some(mem::replace(kv.key_mut(), key)),
330337
GoDown(handle) => {
@@ -1144,13 +1151,16 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
11441151

11451152
let self_iter = mem::replace(self, Self::new_in((*self.alloc).clone())).into_iter();
11461153
let other_iter = mem::replace(other, Self::new_in((*self.alloc).clone())).into_iter();
1147-
let root = self.root.get_or_insert_with(|| Root::new((*self.alloc).clone()));
1148-
root.append_from_sorted_iters(
1149-
self_iter,
1150-
other_iter,
1151-
&mut self.length,
1152-
(*self.alloc).clone(),
1153-
)
1154+
let root = self.root.get_or_insert_with(|| unsafe { Root::new((*self.alloc).clone()) });
1155+
// SAFETY: `self.alloc` is the allocator for the `BTreeMap`.
1156+
unsafe {
1157+
root.append_from_sorted_iters(
1158+
self_iter,
1159+
other_iter,
1160+
&mut self.length,
1161+
(*self.alloc).clone(),
1162+
)
1163+
}
11541164
}
11551165

11561166
/// 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> {
14641474
K: Ord,
14651475
I: IntoIterator<Item = (K, V)>,
14661476
{
1467-
let mut root = Root::new(alloc.clone());
1477+
// SAFETY: `alloc` is the allocator for the returned `BTreeMap`.
1478+
let mut root = unsafe { Root::new(alloc.clone()) };
14681479
let mut length = 0;
1469-
root.bulk_push(DedupSortedIter::new(iter.into_iter()), &mut length, alloc.clone());
1480+
// SAFETY: `alloc` is the allocator for the returned `BTreeMap`.
1481+
unsafe {
1482+
root.bulk_push(DedupSortedIter::new(iter.into_iter()), &mut length, alloc.clone());
1483+
}
14701484
BTreeMap { root: Some(root), length, alloc: ManuallyDrop::new(alloc), _marker: PhantomData }
14711485
}
14721486
}

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

Lines changed: 27 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -342,30 +342,39 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
342342
None => {
343343
// SAFETY: There is no tree yet so no reference to it exists.
344344
let map = unsafe { self.dormant_map.awaken() };
345-
let mut root = NodeRef::new_leaf(self.alloc.clone());
345+
// SAFETY: `self.alloc` is the allocator for the owning `BTreeMap`.
346+
let mut root = unsafe { NodeRef::new_leaf(self.alloc.clone()) };
346347
let val_ptr = root.borrow_mut().push(self.key, value) as *mut V;
347348
map.root = Some(root.forget_type());
348349
map.length = 1;
349350
val_ptr
350351
}
351-
Some(handle) => match handle.insert_recursing(self.key, value, self.alloc.clone()) {
352-
(None, val_ptr) => {
353-
// SAFETY: We have consumed self.handle.
354-
let map = unsafe { self.dormant_map.awaken() };
355-
map.length += 1;
356-
val_ptr
352+
Some(handle) => {
353+
// SAFETY: `self.alloc` is the allocator for the owning `BTreeMap`.
354+
let insert_result =
355+
unsafe { handle.insert_recursing(self.key, value, self.alloc.clone()) };
356+
match insert_result {
357+
(None, val_ptr) => {
358+
// SAFETY: We have consumed self.handle.
359+
let map = unsafe { self.dormant_map.awaken() };
360+
map.length += 1;
361+
val_ptr
362+
}
363+
(Some(ins), val_ptr) => {
364+
drop(ins.left);
365+
// SAFETY: We have consumed self.handle and dropped the
366+
// remaining reference to the tree, ins.left.
367+
let map = unsafe { self.dormant_map.awaken() };
368+
let root = map.root.as_mut().unwrap(); // same as ins.left
369+
// SAFETY: `self.alloc` is the allocator for the owning `BTreeMap`.
370+
unsafe {
371+
root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right)
372+
};
373+
map.length += 1;
374+
val_ptr
375+
}
357376
}
358-
(Some(ins), val_ptr) => {
359-
drop(ins.left);
360-
// SAFETY: We have consumed self.handle and dropped the
361-
// remaining reference to the tree, ins.left.
362-
let map = unsafe { self.dormant_map.awaken() };
363-
let root = map.root.as_mut().unwrap(); // same as ins.left
364-
root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right);
365-
map.length += 1;
366-
val_ptr
367-
}
368-
},
377+
}
369378
};
370379
// Now that we have finished growing the tree using borrowed references,
371380
// dereference the pointer to a part of it, that we picked up along the way.

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,14 @@ impl<K, V> BTreeMap<K, V> {
116116
{
117117
let iter = mem::take(self).into_iter();
118118
if !iter.is_empty() {
119-
self.root.insert(Root::new(*self.alloc)).bulk_push(iter, &mut self.length, *self.alloc);
119+
// SAFETY: `self.alloc` is the allocator for this `BTreeMap`.
120+
unsafe {
121+
self.root.insert(Root::new(*self.alloc)).bulk_push(
122+
iter,
123+
&mut self.length,
124+
*self.alloc,
125+
);
126+
}
120127
}
121128
}
122129
}

0 commit comments

Comments
 (0)