Skip to content

Commit e2242bb

Browse files
committed
BTreeMap: postpone some DrainFilter processing to drop handler
1 parent 577a229 commit e2242bb

File tree

4 files changed

+196
-43
lines changed

4 files changed

+196
-43
lines changed

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

+140-38
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
// ignore-tidy-filelength
2+
13
use core::borrow::Borrow;
24
use core::cmp::Ordering;
35
use core::fmt::{self, Debug};
@@ -1277,16 +1279,28 @@ impl<K: Ord, V> BTreeMap<K, V> {
12771279
}
12781280

12791281
pub(super) fn drain_filter_inner(&mut self) -> DrainFilterInner<'_, K, V> {
1280-
if let Some(root) = self.root.as_mut() {
1281-
let (root, dormant_root) = DormantMutRef::new(root);
1282-
let front = root.node_as_mut().first_leaf_edge();
1282+
let (map, dormant_map) = DormantMutRef::new(self);
1283+
// Leak amplification: wipe root and length in case caller skips drop handler.
1284+
if let Some(root) = map.root.take() {
1285+
let length = mem::replace(&mut map.length, 0);
1286+
let height = root.height();
1287+
1288+
let front = root.into_node_mut().first_leaf_edge();
12831289
DrainFilterInner {
1284-
length: &mut self.length,
1285-
dormant_root: Some(dormant_root),
1286-
cur_leaf_edge: Some(front),
1290+
dormant_map: Some(dormant_map),
1291+
position: Some(DrainFilterPosition::Edge(front)),
1292+
height,
1293+
length,
1294+
remaining: length,
12871295
}
12881296
} else {
1289-
DrainFilterInner { length: &mut self.length, dormant_root: None, cur_leaf_edge: None }
1297+
DrainFilterInner {
1298+
dormant_map: None,
1299+
position: None,
1300+
height: 0,
1301+
length: 0,
1302+
remaining: 0,
1303+
}
12901304
}
12911305
}
12921306

@@ -1671,11 +1685,21 @@ where
16711685
/// Most of the implementation of DrainFilter, independent of the type
16721686
/// of the predicate, thus also serving for BTreeSet::DrainFilter.
16731687
pub(super) struct DrainFilterInner<'a, K: 'a, V: 'a> {
1674-
length: &'a mut usize,
1675-
// dormant_root is wrapped in an Option to be able to `take` it.
1676-
dormant_root: Option<DormantMutRef<'a, node::Root<K, V>>>,
1677-
// cur_leaf_edge is wrapped in an Option because maps without root lack a leaf edge.
1678-
cur_leaf_edge: Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>,
1688+
// wrapped in an Option to be able to `take` it in the drop handler.
1689+
dormant_map: Option<DormantMutRef<'a, BTreeMap<K, V>>>,
1690+
// wrapped in an Option some maps don't have any:
1691+
position: Option<DrainFilterPosition<'a, K, V>>,
1692+
height: usize,
1693+
length: usize,
1694+
remaining: usize,
1695+
}
1696+
enum DrainFilterPosition<'a, K: 'a, V: 'a> {
1697+
// Initial and normal position on a leaf edge.
1698+
Edge(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>),
1699+
// Intermediate position on a KV, that sticks around if a predicate panics.
1700+
KV(Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>),
1701+
// End position on the root node.
1702+
Root(NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>),
16791703
}
16801704

16811705
#[unstable(feature = "btree_drain_filter", issue = "70530")]
@@ -1716,40 +1740,105 @@ where
17161740
}
17171741
}
17181742

