1 files changed, 42 insertions, 11 deletions

diff --git a/rust/src/tree_cache/mod.rs b/rust/src/tree_cache/mod.rs
index 719d0b2cf9..cb45e1f70d 100644
--- a/rust/src/tree_cache/mod.rs
+++ b/rust/src/tree_cache/mod.rs
@@ -54,16 +54,20 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
         }
     }
 
-    pub fn pop(
+    pub fn pop<Q>(
         &mut self,
-        current_key: &K,
-        mut next_keys: impl Iterator<Item = &'a K>,
-    ) -> Result<Option<TreeCacheNode<K, V>>, Error> {
+        current_key: Q,
+        mut next_keys: impl Iterator<Item = Q>,
+    ) -> Result<Option<TreeCacheNode<K, V>>, Error>
+    where
+        Q: Borrow<K>,
+        Q: Hash + Eq + 'a,
+    {
         if let Some(next_key) = next_keys.next() {
             match self {
                 TreeCacheNode::Leaf(_) => bail!("Given key is too long"),
                 TreeCacheNode::Branch(size, map) => {
-                    let node = if let Some(node) = map.get_mut(current_key) {
+                    let node = if let Some(node) = map.get_mut(current_key.borrow()) {
                         node
                     } else {
                         return Ok(None);
@@ -82,7 +86,7 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
             match self {
                 TreeCacheNode::Leaf(_) => bail!("Given key is too long"),
                 TreeCacheNode::Branch(size, map) => {
-                    if let Some(node) = map.remove(current_key) {
+                    if let Some(node) = map.remove(current_key.borrow()) {
                         *size -= node.len();
 
                         Ok(Some(node))
@@ -94,8 +98,8 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
         }
     }
 
-    pub fn items(&self) -> impl Iterator<Item = (Vec<&K>, &V)> {
-        let mut stack = vec![(vec![], self)];
+    pub fn items(&'a self) -> impl Iterator<Item = (Vec<&K>, &V)> {
+        let mut stack = vec![(Vec::new(), self)];
 
         std::iter::from_fn(move || {
             while let Some((prefix, node)) = stack.pop() {
@@ -116,6 +120,29 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
     }
 }
 
+impl<'a, K: Clone + Eq + Hash + 'a, V> TreeCacheNode<K, V> {
+    pub fn into_items(self) -> impl Iterator<Item = (Vec<K>, V)> {
+        let mut stack = vec![(Vec::new(), self)];
+
+        std::iter::from_fn(move || {
+            while let Some((prefix, node)) = stack.pop() {
+                match node {
+                    TreeCacheNode::Leaf(value) => return Some((prefix, value)),
+                    TreeCacheNode::Branch(_, map) => {
+                        stack.extend(map.into_iter().map(|(k, v)| {
+                            let mut prefix = prefix.clone();
+                            prefix.push(k);
+                            (prefix, v)
+                        }));
+                    }
+                }
+            }
+
+            None
+        })
+    }
+}
+
 impl<K, V> Default for TreeCacheNode<K, V> {
     fn default() -> Self {
         TreeCacheNode::new_branch()
@@ -182,10 +209,14 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCache<K, V> {
         }
     }
 
-    pub fn pop_node(
+    pub fn pop_node<Q>(
         &mut self,
-        key: impl IntoIterator<Item = &'a K>,
-    ) -> Result<Option<TreeCacheNode<K, V>>, Error> {
+        key: impl IntoIterator<Item = Q>,
+    ) -> Result<Option<TreeCacheNode<K, V>>, Error>
+    where
+        Q: Borrow<K>,
+        Q: Hash + Eq + 'a,
+    {
         let mut key_iter = key.into_iter();
 
         let k = if let Some(k) = key_iter.next() {