diff --git a/rust/benches/tree_cache.rs b/rust/benches/tree_cache.rs index 191550ad9b..3061968afc 100644 --- a/rust/benches/tree_cache.rs +++ b/rust/benches/tree_cache.rs @@ -58,3 +58,20 @@ fn bench_tree_cache_length(b: &mut Bencher) { b.iter(|| cache.len()); } + +#[bench] +fn tree_cache_iterate(b: &mut Bencher) { + let mut cache: TreeCache = TreeCache::new(); + + for c1 in 0..=10 { + for c2 in 0..=10 { + for c3 in 0..=10 { + for c4 in 0..=10 { + cache.set([c1, c2, c3, c4], 1).unwrap() + } + } + } + } + + b.iter(|| cache.items().count()); +} diff --git a/rust/src/tree_cache/mod.rs b/rust/src/tree_cache/mod.rs index cb45e1f70d..0508895869 100644 --- a/rust/src/tree_cache/mod.rs +++ b/rust/src/tree_cache/mod.rs @@ -99,7 +99,11 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode { } pub fn items(&'a self) -> impl Iterator, &V)> { - let mut stack = vec![(Vec::new(), self)]; + // To avoid a lot of mallocs we guess the length of the key. Ideally + // we'd know this. + let capacity_guesstimate = 10; + + let mut stack = vec![(Vec::with_capacity(capacity_guesstimate), self)]; std::iter::from_fn(move || { while let Some((prefix, node)) = stack.pop() { @@ -107,9 +111,10 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode { TreeCacheNode::Leaf(value) => return Some((prefix, value)), TreeCacheNode::Branch(_, map) => { stack.extend(map.iter().map(|(k, v)| { - let mut prefix = prefix.clone(); - prefix.push(k); - (prefix, v) + let mut new_prefix = Vec::with_capacity(capacity_guesstimate); + new_prefix.extend_from_slice(&prefix); + new_prefix.push(k); + (new_prefix, v) })); } } @@ -118,6 +123,23 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode { None }) } + + pub fn values(&'a self) -> impl Iterator { + let mut stack = vec![self]; + + std::iter::from_fn(move || { + while let Some(node) = stack.pop() { + match node { + TreeCacheNode::Leaf(value) => return Some(value), + TreeCacheNode::Branch(_, map) => { + stack.extend(map.iter().map(|(_k, v)| v)); + } + } + } + + None + }) + } } impl<'a, K: Clone + Eq + Hash + 'a, V> TreeCacheNode {