Skip to content

Commit 74b8924

Browse files
oech3cakebaker
authored andcommitted
shuf: Reduce malloc
1 parent 2923d1b commit 74b8924

1 file changed

Lines changed: 10 additions & 8 deletions

File tree

src/uu/shuf/src/nonrepeating_iterator.rs

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,7 @@
33
// For the full copyright and license information, please view the LICENSE
44
// file that was distributed with this source code.
55

6-
// hijack HashMap for performance
7-
type HashMap<K, V> = std::collections::HashMap<K, V, rustc_hash::FxBuildHasher>;
8-
6+
use rustc_hash::FxHashMap;
97
use std::ops::RangeInclusive;
108

119
use uucore::error::UResult;
@@ -47,12 +45,17 @@ pub(crate) struct NonrepeatingIterator<'a> {
4745

4846
enum Values {
4947
Full(Vec<u64>),
50-
Sparse(RangeInclusive<u64>, HashMap<u64, u64>),
48+
Sparse(RangeInclusive<u64>, FxHashMap<u64, u64>),
5149
}
5250

5351
impl<'a> NonrepeatingIterator<'a> {
5452
pub(crate) fn new(range: RangeInclusive<u64>, rng: &'a mut WrappedRng) -> Self {
55-
let values = Values::Sparse(range, HashMap::default());
53+
const MAX_CAPACITY: usize = 128; // todo: optimize this
54+
let capacity = (range.size_hint().0).min(MAX_CAPACITY);
55+
let values = Values::Sparse(
56+
range,
57+
FxHashMap::with_capacity_and_hasher(capacity, rustc_hash::FxBuildHasher),
58+
);
5659
NonrepeatingIterator { rng, values }
5760
}
5861

@@ -101,8 +104,7 @@ impl Iterator for NonrepeatingIterator<'_> {
101104
Values::Full(_) => (),
102105
Values::Sparse(range, _) if range.is_empty() => return None,
103106
Values::Sparse(range, items) => {
104-
let range_len = range.size_hint().0 as u64;
105-
if items.len() as u64 >= range_len / 8 {
107+
if items.len() as u64 >= items.capacity() as u64 {
106108
self.values = Values::Full(hashmap_to_vec(range.clone(), items));
107109
}
108110
}
@@ -112,7 +114,7 @@ impl Iterator for NonrepeatingIterator<'_> {
112114
}
113115
}
114116

115-
fn hashmap_to_vec(range: RangeInclusive<u64>, map: &HashMap<u64, u64>) -> Vec<u64> {
117+
fn hashmap_to_vec(range: RangeInclusive<u64>, map: &FxHashMap<u64, u64>) -> Vec<u64> {
116118
let lookup = |idx| *map.get(&idx).unwrap_or(&idx);
117119
range.rev().map(lookup).collect()
118120
}

0 commit comments

Comments
 (0)