refactor(tree): simplify tree mutation and entry handling

- Introduced `SubtreeMut` for encapsulating mutable subtree operations
- Simplified `Tree::entry` by delegating to `SubtreeMut::entry`
- Removed redundant and commented-out code for clarity
- Added `Handle::value_mut` for mutable access to node values

src/tree.rs

@@ -623,7 +623,6 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, K, V, Type> {
 }
 
 impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Immut<'a>, K, V, NodeType>, marker::Value> {
-    /// SAFETY: the caller must ensure that `idx < node.len()`
     pub(super) unsafe fn into_value(self) -> &'a V {
         let leaf = self.node.into_leaf();
         let v = leaf.value.as_ref().unwrap();
@@ -631,6 +630,14 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Immut<'a>, K, V, NodeTyp
     }
 }
 
+impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::Value> {
+    pub(super) unsafe fn value_mut(&mut self) -> &mut V {
+        let leaf = self.node.as_leaf_mut();
+        let v = leaf.value.as_mut().unwrap();
+        v
+    }
+}
+
 impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
     pub(super) fn forget_node_type(
         self,
@@ -1053,6 +1060,11 @@ mod search {
     }
 }
 
+enum HandleOrTree<'a, BorrowType, K, V, NodeType, HandleType> {
+    Handle(Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType>),
+    Tree(borrow::DormantMutRef<'a, Tree<K, V>>),
+}
+
 mod entry {
     use core::marker::PhantomData;
 
@@ -1087,7 +1099,7 @@ mod entry {
         fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
             f.debug_struct("VacantEntry")
                 .field("key", &self.key)
-                .field("handle", &self.handle)
+                // .field("handle", &self.handle)
                 .finish()
         }
     }
@@ -1107,15 +1119,13 @@ mod entry {
     {
         pub(super) key: Q,
         pub(super) handle:
-            Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Edge>>,
+            super::HandleOrTree<'a, marker::Mut<'a>, K, V, marker::LeafOrInternal, marker::Edge>,
-        pub(super) dormant_map: DormantMutRef<'a, Tree<K, V>>,
         pub(super) _marker: PhantomData<&'a mut (K, V)>,
     }
 
     pub struct OccupiedEntry<'a, K, V> {
         pub(super) handle:
             Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Value>,
-        pub(super) dormant_map: DormantMutRef<'a, Tree<K, V>>,
         pub(super) _marker: PhantomData<&'a mut (K, V)>,
     }
 
@@ -1123,6 +1133,9 @@ mod entry {
         pub fn get(&self) -> &V {
             unsafe { self.handle.reborrow().into_value() }
         }
+        pub fn get_mut(&mut self) -> &mut V {
+            unsafe { self.handle.value_mut() }
+        }
     }
 
     impl<'a, Q, K: Ord, V> Entry<'a, Q, K, V>
