summary refs log tree commit diff
path: root/rust/src/tree_cache/mod.rs
diff options
context:
space:
mode:
Diffstat (limited to 'rust/src/tree_cache/mod.rs')
-rw-r--r--rust/src/tree_cache/mod.rs30
1 files changed, 26 insertions, 4 deletions
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> {