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<u32, u32> = 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<K, V> {
}
pub fn items(&'a self) -> impl Iterator<Item = (Vec<&K>, &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<K, V> {
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<K, V> {
None
})
}
+
+ pub fn values(&'a self) -> impl Iterator<Item = &V> {
+ 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<K, V> {
|
