Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit b02a619

Browse files
committedFeb 23, 2021
Auto merge of #81937 - ssomers:btree_drainy_refactor_9b, r=Mark-Simulacrum
BTree: move more shared iterator code into navigate.rs The functions in navigate.rs only exist to support iterators, and these look easier on my eyes if there is a shared `struct` with the recurring pair of handles. r? `@Mark-Simulacrum`
2 parents 11f838d + 342aa69 commit b02a619

File tree

2 files changed

+67
-98
lines changed

2 files changed

+67
-98
lines changed
 

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

Lines changed: 35 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use core::ops::{Index, RangeBounds};
99
use core::ptr;
1010

1111
use super::borrow::DormantMutRef;
12+
use super::navigate::LeafRange;
1213
use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};
1314
use super::search::SearchResult::*;
1415

@@ -307,20 +308,15 @@ pub struct IterMut<'a, K: 'a, V: 'a> {
307308
/// [`into_iter`]: IntoIterator::into_iter
308309
#[stable(feature = "rust1", since = "1.0.0")]
309310
pub struct IntoIter<K, V> {
310-
front: Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
311-
back: Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
311+
range: LeafRange<marker::Dying, K, V>,
312312
length: usize,
313313
}
314314

315315
impl<K, V> IntoIter<K, V> {
316316
/// Returns an iterator of references over the remaining items.
317317
#[inline]
318318
pub(super) fn iter(&self) -> Iter<'_, K, V> {
319-
let range = Range {
320-
front: self.front.as_ref().map(|f| f.reborrow()),
321-
back: self.back.as_ref().map(|b| b.reborrow()),
322-
};
323-
319+
let range = Range { inner: self.range.reborrow() };
324320
Iter { range: range, length: self.length }
325321
}
326322
}
@@ -438,8 +434,7 @@ impl<K, V: fmt::Debug> fmt::Debug for IntoValues<K, V> {
438434
/// [`range`]: BTreeMap::range
439435
#[stable(feature = "btree_range", since = "1.17.0")]
440436
pub struct Range<'a, K: 'a, V: 'a> {
441-
front: Option<Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>>,
442-
back: Option<Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>>,
437+
inner: LeafRange<marker::Immut<'a>, K, V>,
443438
}
444439

445440
#[stable(feature = "collection_debug", since = "1.17.0")]
@@ -457,8 +452,7 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for Range<'_, K, V> {
457452
/// [`range_mut`]: BTreeMap::range_mut
458453
#[stable(feature = "btree_range", since = "1.17.0")]
459454
pub struct RangeMut<'a, K: 'a, V: 'a> {
460-
front: Option<Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>>,
461-
back: Option<Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>>,
455+
inner: LeafRange<marker::ValMut<'a>, K, V>,
462456

