shuf: try vec first and fallback to HashMap if it cause OOM (#11169)

Author: oech3

src/uu/shuf/src/nonrepeating_iterator.rs

@@ -49,13 +49,26 @@ enum Values {
 }
 
 impl<'a> NonrepeatingIterator<'a> {
-    pub(crate) fn new(range: RangeInclusive<u64>, rng: &'a mut WrappedRng) -> Self {
-        const MAX_CAPACITY: usize = 128; // todo: optimize this
-        let capacity = (range.size_hint().0).min(MAX_CAPACITY);
-        let values = Values::Sparse(
-            range,
-            FxHashMap::with_capacity_and_hasher(capacity, rustc_hash::FxBuildHasher),
-        );
+    pub(crate) fn new(
+        range: RangeInclusive<u64>,
+        rng: &'a mut WrappedRng,
+        head_count: Option<usize>,
+    ) -> Self {
+        // Save RAM usage with shuf -i 1-huge_number -n small_number
+        const TOO_LARGE_VEC_SIZE: usize = 16_777_216;
+        let range_len = range.size_hint().0;
+        let mut items = Vec::new();
+        let values = if range_len < TOO_LARGE_VEC_SIZE && items.try_reserve(range_len).is_ok() {
+            items.extend(range.rev());
+            Values::Full(items)
+        } else {
+            const MAX_CAPACITY: usize = 128; // todo: optimize this
+            let capacity = head_count.unwrap_or(MAX_CAPACITY).min(range_len);
+            Values::Sparse(
+                range,
+                FxHashMap::with_capacity_and_hasher(capacity, rustc_hash::FxBuildHasher),
+            )
+        };
         NonrepeatingIterator { rng, values }
     }
 

src/uu/shuf/src/shuf.rs

@@ -357,7 +357,7 @@ impl Shufable for RangeInclusive<u64> {
         amount: u64,
     ) -> UResult<impl Iterator<Item = UResult<Self::Item>>> {
         let amount = usize::try_from(amount).unwrap_or(usize::MAX);
-        Ok(NonrepeatingIterator::new(self.clone(), rng).take(amount))
+        Ok(NonrepeatingIterator::new(self.clone(), rng, Some(amount)).take(amount))
     }
 }
 

tests/by-util/test_shuf.rs

@@ -93,23 +93,23 @@ fn test_zero_termination_multi() {
 
 #[test]
 fn test_very_large_range() {
-    let num_samples = 10;
+    let num_samples = 256;
     let result = new_ucmd!()
         .arg("-n")
         .arg(num_samples.to_string())
-        .arg("-i0-1234567890")
+        .arg("-i1-100000000000")
         .succeeds();
     result.no_stderr();
 
-    let result_seq: Vec<isize> = result
+    let result_seq: Vec<u64> = result
         .stdout_str()
         .split('\n')
         .filter(|x| !x.is_empty())
         .map(|x| x.parse().unwrap())
         .collect();
     assert_eq!(result_seq.len(), num_samples, "Miscounted output length!");
     assert!(
-        result_seq.iter().all(|x| (0..=1_234_567_890).contains(x)),
+        result_seq.iter().all(|x| (0..=100_000_000_000).contains(x)),
         "Output includes element not from range: {}",
         result.stdout_str()
     );