Skip to content

Commit cf6ba84

Browse files
committed
IteratorRandom::sample: allow large "amount"
1 parent 8336923 commit cf6ba84

File tree

2 files changed

+9
-1
lines changed

2 files changed

+9
-1
lines changed

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ You may also find the [Upgrade Guide](https://rust-random.github.io/book/update.
1919
- Remove fns `SeedableRng::from_os_rng`, `try_from_os_rng` (#1674)
2020
- Remove `Clone` support for `StdRng`, `ReseedingRng` (#1677)
2121
- Use `postcard` instead of `bincode` to test the serde feature (#1693)
22+
- Avoid excessive allocation in `IteratorRandom::sample` when `amount` is much larger than iterator size
2223

2324
### Additions
2425
- Add fns `IndexedRandom::choose_iter`, `choose_weighted_iter` (#1632)

src/seq/iterator.rs

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,14 @@ pub trait IteratorRandom: Iterator + Sized {
245245
where
246246
R: Rng + ?Sized,
247247
{
248-
let mut reservoir = Vec::with_capacity(amount);
248+
let (size_lower_bound, _size_upper_bound) = self.size_hint();
249+
250+
// If the requested amount is much larger than the iterator size,
251+
// avoid excessive allocation even that it is temporary.
252+
// This also prevents allocation panic for small iterators but large requested amount.
253+
let capacity = core::cmp::min(amount, size_lower_bound);
254+
255+
let mut reservoir = Vec::with_capacity(capacity);
249256
reservoir.extend(self.by_ref().take(amount));
250257

251258
// Continue unless the iterator was exhausted

0 commit comments

Comments
 (0)