Skip to content

Commit 72c4750

Browse files
committed
Option integrated into DrainFilterPosition
1 parent e2242bb commit 72c4750

File tree

1 file changed

+29
-35
lines changed
  • library/alloc/src/collections/btree

1 file changed

+29
-35
lines changed

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

+29-35
Original file line numberDiff line numberDiff line change
@@ -1288,15 +1288,15 @@ impl<K: Ord, V> BTreeMap<K, V> {
12881288
let front = root.into_node_mut().first_leaf_edge();
12891289
DrainFilterInner {
12901290
dormant_map: Some(dormant_map),
1291-
position: Some(DrainFilterPosition::Edge(front)),
1291+
position: DrainFilterPosition::Edge(front),
12921292
height,
12931293
length,
12941294
remaining: length,
12951295
}
12961296
} else {
12971297
DrainFilterInner {
12981298
dormant_map: None,
1299-
position: None,
1299+
position: DrainFilterPosition::None,
13001300
height: 0,
13011301
length: 0,
13021302
remaining: 0,
@@ -1688,12 +1688,14 @@ pub(super) struct DrainFilterInner<'a, K: 'a, V: 'a> {
16881688
// wrapped in an Option to be able to `take` it in the drop handler.
16891689
dormant_map: Option<DormantMutRef<'a, BTreeMap<K, V>>>,
16901690
// wrapped in an Option some maps don't have any:
1691-
position: Option<DrainFilterPosition<'a, K, V>>,
1691+
position: DrainFilterPosition<'a, K, V>,
16921692
height: usize,
16931693
length: usize,
16941694
remaining: usize,
16951695
}
16961696
enum DrainFilterPosition<'a, K: 'a, V: 'a> {
1697+
// There is nothing to position on (or BTreeMap's own code panicked).
1698+
None,
16971699
// Initial and normal position on a leaf edge.
16981700
Edge(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>),
16991701
// Intermediate position on a KV, that sticks around if a predicate panics.
@@ -1743,17 +1745,17 @@ where
17431745
impl<K, V> Drop for DrainFilterInner<'_, K, V> {
17441746
fn drop(&mut self) {
17451747
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)) => {
1748+
let root_node = match mem::replace(&mut self.position, DrainFilterPosition::None) {
1749+
DrainFilterPosition::None => unreachable!(),
1750+
DrainFilterPosition::Edge(edge) => {
17491751
// panic during drop
17501752
edge.into_node().forget_type().root_node()
17511753
}
1752-
Some(DrainFilterPosition::KV(kv)) => {
1754+
DrainFilterPosition::KV(kv) => {
17531755
// panic in predicate
17541756
kv.into_node().root_node()
17551757
}
1756-
Some(DrainFilterPosition::Root(root_node)) => root_node,
1758+
DrainFilterPosition::Root(root_node) => root_node,
17571759
};
17581760
let mut root = node::Root::restore_from_node(root_node);
17591761
for _ in self.height..root.height() {
@@ -1771,7 +1773,7 @@ impl<K, V> Drop for DrainFilterInner<'_, K, V> {
17711773
impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
17721774
/// Allow Debug implementations to predict the next element.
17731775
pub(super) fn peek(&self) -> Option<(&K, &V)> {
1774-
if let Some(DrainFilterPosition::Edge(edge)) = self.position.as_ref() {
1776+
if let DrainFilterPosition::Edge(edge) = &self.position {
17751777
edge.reborrow().next_kv().ok().map(|kv| kv.into_kv())
17761778
} else {
17771779
None
@@ -1783,41 +1785,33 @@ impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
17831785
where
17841786
F: FnMut(&K, &mut V) -> bool,
17851787
{
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,
1788+
if !matches!(&self.position, DrainFilterPosition::Edge(_)) {
1789+
return None;
1790+
}
1791+
let mut cur_leaf_edge = match mem::replace(&mut self.position, DrainFilterPosition::None) {
1792+
DrainFilterPosition::Edge(edge) => edge,
1793+
_ => unreachable!(),
17971794
};
17981795

17991796
loop {
18001797
match cur_leaf_edge.next_kv() {
18011798
Err(root_node) => {
1802-
self.position = Some(DrainFilterPosition::Root(root_node));
1799+
self.position = DrainFilterPosition::Root(root_node);
18031800
return None;
18041801
}
18051802
Ok(kv) => {
18061803
self.remaining -= 1;
18071804
// 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() {
1805+
self.position = DrainFilterPosition::KV(kv);
1806+
let drain = if let DrainFilterPosition::KV(kv) = &mut self.position {
1807+
let (k, v) = kv.kv_mut();
1808+
pred(k, v)
1809+
} else {
1810+
unreachable!()
1811+
};
1812+
let kv = if let DrainFilterPosition::KV(kv) =
1813+
mem::replace(&mut self.position, DrainFilterPosition::None)
1814+
{
18211815
kv
18221816
} else {
18231817
unreachable!()
@@ -1827,7 +1821,7 @@ impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
18271821
let (kv, pos) = kv.remove_kv_tracking(|emptied_internal_node| {
18281822
self.height = emptied_internal_node.height() - 1
18291823
});
1830-
self.position = Some(DrainFilterPosition::Edge(pos));
1824+
self.position = DrainFilterPosition::Edge(pos);
18311825
return Some(kv);
18321826
}
18331827
cur_leaf_edge = kv.next_leaf_edge();

0 commit comments

Comments
 (0)