Skip to content

Commit

Permalink
Improve the performance of HashMap iterators
Browse files Browse the repository at this point in the history
  • Loading branch information
@arthurprs
arthurprs committed Mar 16, 2024
1 parent f97eb5c commit da9ce50
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 108 deletions.
34 changes: 32 additions & 2 deletions benches/hashmap.rs
View file
Expand Up @@ -10,15 +10,15 @@ extern crate test;

use rand::{rngs::SmallRng, Rng, SeedableRng};
use std::iter::FromIterator;
use test::Bencher;
use test::{black_box, Bencher};

use imbl::hashmap::HashMap;

fn random_keys(size: usize) -> Vec<i64> {
let mut gen = SmallRng::from_entropy();
let mut set = Vec::new();
while set.len() < size {
let next = gen.gen::<i64>() % 10000;
let next = gen.gen::<i64>();
if !set.contains(&next) {
set.push(next);
}
Expand Down Expand Up @@ -253,3 +253,33 @@ fn hashmap_lookup_once_1000(b: &mut Bencher) {
fn hashmap_lookup_once_10000(b: &mut Bencher) {
hashmap_lookup_once_n(10000, b)
}

fn hashmap_iter_n(size: usize, b: &mut Bencher) {
let keys = random_keys(size);
let map: HashMap<i64, i64> = HashMap::from_iter(keys.into_iter().map(|i| (i, i)));
b.iter(|| {
for i in map.iter() {
black_box(i);
}
})
}

#[bench]
fn hashmap_iter_10(b: &mut Bencher) {
hashmap_iter_n(10, b)
}

#[bench]
fn hashmap_iter_100(b: &mut Bencher) {
hashmap_iter_n(100, b)
}

#[bench]
fn hashmap_iter_1000(b: &mut Bencher) {
hashmap_iter_n(1000, b)
}

#[bench]
fn hashmap_iter_10000(b: &mut Bencher) {
hashmap_iter_n(10000, b)
}
193 changes: 87 additions & 106 deletions src/nodes/hamt.rs
View file
Expand Up @@ -455,7 +455,6 @@ impl<A: HashValue> CollisionNode<A> {
pub(crate) struct Iter<'a, A> {
count: usize,
stack: Vec<ChunkIter<'a, Entry<A>, HASH_WIDTH>>,
current: ChunkIter<'a, Entry<A>, HASH_WIDTH>,
collision: Option<(HashBits, SliceIter<'a, A>)>,
}

Expand All @@ -465,7 +464,6 @@ impl<'a, A> Clone for Iter<'a, A> {
Self {
count: self.count,
stack: self.stack.clone(),
current: self.current.clone(),
collision: self.collision.clone(),
}
}
Expand All @@ -476,12 +474,13 @@ where
A: 'a,
{
pub(crate) fn new(root: &'a Node<A>, size: usize) -> Self {
Iter {
let mut result = Iter {
count: size,
stack: Vec::with_capacity((HASH_WIDTH / HASH_SHIFT) + 1),
current: root.data.iter(),
collision: None,
}
};
result.stack.push(root.data.iter());
result
}
}

Expand All @@ -492,43 +491,36 @@ where
type Item = (&'a A, HashBits);

fn next(&mut self) -> Option<Self::Item> {
if self.count == 0 {
return None;
}
if self.collision.is_some() {
'outer: loop {
if let Some((hash, ref mut coll)) = self.collision {
match coll.next() {
None => {}
None => self.collision = None,
Some(value) => {
self.count -= 1;
return Some((value, hash));
}
}
}
self.collision = None;
return self.next();
}
match self.current.next() {
Some(Entry::Value(value, hash)) => {
self.count -= 1;
Some((value, *hash))
}
Some(Entry::Node(child)) => {
let current = mem::replace(&mut self.current, child.data.iter());
self.stack.push(current);
self.next()
};
}
Some(Entry::Collision(coll)) => {
self.collision = Some((coll.hash, coll.data.iter()));
self.next()
}
None => match self.stack.pop() {
None => None,
Some(iter) => {
self.current = iter;
self.next()

while let Some(current) = self.stack.last_mut() {
match current.next() {
Some(Entry::Value(value, hash)) => {
self.count -= 1;
return Some((value, *hash));
}
Some(Entry::Node(child)) => {
self.stack.push(child.data.iter());
}
Some(Entry::Collision(coll)) => {
self.collision = Some((coll.hash, coll.data.iter()));
continue 'outer;
}
None => {
self.stack.pop();
}
}
},
}
return None;
}
}

Expand All @@ -547,7 +539,6 @@ pub(crate) struct IterMut<'a, A> {
count: usize,
pool: Pool<Node<A>>,
stack: Vec<ChunkIterMut<'a, Entry<A>, HASH_WIDTH>>,
current: ChunkIterMut<'a, Entry<A>, HASH_WIDTH>,
collision: Option<(HashBits, SliceIterMut<'a, A>)>,
}

Expand All @@ -556,13 +547,14 @@ where
A: 'a,
{
pub(crate) fn new(pool: &Pool<Node<A>>, root: &'a mut Node<A>, size: usize) -> Self {
IterMut {
let mut result = IterMut {
count: size,
pool: pool.clone(),
stack: Vec::with_capacity((HASH_WIDTH / HASH_SHIFT) + 1),
current: root.data.iter_mut(),
collision: None,
}
};
result.stack.push(root.data.iter_mut());
result
}
}

Expand All @@ -573,45 +565,38 @@ where
type Item = (&'a mut A, HashBits);

fn next(&mut self) -> Option<Self::Item> {
if self.count == 0 {
return None;
}
if self.collision.is_some() {
'outer: loop {
if let Some((hash, ref mut coll)) = self.collision {
match coll.next() {
None => {}
None => self.collision = None,
Some(value) => {
self.count -= 1;
return Some((value, hash));
}
}
};
}
self.collision = None;
return self.next();
}
match self.current.next() {
Some(Entry::Value(value, hash)) => {
self.count -= 1;
Some((value, *hash))
}
Some(Entry::Node(child_ref)) => {
let child = PoolRef::make_mut(&self.pool, child_ref);
let current = mem::replace(&mut self.current, child.data.iter_mut());
self.stack.push(current);
self.next()
}
Some(Entry::Collision(coll_ref)) => {
let coll = Ref::make_mut(coll_ref);
self.collision = Some((coll.hash, coll.data.iter_mut()));
self.next()
}
None => match self.stack.pop() {
None => None,
Some(iter) => {
self.current = iter;
self.next()

while let Some(current) = self.stack.last_mut() {
match current.next() {
Some(Entry::Value(value, hash)) => {
self.count -= 1;
return Some((value, *hash));
}
Some(Entry::Node(child_ref)) => {
let child = PoolRef::make_mut(&self.pool, child_ref);
self.stack.push(child.data.iter_mut());
}
Some(Entry::Collision(coll_ref)) => {
let coll = Ref::make_mut(coll_ref);
self.collision = Some((coll.hash, coll.data.iter_mut()));
continue 'outer;
}
None => {
self.stack.pop();
}
}
},
}
return None;
}
}

Expand All @@ -630,19 +615,19 @@ pub(crate) struct Drain<A> {
count: usize,
pool: Pool<Node<A>>,
stack: Vec<PoolRef<Node<A>>>,
current: PoolRef<Node<A>>,
collision: Option<CollisionNode<A>>,
}

impl<A> Drain<A> {
pub(crate) fn new(pool: &Pool<Node<A>>, root: PoolRef<Node<A>>, size: usize) -> Self {
Drain {
let mut result = Drain {
count: size,
pool: pool.clone(),
stack: vec![],
current: root,
stack: Vec::with_capacity((HASH_WIDTH / HASH_SHIFT) + 1),
collision: None,
}
};
result.stack.push(root);
result
}
}

Expand All @@ -653,40 +638,36 @@ where
type Item = (A, HashBits);

fn next(&mut self) -> Option<Self::Item> {
if self.count == 0 {
return None;
}
if self.collision.is_some() {
if let Some(ref mut coll) = self.collision {
if let Some(value) = coll.data.pop() {
self.count -= 1;
return Some((value, coll.hash));
}
}
self.collision = None;
return self.next();
}
match PoolRef::make_mut(&self.pool, &mut self.current).data.pop() {
Some(Entry::Value(value, hash)) => {
self.count -= 1;
Some((value, hash))
}
Some(Entry::Collision(coll_ref)) => {
self.collision = Some(clone_ref(coll_ref));
self.next()
}
Some(Entry::Node(child)) => {
let parent = mem::replace(&mut self.current, child);
self.stack.push(parent);
self.next()
'outer: loop {
if let Some(coll) = &mut self.collision {
match coll.data.pop() {
None => self.collision = None,
Some(value) => {
self.count -= 1;
return Some((value, coll.hash));
}
};
}
None => match self.stack.pop() {
None => None,
Some(parent) => {
self.current = parent;
self.next()

while let Some(current) = self.stack.last_mut() {
match PoolRef::make_mut(&self.pool, current).data.pop() {
Some(Entry::Value(value, hash)) => {
self.count -= 1;
return Some((value, hash));
}
Some(Entry::Node(child)) => {
self.stack.push(child);
}
Some(Entry::Collision(coll)) => {
self.collision = Some(clone_ref(coll));
continue 'outer;
}
None => {
self.stack.pop();
}
}
},
}
return None;
}
}

Expand Down

0 comments on commit da9ce50

Please sign in to comment.