@@ -1288,15 +1288,15 @@ impl<K: Ord, V> BTreeMap<K, V> {
1288
1288
let front = root. into_node_mut ( ) . first_leaf_edge ( ) ;
1289
1289
DrainFilterInner {
1290
1290
dormant_map : Some ( dormant_map) ,
1291
- position : Some ( DrainFilterPosition :: Edge ( front) ) ,
1291
+ position : DrainFilterPosition :: Edge ( front) ,
1292
1292
height,
1293
1293
length,
1294
1294
remaining : length,
1295
1295
}
1296
1296
} else {
1297
1297
DrainFilterInner {
1298
1298
dormant_map : None ,
1299
- position : None ,
1299
+ position : DrainFilterPosition :: None ,
1300
1300
height : 0 ,
1301
1301
length : 0 ,
1302
1302
remaining : 0 ,
@@ -1688,12 +1688,14 @@ pub(super) struct DrainFilterInner<'a, K: 'a, V: 'a> {
1688
1688
// wrapped in an Option to be able to `take` it in the drop handler.
1689
1689
dormant_map : Option < DormantMutRef < ' a , BTreeMap < K , V > > > ,
1690
1690
// wrapped in an Option some maps don't have any:
1691
- position : Option < DrainFilterPosition < ' a , K , V > > ,
1691
+ position : DrainFilterPosition < ' a , K , V > ,
1692
1692
height : usize ,
1693
1693
length : usize ,
1694
1694
remaining : usize ,
1695
1695
}
1696
1696
enum DrainFilterPosition < ' a , K : ' a , V : ' a > {
1697
+ // There is nothing to position on (or BTreeMap's own code panicked).
1698
+ None ,
1697
1699
// Initial and normal position on a leaf edge.
1698
1700
Edge ( Handle < NodeRef < marker:: Mut < ' a > , K , V , marker:: Leaf > , marker:: Edge > ) ,
1699
1701
// Intermediate position on a KV, that sticks around if a predicate panics.
@@ -1743,17 +1745,17 @@ where
1743
1745
impl < K , V > Drop for DrainFilterInner < ' _ , K , V > {
1744
1746
fn drop ( & mut self ) {
1745
1747
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) => {
1749
1751
// panic during drop
1750
1752
edge. into_node ( ) . forget_type ( ) . root_node ( )
1751
1753
}
1752
- Some ( DrainFilterPosition :: KV ( kv) ) => {
1754
+ DrainFilterPosition :: KV ( kv) => {
1753
1755
// panic in predicate
1754
1756
kv. into_node ( ) . root_node ( )
1755
1757
}
1756
- Some ( DrainFilterPosition :: Root ( root_node) ) => root_node,
1758
+ DrainFilterPosition :: Root ( root_node) => root_node,
1757
1759
} ;
1758
1760
let mut root = node:: Root :: restore_from_node ( root_node) ;
1759
1761
for _ in self . height ..root. height ( ) {
@@ -1771,7 +1773,7 @@ impl<K, V> Drop for DrainFilterInner<'_, K, V> {
1771
1773
impl < ' a , K : ' a , V : ' a > DrainFilterInner < ' a , K , V > {
1772
1774
/// Allow Debug implementations to predict the next element.
1773
1775
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 {
1775
1777
edge. reborrow ( ) . next_kv ( ) . ok ( ) . map ( |kv| kv. into_kv ( ) )
1776
1778
} else {
1777
1779
None
@@ -1783,41 +1785,33 @@ impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
1783
1785
where
1784
1786
F : FnMut ( & K , & mut V ) -> bool ,
1785
1787
{
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 ! ( ) ,
1797
1794
} ;
1798
1795
1799
1796
loop {
1800
1797
match cur_leaf_edge. next_kv ( ) {
1801
1798
Err ( root_node) => {
1802
- self . position = Some ( DrainFilterPosition :: Root ( root_node) ) ;
1799
+ self . position = DrainFilterPosition :: Root ( root_node) ;
1803
1800
return None ;
1804
1801
}
1805
1802
Ok ( kv) => {
1806
1803
self . remaining -= 1 ;
1807
1804
// 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
+ {
1821
1815
kv
1822
1816
} else {
1823
1817
unreachable ! ( )
@@ -1827,7 +1821,7 @@ impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
1827
1821
let ( kv, pos) = kv. remove_kv_tracking ( |emptied_internal_node| {
1828
1822
self . height = emptied_internal_node. height ( ) - 1
1829
1823
} ) ;
1830
- self . position = Some ( DrainFilterPosition :: Edge ( pos) ) ;
1824
+ self . position = DrainFilterPosition :: Edge ( pos) ;
1831
1825
return Some ( kv) ;
1832
1826
}
1833
1827
cur_leaf_edge = kv. next_leaf_edge ( ) ;
0 commit comments