463457
// Be invariant in `K` and `V`
464458
_marker: PhantomData<&'a mut (K, V)>,
@@ -467,10 +461,7 @@ pub struct RangeMut<'a, K: 'a, V: 'a> {
467461
#[stable(feature = "collection_debug", since = "1.17.0")]
468462
impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for RangeMut<'_, K, V> {
469463
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
470-
let range = Range {
471-
front: self.front.as_ref().map(|f| f.reborrow()),
472-
back: self.back.as_ref().map(|b| b.reborrow()),
473-
};
464+
let range = Range { inner: self.inner.reborrow() };
474465
f.debug_list().entries(range).finish()
475466
}
476467
}
@@ -1018,11 +1009,9 @@ impl<K, V> BTreeMap<K, V> {
10181009
R: RangeBounds<T>,
10191010
{
10201011
if let Some(root) = &self.root {
1021-
let (f, b) = root.reborrow().range_search(range);
1022-
1023-
Range { front: Some(f), back: Some(b) }
1012+
Range { inner: root.reborrow().range_search(range) }
10241013
} else {
1025-
Range { front: None, back: None }
1014+
Range { inner: LeafRange::none() }
10261015
}
10271016
}
10281017

@@ -1064,11 +1053,9 @@ impl<K, V> BTreeMap<K, V> {
10641053
R: RangeBounds<T>,
10651054
{
10661055
if let Some(root) = &mut self.root {
1067-
let (f, b) = root.borrow_valmut().range_search(range);
1068-
1069-
RangeMut { front: Some(f), back: Some(b), _marker: PhantomData }
1056+
RangeMut { inner: root.borrow_valmut().range_search(range), _marker: PhantomData }
10701057
} else {
1071-
RangeMut { front: None, back: None, _marker: PhantomData }
1058+
RangeMut { inner: LeafRange::none(), _marker: PhantomData }
10721059
}
10731060
}
10741061

@@ -1407,11 +1394,11 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
14071394
fn into_iter(self) -> IntoIter<K, V> {
14081395
let mut me = ManuallyDrop::new(self);
14091396
if let Some(root) = me.root.take() {
1410-
let (f, b) = root.into_dying().full_range();
1397+
let full_range = root.into_dying().full_range();
14111398

1412-
IntoIter { front: Some(f), back: Some(b), length: me.length }
1399+
IntoIter { range: full_range, length: me.length }
14131400
} else {
1414-
IntoIter { front: None, back: None, length: 0 }
1401+
IntoIter { range: LeafRange::none(), length: 0 }
14151402
}
14161403
}
14171404
}
@@ -1450,7 +1437,7 @@ impl<K, V> Drop for Dropper<K, V> {
14501437
#[stable(feature = "btree_drop", since = "1.7.0")]
14511438
impl<K, V> Drop for IntoIter<K, V> {
14521439
fn drop(&mut self) {
1453-
if let Some(front) = self.front.take() {
1440+
if let Some(front) = self.range.front.take() {
14541441
Dropper { front, remaining_length: self.length };
14551442
}
14561443
}
@@ -1465,7 +1452,7 @@ impl<K, V> Iterator for IntoIter<K, V> {
14651452
None
14661453
} else {
14671454
self.length -= 1;
1468-
Some(unsafe { self.front.as_mut().unwrap().deallocating_next_unchecked() })
1455+
Some(unsafe { self.range.front.as_mut().unwrap().deallocating_next_unchecked() })
14691456
}
14701457
}
14711458

@@ -1481,7 +1468,7 @@ impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
14811468
None
14821469
} else {
14831470
self.length -= 1;
1484-
Some(unsafe { self.back.as_mut().unwrap().deallocating_next_back_unchecked() })
1471+
Some(unsafe { self.range.back.as_mut().unwrap().deallocating_next_back_unchecked() })
14851472
}
14861473
}
14871474
}
@@ -1698,7 +1685,7 @@ impl<'a, K, V> Iterator for Range<'a, K, V> {
16981685
type Item = (&'a K, &'a V);
16991686

17001687
fn next(&mut self) -> Option<(&'a K, &'a V)> {
1701-
if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
1688+
if self.inner.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
17021689
}
17031690

17041691
fn last(mut self) -> Option<(&'a K, &'a V)> {
@@ -1749,12 +1736,8 @@ impl<K, V> ExactSizeIterator for ValuesMut<'_, K, V> {
17491736
impl<K, V> FusedIterator for ValuesMut<'_, K, V> {}
17501737

17511738
impl<'a, K, V> Range<'a, K, V> {
1752-
fn is_empty(&self) -> bool {
1753-
self.front == self.back
1754-
}
1755-
17561739
unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
1757-
unsafe { self.front.as_mut().unwrap_unchecked().next_unchecked() }
1740+
unsafe { self.inner.front.as_mut().unwrap_unchecked().next_unchecked() }
17581741
}
17591742
}
17601743

@@ -1837,13 +1820,13 @@ impl<K, V> FusedIterator for IntoValues<K, V> {}
18371820
#[stable(feature = "btree_range", since = "1.17.0")]
18381821
impl<'a, K, V> DoubleEndedIterator for Range<'a, K, V> {
18391822
fn next_back(&mut self) -> Option<(&'a K, &'a V)> {
1840-
if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
1823+
if self.inner.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
18411824
}
18421825
}
18431826

18441827
impl<'a, K, V> Range<'a, K, V> {
18451828
unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
1846-
unsafe { self.back.as_mut().unwrap_unchecked().next_back_unchecked() }
1829+
unsafe { self.inner.back.as_mut().unwrap_unchecked().next_back_unchecked() }
18471830
}
18481831
}
18491832

@@ -1853,7 +1836,7 @@ impl<K, V> FusedIterator for Range<'_, K, V> {}
18531836
#[stable(feature = "btree_range", since = "1.17.0")]
18541837
impl<K, V> Clone for Range<'_, K, V> {
18551838
fn clone(&self) -> Self {
1856-
Range { front: self.front, back: self.back }
1839+
Range { inner: LeafRange { front: self.inner.front, back: self.inner.back } }
18571840
}
18581841
}
18591842

