Skip to content

Commit b6269f2

Browse files
committed
Auto merge of #64383 - pcpthm:btreeset-size-hint, r=dtolnay
Improve BTreeSet::Intersection::size_hint A comment on `IntersectionInner` mentions `small_iter` should be smaller than `other_iter` but this condition is broken while iterating because those two iterators can be consumed at a different rate. I added a test to demonstrate this situation. <del>I made `small_iter.len() < other_iter.len()` always true by swapping two iterators when that condition became false. This change affects the return value of `size_hint`. The previous result was also correct but this new version always returns smaller upper bound than the previous version.</del> I changed `size_hint` to taking minimum of both lengths of iterators and renamed fields to `a` and `b` to match `Union` iterator.
2 parents 3e3e06d + 4333b86 commit b6269f2

File tree

2 files changed

+33
-22
lines changed

2 files changed

+33
-22
lines changed

src/liballoc/collections/btree/set.rs

+22-22
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
use core::borrow::Borrow;
55
use core::cmp::Ordering::{self, Less, Greater, Equal};
6-
use core::cmp::max;
6+
use core::cmp::{max, min};
77
use core::fmt::{self, Debug};
88
use core::iter::{Peekable, FromIterator, FusedIterator};
99
use core::ops::{BitOr, BitAnd, BitXor, Sub, RangeBounds};
@@ -187,8 +187,8 @@ pub struct Intersection<'a, T: 'a> {
187187
}
188188
enum IntersectionInner<'a, T: 'a> {
189189
Stitch {
190-
small_iter: Iter<'a, T>, // for size_hint, should be the smaller of the sets
191-
other_iter: Iter<'a, T>,
190+
a: Iter<'a, T>,
191+
b: Iter<'a, T>,
192192
},
193193
Search {
194194
small_iter: Iter<'a, T>,
@@ -201,12 +201,12 @@ impl<T: fmt::Debug> fmt::Debug for Intersection<'_, T> {
201201
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
202202
match &self.inner {
203203
IntersectionInner::Stitch {
204-
small_iter,
205-
other_iter,
204+
a,
205+
b,
206206
} => f
207207
.debug_tuple("Intersection")
208-
.field(&small_iter)
209-
.field(&other_iter)
208+
.field(&a)
209+
.field(&b)
210210
.finish(),
211211
IntersectionInner::Search {
212212
small_iter,
@@ -397,8 +397,8 @@ impl<T: Ord> BTreeSet<T> {
397397
// Iterate both sets jointly, spotting matches along the way.
398398
Intersection {
399399
inner: IntersectionInner::Stitch {
400-
small_iter: small.iter(),
401-
other_iter: other.iter(),
400+
a: small.iter(),
401+
b: other.iter(),
402402
},
403403
}
404404
} else {
@@ -1221,11 +1221,11 @@ impl<T> Clone for Intersection<'_, T> {
12211221
Intersection {
12221222
inner: match &self.inner {
12231223
IntersectionInner::Stitch {
1224-
small_iter,
1225-
other_iter,
1224+
a,
1225+
b,
12261226
} => IntersectionInner::Stitch {
1227-
small_iter: small_iter.clone(),
1228-
other_iter: other_iter.clone(),
1227+
a: a.clone(),
1228+
b: b.clone(),
12291229
},
12301230
IntersectionInner::Search {
12311231
small_iter,
@@ -1245,16 +1245,16 @@ impl<'a, T: Ord> Iterator for Intersection<'a, T> {
12451245
fn next(&mut self) -> Option<&'a T> {
12461246
match &mut self.inner {
12471247
IntersectionInner::Stitch {
1248-
small_iter,
1249-
other_iter,
1248+
a,
1249+
b,
12501250
} => {
1251-
let mut small_next = small_iter.next()?;
1252-
let mut other_next = other_iter.next()?;
1251+
let mut a_next = a.next()?;
1252+
let mut b_next = b.next()?;
12531253
loop {
1254-
match Ord::cmp(small_next, other_next) {
1255-
Less => small_next = small_iter.next()?,
1256-
Greater => other_next = other_iter.next()?,
1257-
Equal => return Some(small_next),
1254+
match Ord::cmp(a_next, b_next) {
1255+
Less => a_next = a.next()?,
1256+
Greater => b_next = b.next()?,
1257+
Equal => return Some(a_next),
12581258
}
12591259
}
12601260
}
@@ -1272,7 +1272,7 @@ impl<'a, T: Ord> Iterator for Intersection<'a, T> {
12721272

12731273
fn size_hint(&self) -> (usize, Option<usize>) {
12741274
let min_len = match &self.inner {
1275-
IntersectionInner::Stitch { small_iter, .. } => small_iter.len(),
1275+
IntersectionInner::Stitch { a, b } => min(a.len(), b.len()),
12761276
IntersectionInner::Search { small_iter, .. } => small_iter.len(),
12771277
};
12781278
(0, Some(min_len))

src/liballoc/tests/btree/set.rs

+11
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,17 @@ fn test_intersection() {
9090
&[1, 3, 11, 77, 103]);
9191
}
9292

93+
#[test]
94+
fn test_intersection_size_hint() {
95+
let x: BTreeSet<i32> = [3, 4].iter().copied().collect();
96+
let y: BTreeSet<i32> = [1, 2, 3].iter().copied().collect();
97+
let mut iter = x.intersection(&y);
98+
assert_eq!(iter.size_hint(), (0, Some(2)));
99+
assert_eq!(iter.next(), Some(&3));
100+
assert_eq!(iter.size_hint(), (0, Some(0)));
101+
assert_eq!(iter.next(), None);
102+
}
103+
93104
#[test]
94105
fn test_difference() {
95106
fn check_difference(a: &[i32], b: &[i32], expected: &[i32]) {

0 commit comments

Comments
 (0)