1743+
impl<K, V> Drop for DrainFilterInner<'_, K, V> {
1744+
fn drop(&mut self) {
1745+
if let Some(dormant_map) = self.dormant_map.take() {
1746+
let root_node = match self.position.take() {
1747+
None => unreachable!(),
1748+
Some(DrainFilterPosition::Edge(edge)) => {
1749+
// panic during drop
1750+
edge.into_node().forget_type().root_node()
1751+
}
1752+
Some(DrainFilterPosition::KV(kv)) => {
1753+
// panic in predicate
1754+
kv.into_node().root_node()
1755+
}
1756+
Some(DrainFilterPosition::Root(root_node)) => root_node,
1757+
};
1758+
let mut root = node::Root::restore_from_node(root_node);
1759+
for _ in self.height..root.height() {
1760+
root.pop_internal_level();
1761+
}
1762+
debug_assert_eq!(self.height, root.height());
1763+
1764+
let map = unsafe { dormant_map.awaken() };
1765+
map.root = Some(root);
1766+
map.length = self.length;
1767+
}
1768+
}
1769+
}
1770+
17191771
impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
17201772
/// Allow Debug implementations to predict the next element.
17211773
pub(super) fn peek(&self) -> Option<(&K, &V)> {
1722-
let edge = self.cur_leaf_edge.as_ref()?;
1723-
edge.reborrow().next_kv().ok().map(|kv| kv.into_kv())
1774+
if let Some(DrainFilterPosition::Edge(edge)) = self.position.as_ref() {
1775+
edge.reborrow().next_kv().ok().map(|kv| kv.into_kv())
1776+
} else {
1777+
None
1778+
}
17241779
}
17251780