@@ -1141,29 +1154,93 @@ mod entry {
     where
         Q: Iterator<Item = K>,
     {
-        pub fn insert_entry(mut self, value: V) -> OccupiedEntry<'a, K, V> {
+        pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V> {
+            use super::HandleOrTree;
             let handle = match self.handle {
                 // no root node yet
-                None => {
+                HandleOrTree::Tree(mut tree) => {
-                    let tree = unsafe { self.dormant_map.reborrow() };
+                    // SAFETY: there are no nodes in the tree yet
+                    let tree = unsafe { tree.reborrow() };
                     let root = tree.root.insert(NodeRef::new_internal());
 
                     unsafe {
                         Handle::new_edge(root.borrow_mut(), 0).insert_recursing(self.key, value)
                     }
                 }
-                Some(handle) => unsafe { handle.insert_recursing(self.key, value) },
+                HandleOrTree::Handle(handle) => unsafe { handle.insert_recursing(self.key, value) },
             };
 
             OccupiedEntry {
                 handle,
-                dormant_map: self.dormant_map,
                 _marker: PhantomData,
             }
         }
     }
 }
 
+mod subtree {
+    use core::marker::PhantomData;
+
+    use crate::tree::{HandleOrTree, search};
+
+    use super::{
+        NodeRef, OnceAndIter,
+        entry::{Entry, OccupiedEntry, VacantEntry},
+        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)>,
+    }
+
+    impl<'a, K, V> SubtreeMut<'a, K, V>
+    where
+        K: Ord,
+    {
+        pub fn get<Q>(&self, mut key_seq: Q) -> Option<&V>
+        where
+            Q: Iterator<Item = K>,
+        {
+            let root = self.root.reborrow();
+            match root.search_tree(&mut key_seq) {
+                search::SearchResult::Found(handle) => Some(unsafe { handle.into_value() }),
+                _ => {
+                    // 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>,
+        {
+            use Entry::*;
+
+            // SAFETY: this is actually our borrow?
+            let root = unsafe { self.root.reborrow_mut() };
+
+            match root.search_tree(&mut key_seq) {
+                search::SearchResult::Found(handle) => Occupied(OccupiedEntry {
+                    handle,
+                    _marker: PhantomData,
+                }),
+                search::SearchResult::GoDown(handle) => Vacant(VacantEntry {
+                    key: key_seq.into(),
+                    handle: HandleOrTree::Handle(handle.forget_node_type()),
+                    _marker: PhantomData,
+                }),
+                search::SearchResult::Insert(key, handle) => Vacant(VacantEntry {
+                    key: OnceAndIter::once(key, key_seq),
+                    handle: HandleOrTree::Handle(handle),
+                    _marker: PhantomData,
+                }),
+            }
+        }
+    }
+}
+
 mod borrow {
     use core::{marker::PhantomData, ptr::NonNull};
 
@@ -1361,89 +1438,51 @@ impl<K, V> Tree<K, V>
 where
     K: Ord,
 {
-    // pub fn closest_entry<Q>(
+    fn as_subtree_mut<'a>(&'a mut self) -> Option<subtree::SubtreeMut<'a, K, V>> {
-    //     &mut self,
+        let root = self
-    //     mut key_seq: Q,
+            .root
-    // ) -> Option<entry::Entry<OnceAndIter<Q, K>, K, V>>
+            .as_mut()?
-    // where
+            .borrow_mut()
-    //     Q: Iterator<Item = K>,
+            .dormant()
-    // {
+            .as_leaf_or_internal();
-    //     use borrow::DormantMutRef;
-    //     use entry::Entry;
 
-    //     let (tree, dormant) = DormantMutRef::new(self);
+        Some(subtree::SubtreeMut {
+            root: unsafe { root.awaken() },
+            _marker: PhantomData,
+        })
+    }
 
-    //     match tree.root {
+    pub fn entry<'a, Q>(&'a mut self, key_seq: Q) -> entry::Entry<'a, OnceAndIter<Q, K>, K, V>
-    //         None => None,
-    //         Some(ref mut root) => match root.borrow_mut().forget_type().search_tree(&mut key_seq) {
-    //             search::SearchResult::Found(handle) => {
-    //                 Some(Entry::Occupied(entry::OccupiedEntry {
-    //                     handle,
-    //                     dormant_map: dormant,
-    //                     _marker: PhantomData,
-    //                 }))
-    //             }
-    //             search::SearchResult::GoDown(handle) => match key_seq.next() {
-    //                 None => Some(Entry::Occupied(entry::OccupiedEntry {
-    //                     handle: handle.forget_node_type().into_value(),
-    //                     dormant_map: dormant,
-    //                     _marker: PhantomData,
-    //                 })),
-    //                 Some(key) => Some(Entry::Vacant(entry::VacantEntry {
-    //                     key: OnceAndIter::once(key, key_seq),
-    //                     handle: Some(handle.forget_node_type()),
-    //                     dormant_map: dormant,
-    //                     _marker: PhantomData,
-    //                 })),
-    //             },
-    //             search::SearchResult::Insert(key, handle) => {
-    //                 Some(Entry::Vacant(entry::VacantEntry {
-    //                     key: OnceAndIter::once(key, key_seq),
-    //                     handle: Some(handle),
-    //                     dormant_map: dormant,
-    //                     _marker: PhantomData,
-    //                 }))
-    //             }
-    //         },
-    //     }
-    // }
-
-    pub fn entry<Q>(&'_ mut self, mut key_seq: Q) -> entry::Entry<'_, OnceAndIter<Q, K>, K, V>
     where
         Q: Iterator<Item = K>,
     {
         use borrow::DormantMutRef;
-        use entry::{Entry::*, OccupiedEntry, VacantEntry};
+        use entry::{Entry::*, VacantEntry};
 
         let (tree, dormant) = DormantMutRef::new(self);
 
-        match tree.root {
+        let entry = match tree.as_subtree_mut() {
+            Some(mut subtree) => {
+                let entry = subtree.entry(key_seq);
+                // SAFETY: extending the lifetime is fine because we borrow the tree for 'a,
+                // and no references to the subtree are live after this.
+                // The same could be achieved using `dormant.reborrow()` a bunch
+                // of times while destructuring the entry.
+                unsafe {
+                    core::mem::transmute::<
+                        entry::Entry<'_, OnceAndIter<Q, K>, K, V>,
+                        entry::Entry<'a, OnceAndIter<Q, K>, K, V>,
+                    >(entry)
+                }
+            }
             None => Vacant(VacantEntry {
                 key: key_seq.into(),
-                handle: None,
+                handle: HandleOrTree::Tree(dormant),
-                dormant_map: dormant,
                 _marker: PhantomData,
             }),
-            Some(ref mut root) => match root.borrow_mut().forget_type().search_tree(&mut key_seq) {
+        };
-                search::SearchResult::Found(handle) => Occupied(OccupiedEntry {
+
-                    handle,
+        entry
-                    dormant_map: dormant,
-                    _marker: PhantomData,
-                }),
-                search::SearchResult::GoDown(handle) => Vacant(VacantEntry {
-                    key: key_seq.into(),
-                    handle: Some(handle.forget_node_type()),
-                    dormant_map: dormant,
-                    _marker: PhantomData,
-                }),
-                search::SearchResult::Insert(key, handle) => Vacant(VacantEntry {
-                    key: OnceAndIter::once(key, key_seq),
-                    handle: Some(handle),
-                    dormant_map: dormant,
-                    _marker: PhantomData,
-                }),
-            },
-        }
     }
 }