1
+ // ignore-tidy-filelength
2
+
1
3
use core:: borrow:: Borrow ;
2
4
use core:: cmp:: Ordering ;
3
5
use core:: fmt:: { self , Debug } ;
@@ -1277,16 +1279,28 @@ impl<K: Ord, V> BTreeMap<K, V> {
1277
1279
}
1278
1280
1279
1281
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 ( ) ;
1283
1289
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,
1287
1295
}
1288
1296
} 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
+ }
1290
1304
}
1291
1305
}
1292
1306
@@ -1671,11 +1685,21 @@ where
1671
1685
/// Most of the implementation of DrainFilter, independent of the type
1672
1686
/// of the predicate, thus also serving for BTreeSet::DrainFilter.
1673
1687
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 > ) ,
1679
1703
}
1680
1704
1681
1705
#[ unstable( feature = "btree_drain_filter" , issue = "70530" ) ]
@@ -1716,40 +1740,105 @@ where
1716
1740
}
1717
1741
}
1718
1742
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
+
1719
1771
impl < ' a , K : ' a , V : ' a > DrainFilterInner < ' a , K , V > {
1720
1772
/// Allow Debug implementations to predict the next element.
1721
1773
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
+ }
1724
1779
}
1725
1780
1726
1781
/// Implementation of a typical `DrainFilter::next` method, given the predicate.
1727
1782
pub ( super ) fn next < F > ( & mut self , pred : & mut F ) -> Option < ( K , V ) >
1728
1783
where
1729
1784
F : FnMut ( & K , & mut V ) -> bool ,
1730
1785
{
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
+ }
1744
1835
}
1745
- self . cur_leaf_edge = Some ( kv. next_leaf_edge ( ) ) ;
1746
1836
}
1747
- None
1748
1837
}
1749
1838
1750
1839
/// Implementation of a typical `DrainFilter::size_hint` method.
1751
1840
pub ( super ) fn size_hint ( & self ) -> ( usize , Option < usize > ) {
1752
- ( 0 , Some ( * self . length ) )
1841
+ ( 0 , Some ( self . remaining ) )
1753
1842
}
1754
1843
}
1755
1844
@@ -2767,7 +2856,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
2767
2856
// Body of `remove_entry`, separate to keep the above implementations short.
2768
2857
fn remove_kv ( self ) -> ( K , V ) {
2769
2858
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 ) ;
2771
2860
// SAFETY: we consumed the intermediate root borrow, `self.handle`.
2772
2861
let map = unsafe { self . dormant_map . awaken ( ) } ;
2773
2862
map. length -= 1 ;
@@ -2782,10 +2871,13 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
2782
2871
impl < ' a , K : ' a , V : ' a > Handle < NodeRef < marker:: Mut < ' a > , K , V , marker:: LeafOrInternal > , marker:: KV > {
2783
2872
/// Removes a key/value-pair from the map, and returns that pair, as well as
2784
2873
/// the leaf edge corresponding to that former pair.
2785
- fn remove_kv_tracking < F : FnOnce ( ) > (
2874
+ fn remove_kv_tracking < F > (
2786
2875
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
+ {
2789
2881
let ( old_kv, mut pos, was_internal) = match self . force ( ) {
2790
2882
Leaf ( leaf) => {
2791
2883
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
2836
2928
if parent. len ( ) == 0 {
2837
2929
// The parent that was just emptied must be the root,
2838
2930
// 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) ;
2841
2935
break ;
2842
2936
} else {
2843
2937
cur_node = parent. forget_type ( ) ;
@@ -2948,8 +3042,16 @@ fn handle_underfull_node<K, V>(
2948
3042
let ( is_left, mut handle) = match parent. left_kv ( ) {
2949
3043
Ok ( left) => ( true , left) ,
2950
3044
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
+ }
2953
3055
}
2954
3056
} ;
2955
3057
0 commit comments