From 07c0b01b52a26351ae51c209d41d3ffe3851e740 Mon Sep 17 00:00:00 2001 From: oech3 <79379754+oech3@users.noreply.github.com> Date: Wed, 18 Feb 2026 02:13:25 +0900 Subject: [PATCH] shuf: Reduce malloc --- src/uu/shuf/src/nonrepeating_iterator.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/uu/shuf/src/nonrepeating_iterator.rs b/src/uu/shuf/src/nonrepeating_iterator.rs index 4cca9ea8c03..f300d85a374 100644 --- a/src/uu/shuf/src/nonrepeating_iterator.rs +++ b/src/uu/shuf/src/nonrepeating_iterator.rs @@ -3,9 +3,7 @@ // For the full copyright and license information, please view the LICENSE // file that was distributed with this source code. -// hijack HashMap for performance -type HashMap = std::collections::HashMap; - +use rustc_hash::FxHashMap; use std::ops::RangeInclusive; use uucore::error::UResult; @@ -47,12 +45,17 @@ pub(crate) struct NonrepeatingIterator<'a> { enum Values { Full(Vec), - Sparse(RangeInclusive, HashMap), + Sparse(RangeInclusive, FxHashMap), } impl<'a> NonrepeatingIterator<'a> { pub(crate) fn new(range: RangeInclusive, rng: &'a mut WrappedRng) -> Self { - let values = Values::Sparse(range, HashMap::default()); + 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), + ); NonrepeatingIterator { rng, values } } @@ -101,8 +104,7 @@ impl Iterator for NonrepeatingIterator<'_> { Values::Full(_) => (), Values::Sparse(range, _) if range.is_empty() => return None, Values::Sparse(range, items) => { - let range_len = range.size_hint().0 as u64; - if items.len() as u64 >= range_len / 8 { + if items.len() as u64 >= items.capacity() as u64 { self.values = Values::Full(hashmap_to_vec(range.clone(), items)); } } @@ -112,7 +114,7 @@ impl Iterator for NonrepeatingIterator<'_> { } } -fn hashmap_to_vec(range: RangeInclusive, map: &HashMap) -> Vec { +fn hashmap_to_vec(range: RangeInclusive, map: &FxHashMap) -> Vec { let lookup = |idx| *map.get(&idx).unwrap_or(&idx); range.rev().map(lookup).collect() }