17261781
/// Implementation of a typical `DrainFilter::next` method, given the predicate.
17271782
pub(super) fn next<F>(&mut self, pred: &mut F) -> Option<(K, V)>
17281783
where
17291784
F: FnMut(&K, &mut V) -> bool,
17301785
{
1731-
while let Ok(mut kv) = self.cur_leaf_edge.take()?.next_kv() {
1732-
let (k, v) = kv.kv_mut();
1733-
if pred(k, v) {
1734-
*self.length -= 1;
1735-
let (kv, pos) = kv.remove_kv_tracking(|| {
1736-
// SAFETY: we will touch the root in a way that will not
1737-
// invalidate the position returned.
1738-
let root = unsafe { self.dormant_root.take().unwrap().awaken() };
1739-
root.pop_internal_level();
1740-
self.dormant_root = Some(DormantMutRef::new(root).1);
1741-
});
1742-
self.cur_leaf_edge = Some(pos);
1743-
return Some(kv);
1786+
let mut cur_leaf_edge = match self.position.take() {
1787+
None => return None,
1788+
Some(DrainFilterPosition::Root(root_node)) => {
1789+
self.position = Some(DrainFilterPosition::Root(root_node));
1790+
return None;
1791+
}
1792+
Some(DrainFilterPosition::KV(kv)) => {
1793+
self.position = Some(DrainFilterPosition::KV(kv));
1794+
return None;
1795+
}
1796+
Some(DrainFilterPosition::Edge(edge)) => edge,
1797+
};
1798+
1799+
loop {
1800+
match cur_leaf_edge.next_kv() {
1801+
Err(root_node) => {
1802+
self.position = Some(DrainFilterPosition::Root(root_node));
1803+
return None;
1804+
}
1805+
Ok(kv) => {
1806+
self.remaining -= 1;
1807+
// Store the intermediate position in case the predicate panics.
1808+
self.position = Some(DrainFilterPosition::KV(kv));
1809+
let drain = self
1810+
.position
1811+
.as_mut()
1812+
.map(|pos| match pos {
1813+
DrainFilterPosition::KV(kv) => {
1814+
let (k, v) = kv.kv_mut();
1815+
pred(k, v)
1816+
}
1817+
_ => unreachable!(),
1818+
})
1819+
.unwrap();
1820+
let kv = if let Some(DrainFilterPosition::KV(kv)) = self.position.take() {
1821+
kv
1822+
} else {
1823+
unreachable!()
1824+
};
1825+
if drain {
1826+
self.length -= 1;
1827+
let (kv, pos) = kv.remove_kv_tracking(|emptied_internal_node| {
1828+
self.height = emptied_internal_node.height() - 1
1829+
});
1830+
self.position = Some(DrainFilterPosition::Edge(pos));
1831+
return Some(kv);
1832+
}
1833+
cur_leaf_edge = kv.next_leaf_edge();
1834+
}
17441835
}
1745-
self.cur_leaf_edge = Some(kv.next_leaf_edge());
17461836
}
1747-
None
17481837
}
17491838

17501839
/// Implementation of a typical `DrainFilter::size_hint` method.
17511840
pub(super) fn size_hint(&self) -> (usize, Option<usize>) {
1752-
(0, Some(*self.length))
1841+
(0, Some(self.remaining))
17531842
}
17541843
}
17551844

@@ -2767,7 +2856,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
27672856
// Body of `remove_entry`, separate to keep the above implementations short.
27682857
fn remove_kv(self) -> (K, V) {
27692858
let mut emptied_internal_root = false;
2770-
let (old_kv, _) = self.handle.remove_kv_tracking(|| emptied_internal_root = true);
2859+
let (old_kv, _) = self.handle.remove_kv_tracking(|_| emptied_internal_root = true);
27712860
// SAFETY: we consumed the intermediate root borrow, `self.handle`.
27722861
let map = unsafe { self.dormant_map.awaken() };
27732862
map.length -= 1;
@@ -2782,10 +2871,13 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
27822871
impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
27832872
/// Removes a key/value-pair from the map, and returns that pair, as well as
27842873
/// the leaf edge corresponding to that former pair.
2785-
fn remove_kv_tracking<F: FnOnce()>(
2874+
fn remove_kv_tracking<F>(
27862875
self,
2787-
handle_emptied_internal_root: F,
2788-
) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
2876+
handle_emptied_parent: F,
2877+
) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>)
2878+
where
2879+
F: FnOnce(NodeRef<marker::Mut<'a>, K, V, marker::Internal>),
2880+
{
27892881
let (old_kv, mut pos, was_internal) = match self.force() {
27902882
Leaf(leaf) => {
27912883
let (old_kv, pos) = leaf.remove();
@@ -2836,8 +2928,10 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInter
28362928
if parent.len() == 0 {
28372929
// The parent that was just emptied must be the root,
28382930
// because nodes on a lower level would not have been
2839-
// left with a single child.
2840-
handle_emptied_internal_root();
2931+
// left with a single child. Or, during `DrainFilter`
2932+
// use, its ancestors have been emptied earlier.
2933+
// The parent and its ancestors should be popped off.
2934+
handle_emptied_parent(parent);
28412935
break;
28422936
} else {
28432937
cur_node = parent.forget_type();
@@ -2948,8 +3042,16 @@ fn handle_underfull_node<K, V>(
29483042
let (is_left, mut handle) = match parent.left_kv() {
29493043
Ok(left) => (true, left),
29503044
Err(parent) => {
2951-
let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
2952-
(false, right)
3045+
match parent.right_kv() {
3046+
Ok(right) => (false, right),
3047+
Err(_) => {
3048+
// The underfull node has an empty parent, so it is the only child
3049+
// of an empty root. It is destined to become the new root, thus
3050+
// allowed to be underfull. The empty parent should be removed later
3051+
// by `pop_internal_level`.
3052+
return AtRoot;
3053+
}
3054+
}
29533055
}
29543056
};
29553057

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

+31
Original file line numberDiff line numberDiff line change
@@ -1147,6 +1147,37 @@ mod test_drain_filter {
11471147
assert_eq!(map.last_entry().unwrap().key(), &8);
11481148
map.check();
11491149
}
1150+
1151+
#[test]
1152+
#[cfg_attr(miri, ignore)] // Miri reports the expected leak
1153+
fn forget_height_1() {
1154+
static DROPS: AtomicUsize = AtomicUsize::new(0);
1155+
1156+
struct D;
1157+
impl Drop for D {
1158+
fn drop(&mut self) {
1159+
DROPS.fetch_add(1, Ordering::SeqCst);
1160+
}
1161+
}
1162+
1163+
let mut map: BTreeMap<_, _> = (0..MIN_INSERTS_HEIGHT_1).map(|i| (i, D)).collect();
1164+
map.remove(&0); // leaving 2 levels on the verge of collapse, 5 + 1 + 5 elements.
1165+
map.check();
1166+
assert_eq!(map.height(), Some(1));
1167+
assert_eq!(DROPS.load(Ordering::SeqCst), 1);
1168+
1169+
let mut it = map.drain_filter(|_, _| true);
1170+
assert_eq!(it.size_hint(), (0, Some(MIN_INSERTS_HEIGHT_1 - 1)));
1171+
assert_eq!(it.next().unwrap().0, 1);
1172+
assert_eq!(it.size_hint(), (0, Some(MIN_INSERTS_HEIGHT_1 - 2)));
1173+
assert_eq!(DROPS.load(Ordering::SeqCst), 2);
1174+
mem::forget(it);
1175+
assert_eq!(map.len(), 0);
1176+
assert_eq!(map.height(), None);
1177+
map.check();
1178+
drop(map);
1179+
assert_eq!(DROPS.load(Ordering::SeqCst), 2); // more would be nice to have.
1180+
}
11501181
}
11511182

11521183
#[test]

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

+11
Original file line numberDiff line numberDiff line change
@@ -333,3 +333,14 @@ impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
333333
}
334334
}
335335
}
336+
337+
impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
338+
pub fn root_node(mut self) -> Self {
339+
loop {
340+
self = match self.ascend() {
341+
Ok(edge) => edge.into_node().forget_type(),
342+
Err(root) => return root,
343+
}
344+
}
345+
}
346+
}

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

+14-5
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,12 @@ impl<K, V> Root<K, V> {
166166
Root { node: BoxedNode::from_leaf(Box::new(unsafe { LeafNode::new() })), height: 0 }
167167
}
168168

169+
pub fn restore_from_node<'a>(
170+
root_node: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
171+
) -> Self {
172+
Root { node: unsafe { BoxedNode::from_ptr(root_node.node) }, height: root_node.height() }
173+
}
174+
169175
/// Borrows and returns an immutable reference to the node owned by the root.
170176
pub fn node_as_ref(&self) -> NodeRef<marker::Immut<'_>, K, V, marker::LeafOrInternal> {
171177
NodeRef { height: self.height, node: self.node.as_ptr(), _marker: PhantomData }
@@ -176,6 +182,11 @@ impl<K, V> Root<K, V> {
176182
NodeRef { height: self.height, node: self.node.as_ptr(), _marker: PhantomData }
177183
}
178184

185+
/// Converts root into a mutable reference to its node.
186+
pub fn into_node_mut<'a>(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
187+
NodeRef { height: self.height, node: self.node.as_ptr(), _marker: PhantomData }
188+
}
189+
179190
pub fn into_ref(self) -> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
180191
NodeRef { height: self.height, node: self.node.as_ptr(), _marker: PhantomData }
181192
}
@@ -205,8 +216,7 @@ impl<K, V> Root<K, V> {
205216
/// no cleanup is done on any of the other children.
206217
/// This decreases the height by 1 and is the opposite of `push_internal_level`.
207218
///
208-
/// Requires exclusive access to the `Root` object but not to the root node;
209-
/// it will not invalidate existing handles or references to the root node.
219+
/// Requires exclusive access to the `Root` object and to the root node.
210220
///
211221
/// Panics if there is no internal level, i.e., if the root node is a leaf.
212222
pub fn pop_internal_level(&mut self) {
@@ -220,9 +230,8 @@ impl<K, V> Root<K, V> {
220230
)
221231
};
222232
self.height -= 1;
223-
unsafe {
224-
(*self.node_as_mut().as_leaf_mut()).parent = ptr::null();
225-
}
233+
let root_node = unsafe { &mut *self.node_as_mut().as_leaf_mut() };
234+
root_node.parent = ptr::null();
226235

227236
unsafe {
228237
Global.dealloc(NonNull::from(top).cast(), Layout::new::<InternalNode<K, V>>());

0 commit comments

Comments
 (0)