Skip to content

Commit

Permalink
Merge pull request #1059 from kazcw/iteratorrandom-choose
Browse files Browse the repository at this point in the history
fix accuracy of IteratorRandom::choose()
  • Loading branch information
dhardy committed Oct 17, 2020
2 parents c42d027 + a37a599 commit 2dda0c9
Showing 1 changed file with 5 additions and 6 deletions.
11 changes: 5 additions & 6 deletions src/seq/mod.rs
Expand Up @@ -296,15 +296,15 @@ pub trait IteratorRandom: Iterator + Sized {
/// depends on size hints. In particular, `Iterator` combinators that don't
/// change the values yielded but change the size hints may result in
/// `choose` returning different elements.
///
/// For slices, prefer [`SliceRandom::choose`] which guarantees `O(1)`
/// performance.
fn choose<R>(mut self, rng: &mut R) -> Option<Self::Item>
where R: Rng + ?Sized {
let (mut lower, mut upper) = self.size_hint();
let mut consumed = 0;
let mut result = None;

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

assert_eq!(choose([].iter().cloned()), None);
assert_eq!(choose(0..100), Some(33));
assert_eq!(choose(UnhintedIterator { iter: 0..100 }), Some(76));
assert_eq!(choose(UnhintedIterator { iter: 0..100 }), Some(40));
assert_eq!(
choose(ChunkHintedIterator {
iter: 0..100,
Expand Down

0 comments on commit 2dda0c9

Please sign in to comment.