nav: centralize find_leaf_for_key in tree_structure; refactor get/get_mut to use it; inline for hot path

Author: KentBeck

rust/src/get_operations.rs

@@ -32,8 +32,16 @@ impl<K: Ord + Clone, V: Clone> BPlusTreeMap<K, V> {
     /// assert_eq!(tree.get(&2), None);
     /// ```
     pub fn get(&self, key: &K) -> Option<&V> {
-        let node = &self.root;
-        self.get_recursive(node, key)
+        let (leaf_id, index) = self.find_leaf_for_key(key)?;
+        let leaf = self.get_leaf(leaf_id)?;
+        if index < leaf.keys_len() {
+            if let Some(k) = leaf.get_key(index) {
+                if k == key {
+                    return leaf.get_value(index);
+                }
+            }
+        }
+        None
     }
 
     /// Check if key exists in the tree.
@@ -133,8 +141,16 @@ impl<K: Ord + Clone, V: Clone> BPlusTreeMap<K, V> {
     /// assert_eq!(tree.get(&1), Some(&"ONE"));
     /// ```
     pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
-        let root = self.root;
-        self.get_mut_recursive(&root, key)
+        let (leaf_id, index) = self.find_leaf_for_key(key)?;
+        let leaf = self.get_leaf_mut(leaf_id)?;
+        if index < leaf.keys_len() {
+            if let Some(k) = leaf.get_key(index) {
+                if k == key {
+                    return leaf.get_value_mut(index);
+                }
+            }
+        }
+        None
     }
 
     /// Try to get a value, returning detailed error context on failure.
@@ -204,28 +220,7 @@ impl<K: Ord + Clone, V: Clone> BPlusTreeMap<K, V> {
     // PRIVATE HELPER METHODS FOR GET OPERATIONS
     // ============================================================================
 
-    /// Recursively search for a key in the tree.
-    #[inline]
-    fn get_recursive<'a>(&'a self, node: &'a NodeRef<K, V>, key: &K) -> Option<&'a V> {
-        match node {
-            NodeRef::Leaf(id, _) => self.get_leaf(*id).and_then(|leaf| leaf.get(key)),
-            NodeRef::Branch(id, _) => self
-                .get_branch(*id)
-                .and_then(|branch| branch.get_child(key))
-                .and_then(|child| self.get_recursive(child, key)),
-        }
-    }
-
-    /// Get mutable reference recursively.
-    fn get_mut_recursive(&mut self, node: &NodeRef<K, V>, key: &K) -> Option<&mut V> {
-        match node {
-            NodeRef::Leaf(id, _) => self.get_leaf_mut(*id).and_then(|leaf| leaf.get_mut(key)),
-            NodeRef::Branch(id, _) => {
-                let (_child_index, child_ref) = self.get_child_for_key(*id, key)?;
-                self.get_mut_recursive(&child_ref, key)
-            }
-        }
-    }
+    // Removed old recursive get helpers in favor of direct leaf-position lookup
 
     /// Helper to get child info for a key in a branch.
     #[inline]

rust/src/range_queries.rs

@@ -4,7 +4,7 @@
 //! bounds resolution, and range optimization algorithms.
 
 use crate::iteration::RangeIterator;
-use crate::types::{BPlusTreeMap, NodeId, NodeRef};
+use crate::types::{BPlusTreeMap, NodeId};
 use std::ops::{Bound, RangeBounds};
 
 /// Type alias for complex range analysis result
@@ -75,8 +75,8 @@ impl<K: Ord + Clone, V: Clone> BPlusTreeMap<K, V> {
     {
         // Optimize start bound resolution - eliminate redundant Option handling
         let (start_info, skip_first) = match range.start_bound() {
-            Bound::Included(key) => (self.find_range_start(key), false),
-            Bound::Excluded(key) => (self.find_range_start(key), true),
+            Bound::Included(key) => (self.find_leaf_for_key(key), false),
+            Bound::Excluded(key) => (self.find_leaf_for_key(key), true),
             Bound::Unbounded => (self.get_first_leaf_id().map(|id| (id, 0)), false),
         };
 
@@ -90,45 +90,6 @@ impl<K: Ord + Clone, V: Clone> BPlusTreeMap<K, V> {
         (start_info, skip_first, end_info)
     }
 
-    /// Find the starting position for a range query.
-    fn find_range_start(&self, key: &K) -> Option<(NodeId, usize)> {
-        self.find_leaf_for_key(key)
-    }
-
-    /// Find the leaf node and index where a key should be located.
-    fn find_leaf_for_key(&self, key: &K) -> Option<(NodeId, usize)> {
-        let mut current = &self.root;
-
-        loop {
-            match current {
-                NodeRef::Leaf(leaf_id, _) => {
-                    if let Some(leaf) = self.get_leaf(*leaf_id) {
-                        // Find the position where this key would be inserted
-                        let index = match leaf.binary_search_keys(key) {
-                            Ok(idx) => idx,  // Key found at exact position
-                            Err(idx) => idx, // Key would be inserted at this position
-                        };
-                        return Some((*leaf_id, index));
-                    } else {
-                        return None;
-                    }
-                }
-                NodeRef::Branch(branch_id, _) => {
-                    if let Some(branch) = self.get_branch(*branch_id) {
-                        let child_index = branch.find_child_index(key);
-                        if let Some(child) = branch.children.get(child_index) {
-                            current = child;
-                        } else {
-                            return None;
-                        }
-                    } else {
-                        return None;
-                    }
-                }
-            }
-        }
-    }
-
     // ============================================================================
     // RANGE OPTIMIZATION HELPERS
     // ============================================================================

rust/src/tree_structure.rs

@@ -138,6 +138,42 @@ impl<K: Ord + Clone, V: Clone> BPlusTreeMap<K, V> {
         }
     }
 
+    /// Find the leaf node and index where a key should be located.
+    /// Returns the leaf `NodeId` and the insertion index within that leaf.
+    #[inline]
+    pub(crate) fn find_leaf_for_key(&self, key: &K) -> Option<(NodeId, usize)> {
+        let mut current = &self.root;
+
+        loop {
+            match current {
+                NodeRef::Leaf(leaf_id, _) => {
+                    if let Some(leaf) = self.get_leaf(*leaf_id) {
+                        // Find the position where this key would be inserted
+                        let index = match leaf.binary_search_keys(key) {
+                            Ok(idx) => idx,  // Key found at exact position
+                            Err(idx) => idx, // Key would be inserted at this position
+                        };
+                        return Some((*leaf_id, index));
+                    } else {
+                        return None;
+                    }
+                }
+                NodeRef::Branch(branch_id, _) => {
+                    if let Some(branch) = self.get_branch(*branch_id) {
+                        let child_index = branch.find_child_index(key);
+                        if let Some(child) = branch.children.get(child_index) {
+                            current = child;
+                        } else {
+                            return None;
+                        }
+                    } else {
+                        return None;
+                    }
+                }
+            }
+        }
+    }
+
     // Arena statistics and management methods moved to arena.rs module
 
     // ============================================================================