From cf6ba847ae0a0f2b4afcdb1e4d2ba763661f3316 Mon Sep 17 00:00:00 2001 From: Stepan Koltsov Date: Sat, 20 Dec 2025 07:43:01 +0000 Subject: [PATCH] IteratorRandom::sample: allow large "amount" --- CHANGELOG.md | 1 + src/seq/iterator.rs | 9 ++++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 66a4e280f5..ba22f10ba6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,7 @@ You may also find the [Upgrade Guide](https://rust-random.github.io/book/update. - Remove fns `SeedableRng::from_os_rng`, `try_from_os_rng` (#1674) - Remove `Clone` support for `StdRng`, `ReseedingRng` (#1677) - Use `postcard` instead of `bincode` to test the serde feature (#1693) +- Avoid excessive allocation in `IteratorRandom::sample` when `amount` is much larger than iterator size ### Additions - Add fns `IndexedRandom::choose_iter`, `choose_weighted_iter` (#1632) diff --git a/src/seq/iterator.rs b/src/seq/iterator.rs index 187ca3c055..f6addb9fcf 100644 --- a/src/seq/iterator.rs +++ b/src/seq/iterator.rs @@ -245,7 +245,14 @@ pub trait IteratorRandom: Iterator + Sized { where R: Rng + ?Sized, { - let mut reservoir = Vec::with_capacity(amount); + let (size_lower_bound, _size_upper_bound) = self.size_hint(); + + // If the requested amount is much larger than the iterator size, + // avoid excessive allocation even that it is temporary. + // This also prevents allocation panic for small iterators but large requested amount. + let capacity = core::cmp::min(amount, size_lower_bound); + + let mut reservoir = Vec::with_capacity(capacity); reservoir.extend(self.by_ref().take(amount)); // Continue unless the iterator was exhausted