@@ -1862,7 +1845,7 @@ impl<'a, K, V> Iterator for RangeMut<'a, K, V> {
18621845
type Item = (&'a K, &'a mut V);
18631846

18641847
fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
1865-
if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
1848+
if self.inner.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
18661849
}
18671850

18681851
fn last(mut self) -> Option<(&'a K, &'a mut V)> {
@@ -1879,28 +1862,21 @@ impl<'a, K, V> Iterator for RangeMut<'a, K, V> {
18791862
}
18801863

18811864
impl<'a, K, V> RangeMut<'a, K, V> {
1882-
fn is_empty(&self) -> bool {
1883-
self.front == self.back
1884-
}
1885-
18861865
unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
1887-
unsafe { self.front.as_mut().unwrap_unchecked().next_unchecked() }
1866+
unsafe { self.inner.front.as_mut().unwrap_unchecked().next_unchecked() }
18881867
}
18891868

18901869
/// Returns an iterator of references over the remaining items.
18911870
#[inline]
18921871
pub(super) fn iter(&self) -> Range<'_, K, V> {
1893-
Range {
1894-
front: self.front.as_ref().map(|f| f.reborrow()),
1895-
back: self.back.as_ref().map(|b| b.reborrow()),
1896-
}
1872+
Range { inner: self.inner.reborrow() }
18971873
}
18981874
}
18991875

19001876
#[stable(feature = "btree_range", since = "1.17.0")]
19011877
impl<'a, K, V> DoubleEndedIterator for RangeMut<'a, K, V> {
19021878
fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> {
1903-
if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
1879+
if self.inner.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
19041880
}
19051881
}
19061882

@@ -1909,7 +1885,7 @@ impl<K, V> FusedIterator for RangeMut<'_, K, V> {}
19091885

19101886
impl<'a, K, V> RangeMut<'a, K, V> {
19111887
unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
1912-
unsafe { self.back.as_mut().unwrap_unchecked().next_back_unchecked() }
1888+
unsafe { self.inner.back.as_mut().unwrap_unchecked().next_back_unchecked() }
19131889
}
19141890
}
19151891

@@ -2043,11 +2019,11 @@ impl<K, V> BTreeMap<K, V> {
20432019
#[stable(feature = "rust1", since = "1.0.0")]
20442020
pub fn iter(&self) -> Iter<'_, K, V> {
20452021
if let Some(root) = &self.root {
2046-
let (f, b) = root.reborrow().full_range();
2022+
let full_range = root.reborrow().full_range();
20472023

2048-
Iter { range: Range { front: Some(f), back: Some(b) }, length: self.length }
2024+
Iter { range: Range { inner: full_range }, length: self.length }
20492025
} else {
2050-
Iter { range: Range { front: None, back: None }, length: 0 }
2026+
Iter { range: Range { inner: LeafRange::none() }, length: 0 }
20512027
}
20522028
}
20532029

@@ -2075,14 +2051,17 @@ impl<K, V> BTreeMap<K, V> {
20752051
#[stable(feature = "rust1", since = "1.0.0")]
20762052
pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
20772053
if let Some(root) = &mut self.root {
2078-
let (f, b) = root.borrow_valmut().full_range();
2054+
let full_range = root.borrow_valmut().full_range();
20792055

20802056
IterMut {
2081-
range: RangeMut { front: Some(f), back: Some(b), _marker: PhantomData },
2057+
range: RangeMut { inner: full_range, _marker: PhantomData },
20822058
length: self.length,
20832059
}
20842060
} else {
2085-
IterMut { range: RangeMut { front: None, back: None, _marker: PhantomData }, length: 0 }
2061+
IterMut {
2062+
range: RangeMut { inner: LeafRange::none(), _marker: PhantomData },
2063+
length: 0,
2064+
}
20862065
}
20872066
}
20882067

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

Lines changed: 32 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,29 @@ use core::ptr;
77
use super::node::{marker, ForceResult::*, Handle, NodeRef};
88
use super::search::SearchResult;
99

