Skip to content

Commit 2dda0c9

Browse files
authored
Merge pull request #1059 from kazcw/iteratorrandom-choose
fix accuracy of IteratorRandom::choose()
2 parents c42d027 + a37a599 commit 2dda0c9

File tree

1 file changed

+5
-6
lines changed

1 file changed

+5
-6
lines changed

src/seq/mod.rs

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -296,15 +296,15 @@ pub trait IteratorRandom: Iterator + Sized {
296296
/// depends on size hints. In particular, `Iterator` combinators that don't
297297
/// change the values yielded but change the size hints may result in
298298
/// `choose` returning different elements.
299-
///
300-
/// For slices, prefer [`SliceRandom::choose`] which guarantees `O(1)`
301-
/// performance.
302299
fn choose<R>(mut self, rng: &mut R) -> Option<Self::Item>
303300
where R: Rng + ?Sized {
304301
let (mut lower, mut upper) = self.size_hint();
305302
let mut consumed = 0;
306303
let mut result = None;
307304

305+
// Handling for this condition outside the loop allows the optimizer to eliminate the loop
306+
// when the Iterator is an ExactSizeIterator. This has a large performance impact on e.g.
307+
// seq_iter_choose_from_1000.
308308
if upper == Some(lower) {
309309
return if lower == 0 {
310310
None
@@ -336,8 +336,7 @@ pub trait IteratorRandom: Iterator + Sized {
336336
return result;
337337
}
338338
consumed += 1;
339-
let denom = consumed as f64; // accurate to 2^53 elements
340-
if rng.gen_bool(1.0 / denom) {
339+
if gen_index(rng, consumed) == 0 {
341340
result = elem;
342341
}
343342
}
@@ -963,7 +962,7 @@ mod test {
963962

964963
assert_eq!(choose([].iter().cloned()), None);
965964
assert_eq!(choose(0..100), Some(33));
966-
assert_eq!(choose(UnhintedIterator { iter: 0..100 }), Some(76));
965+
assert_eq!(choose(UnhintedIterator { iter: 0..100 }), Some(40));
967966
assert_eq!(
968967
choose(ChunkHintedIterator {
969968
iter: 0..100,

0 commit comments

Comments
 (0)