From affa52ba738713ba2fcabb7a48aa8007c0d9d252 Mon Sep 17 00:00:00 2001
From: Janis <janis@nirgendwo.xyz>
Date: Mon, 11 Aug 2025 22:43:34 +0200
Subject: [PATCH] feat(tree): add mutable subtree access and traversal methods

- Introduced `into_value_mut` for mutable value access in `Handle`
- Added `into_subtree` for converting OccupiedEntries into mutable subtrees
- Refactored `SubtreeMut` into a generic `Subtree` supporting borrow types
- Implemented `get_subtree` and `get_subtree_mut` for subtree traversal
- Updated `Tree` to use the new `Subtree` abstraction
---
 src/tree.rs | 77 ++++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 70 insertions(+), 7 deletions(-)

diff --git a/src/tree.rs b/src/tree.rs
index 863e452..e31dee5 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -636,6 +636,11 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>
         let v = leaf.value.as_mut().unwrap();
         v
     }
+
+    pub fn into_value_mut(self) -> &'a mut V {
+        let leaf = self.node.into_leaf_mut();
+        leaf.value.as_mut().unwrap()
+    }
 }
 
 impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
@@ -1136,6 +1141,12 @@ mod entry {
         pub fn get_mut(&mut self) -> &mut V {
             unsafe { self.handle.value_mut() }
         }
+
+        pub fn into_subtree(self) -> super::subtree::Subtree<K, V, marker::Mut<'a>> {
+            super::subtree::Subtree {
+                root: self.handle.node.as_leaf_or_internal(),
+            }
+        }
     }
 
     impl<'a, Q, K: Ord, V> Entry<'a, Q, K, V>
@@ -1189,12 +1200,11 @@ mod subtree {
         marker,
     };
 
-    pub struct SubtreeMut<'a, K, V> {
-        pub(super) root: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
-        pub(super) _marker: core::marker::PhantomData<&'a mut (K, V)>,
+    pub struct Subtree<K, V, BorrowType> {
+        pub(super) root: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
     }
 
-    impl<'a, K, V> SubtreeMut<'a, K, V>
+    impl<K, V, BorrowType> Subtree<K, V, BorrowType>
     where
         K: Ord,
     {
@@ -1212,6 +1222,60 @@ mod subtree {
             }
         }
 
+        pub fn get_subtree<Q>(&self, mut key_seq: Q) -> Option<Subtree<K, V, marker::Immut<'_>>>
+        where
+            Q: Iterator<Item = K>,
+        {
+            let root = self.root.reborrow();
+            match root.search_tree(&mut key_seq) {
+                search::SearchResult::Found(handle) => Some(Subtree {
+                    root: handle.node.as_leaf_or_internal(),
+                }),
+                _ => {
+                    // key not found
+                    None
+                }
+            }
+        }
+    }
+
+    impl<'a, K, V> Subtree<K, V, marker::Mut<'a>>
+    where
+        K: Ord,
+    {
+        pub fn get_mut<Q>(&mut self, mut key_seq: Q) -> Option<&mut V>
+        where
+            Q: Iterator<Item = K>,
+        {
+            let root = unsafe { self.root.reborrow_mut() };
+            match root.search_tree(&mut key_seq) {
+                search::SearchResult::Found(handle) => Some(handle.into_value_mut()),
+                _ => {
+                    // key not found
+                    None
+                }
+            }
+        }
+
+        pub fn get_subtree_mut<Q>(
+            &'_ mut self,
+            mut key_seq: Q,
+        ) -> Option<Subtree<K, V, marker::Mut<'_>>>
+        where
+            Q: Iterator<Item = K>,
+        {
+            let root = unsafe { self.root.reborrow_mut() };
+            match root.search_tree(&mut key_seq) {
+                search::SearchResult::Found(handle) => Some(Subtree {
+                    root: handle.node.as_leaf_or_internal(),
+                }),
+                _ => {
+                    // key not found
+                    None
+                }
+            }
+        }
+
         pub fn entry<Q>(&'_ mut self, mut key_seq: Q) -> Entry<'_, OnceAndIter<Q, K>, K, V>
         where
             Q: Iterator<Item = K>,
@@ -1438,7 +1502,7 @@ impl<K, V> Tree<K, V>
 where
     K: Ord,
 {
-    fn as_subtree_mut<'a>(&'a mut self) -> Option<subtree::SubtreeMut<'a, K, V>> {
+    fn as_subtree_mut<'a>(&'a mut self) -> Option<subtree::Subtree<K, V, marker::Mut<'a>>> {
         let root = self
             .root
             .as_mut()?
@@ -1446,9 +1510,8 @@ where
             .dormant()
             .as_leaf_or_internal();
 
-        Some(subtree::SubtreeMut {
+        Some(subtree::Subtree {
             root: unsafe { root.awaken() },
-            _marker: PhantomData,
         })
     }