IteratorRandom::sample: allow large "amount"

Author: stepancheg

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 `IteratorRandom::sample` when `amount` is much larger than iterator size
 
 ### Additions
 - Add fns `IndexedRandom::choose_iter`, `choose_weighted_iter` (#1632)

src/seq/iterator.rs

@@ -245,7 +245,22 @@ 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,
+            core::cmp::max(
+                size_lower_bound,
+                (32usize << 10)
+                    .checked_div(size_of::<Self::Item>())
+                    .unwrap_or(0),
+            ),
+        );
+
+        let mut reservoir = Vec::with_capacity(capacity);
         reservoir.extend(self.by_ref().take(amount));
 
         // Continue unless the iterator was exhausted