10+
pub struct LeafRange<BorrowType, K, V> {
11+
pub front: Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
12+
pub back: Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
13+
}
14+
15+
impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {
16+
pub fn none() -> Self {
17+
LeafRange { front: None, back: None }
18+
}
19+
20+
pub fn is_empty(&self) -> bool {
21+
self.front == self.back
22+
}
23+
24+
/// Temporarily takes out another, immutable equivalent of the same range.
25+
pub fn reborrow(&self) -> LeafRange<marker::Immut<'_>, K, V> {
26+
LeafRange {
27+
front: self.front.as_ref().map(|f| f.reborrow()),
28+
back: self.back.as_ref().map(|b| b.reborrow()),
29+
}
30+
}
31+
}
32+
1033
/// Finds the leaf edges delimiting a specified range in or underneath a node.
1134
///
1235
/// The result is meaningful only if the tree is ordered by key, like the tree
@@ -15,10 +38,7 @@ fn range_search<BorrowType: marker::BorrowType, K, V, Q, R>(
1538
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
1639
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
1740
range: R,
18-
) -> (
19-
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
20-
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
21-
)
41+
) -> LeafRange<BorrowType, K, V>
2242
where
2343
Q: ?Sized + Ord,
2444
K: Borrow<Q>,
@@ -92,7 +112,7 @@ where
92112
}
93113
match (front.force(), back.force()) {
94114
(Leaf(f), Leaf(b)) => {
95-
return (f, b);
115+
return LeafRange { front: Some(f), back: Some(b) };
96116
}
97117
(Internal(min_int), Internal(max_int)) => {
98118
min_node = min_int.descend();
@@ -108,18 +128,15 @@ where
108128
fn full_range<BorrowType: marker::BorrowType, K, V>(
109129
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
110130
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
111-
) -> (
112-
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
113-
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
114-
) {
131+
) -> LeafRange<BorrowType, K, V> {
115132
let mut min_node = root1;
116133
let mut max_node = root2;
117134
loop {
118135
let front = min_node.first_edge();
119136
let back = max_node.last_edge();
120137
match (front.force(), back.force()) {
121138
(Leaf(f), Leaf(b)) => {
122-
return (f, b);
139+
return LeafRange { front: Some(f), back: Some(b) };
123140
}
124141
(Internal(min_int), Internal(max_int)) => {
125142
min_node = min_int.descend();
@@ -135,13 +152,7 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>
135152
///
136153
/// The result is meaningful only if the tree is ordered by key, like the tree
137154
/// in a `BTreeMap` is.
138-
pub fn range_search<Q, R>(
139-
self,
140-
range: R,
141-
) -> (
142-
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
143-
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
144-
)
155+
pub fn range_search<Q, R>(self, range: R) -> LeafRange<marker::Immut<'a>, K, V>
145156
where
146157
Q: ?Sized + Ord,
147158
K: Borrow<Q>,
@@ -151,12 +162,7 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>
151162
}
152163

153164
/// Finds the pair of leaf edges delimiting an entire tree.
154-
pub fn full_range(
155-
self,
156-
) -> (
157-
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
158-
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
159-
) {
165+
pub fn full_range(self) -> LeafRange<marker::Immut<'a>, K, V> {
160166
full_range(self, self)
161167
}
162168
}
@@ -168,13 +174,7 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::ValMut<'a>, K, V, marker::LeafOrInternal>
168174
///
169175
/// The result is meaningful only if the tree is ordered by key, like the tree
170176
/// in a `BTreeMap` is.
171-
pub fn range_search<Q, R>(
172-
self,
173-
range: R,
174-
) -> (
175-
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
176-
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
177-
)
177+
pub fn range_search<Q, R>(self, range: R) -> LeafRange<marker::ValMut<'a>, K, V>
178178
where
179179
Q: ?Sized + Ord,
180180
K: Borrow<Q>,
@@ -189,12 +189,7 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::ValMut<'a>, K, V, marker::LeafOrInternal>
189189
/// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
190190
/// The results are non-unique references allowing mutation (of values only), so must be used
191191
/// with care.
192-
pub fn full_range(
193-
self,
194-
) -> (
195-
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
196-
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
197-
) {
192+
pub fn full_range(self) -> LeafRange<marker::ValMut<'a>, K, V> {
198193
// We duplicate the root NodeRef here -- we will never visit the same KV
199194
// twice, and never end up with overlapping value references.
200195
let self2 = unsafe { ptr::read(&self) };
@@ -206,12 +201,7 @@ impl<K, V> NodeRef<marker::Dying, K, V, marker::LeafOrInternal> {
206201
/// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
207202
/// The results are non-unique references allowing massively destructive mutation, so must be
208203
/// used with the utmost care.
209-
pub fn full_range(
210-
self,
211-
) -> (
212-
Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
213-
Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
214-
) {
204+
pub fn full_range(self) -> LeafRange<marker::Dying, K, V> {
215205
// We duplicate the root NodeRef here -- we will never access it in a way
216206
// that overlaps references obtained from the root.
217207
let self2 = unsafe { ptr::read(&self) };

0 commit comments

Comments
 (0)
Please sign in to comment.