diff --git a/src/uu/shuf/src/nonrepeating_iterator.rs b/src/uu/shuf/src/nonrepeating_iterator.rs index 10d143152..0c304f1d6 100644 --- a/src/uu/shuf/src/nonrepeating_iterator.rs +++ b/src/uu/shuf/src/nonrepeating_iterator.rs @@ -49,13 +49,26 @@ enum Values { } impl<'a> NonrepeatingIterator<'a> { - pub(crate) fn new(range: RangeInclusive, 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, + rng: &'a mut WrappedRng, + head_count: Option, + ) -> 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 } } diff --git a/src/uu/shuf/src/shuf.rs b/src/uu/shuf/src/shuf.rs index c7e5689ab..5661ec2ee 100644 --- a/src/uu/shuf/src/shuf.rs +++ b/src/uu/shuf/src/shuf.rs @@ -357,7 +357,7 @@ impl Shufable for RangeInclusive { amount: u64, ) -> UResult>> { 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)) } } diff --git a/tests/by-util/test_shuf.rs b/tests/by-util/test_shuf.rs index 01c53c2e1..d9fd4d239 100644 --- a/tests/by-util/test_shuf.rs +++ b/tests/by-util/test_shuf.rs @@ -93,15 +93,15 @@ 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 = result + let result_seq: Vec = result .stdout_str() .split('\n') .filter(|x| !x.is_empty()) @@ -109,7 +109,7 @@ fn test_very_large_range() { .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() );