Optimize lookup for any descendant leaf

src/raw/representation.rs

@@ -8,7 +8,7 @@ use core::{
     ops::{Bound, RangeBounds, RangeInclusive},
 };
 
-use crate::{raw::minimum_unchecked, rust_nightly_apis::likely, AsBytes};
+use crate::{rust_nightly_apis::likely, AsBytes};
 
 mod header;
 pub use header::*;
@@ -472,16 +472,12 @@ pub unsafe trait InnerNodeCommon<K, V, const PREFIX_LEN: usize>:
         if likely!(len <= PREFIX_LEN) {
             (header.read_prefix(), None)
         } else {
-            // SAFETY: By construction a InnerNode, must have >= 1 children, this
-            // is even more strict since in the case of 1 child the node can be
-            // collapsed, so a InnerNode must have >= 2 children, so it's safe
-            // to search for the minium. And the same applies to the `minimum_unchecked`
-            // function
-            let (_, min_child) = self.min();
-            let leaf_ptr = unsafe { minimum_unchecked(min_child) };
-
-            // SAFETY: Since have a shared reference
-            // is safe to create a shared reference from it
+            // SAFETY: Given we have a shared reference to this leaf node, there must be no
+            // concurrent mutation on this node or any child node of this node.
+            let leaf_ptr = unsafe { any_subtree_leaf(self) };
+
+            // SAFETY: Since have a shared reference is safe to create a shared reference
+            // from it
             let leaf = unsafe { leaf_ptr.as_ref() };
             let leaf = leaf.key_ref().as_bytes();
 
@@ -519,6 +515,16 @@ pub unsafe trait InnerNodeCommon<K, V, const PREFIX_LEN: usize>:
     /// because if we had, no children this current node should have been
     /// deleted.
     fn max(&self) -> (u8, OpaqueNodePtr<K, V, PREFIX_LEN>);
+
+    /// Returns any child pointer from this node, preferring children that are
+    /// pointers to leaf nodes.
+    ///
+    /// This function is useful in context where we need to quickly find a leaf
+    /// of a subtree, but don't care how we get there.
+    fn any_child_prefer_leaf(&self) -> OpaqueNodePtr<K, V, PREFIX_LEN> {
+        // The minimum child is a valid choice
+        self.min().1
+    }
 }
 
 /// Common methods implemented by all inner nodes that will be used in the tree.
@@ -807,6 +813,33 @@ fn assert_valid_range_bounds(bound: &impl RangeBounds<u8>) {
     }
 }
 
+/// Find any leaf in the subtree rooted by `root`.
+///
+/// # Safety
+///  - This function cannot be called concurrently with any mutating operation
+///    on `root` or any child node of `root`. This function will arbitrarily
+///    read to any child in the given tree.
+#[inline]
+unsafe fn any_subtree_leaf<N, K, V, const PREFIX_LEN: usize>(
+    root: &N,
+) -> NodePtr<PREFIX_LEN, LeafNode<K, V, PREFIX_LEN>>
+where
+    N: InnerNodeCommon<K, V, PREFIX_LEN>,
+{
+    let mut current_node = root.any_child_prefer_leaf();
+
+    loop {
+        current_node = match_concrete_node_ptr!(match (current_node.to_node_ptr()) {
+            // SAFETY: No other concurrent mutation will happen, the reference returned from
+            // `as_ref` is also bounded to this loop, not returned outside.
+            InnerNode(inner_node) => unsafe { inner_node.as_ref().any_child_prefer_leaf() },
+            LeafNode(leaf_node) => {
+                return leaf_node;
+            },
+        });
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use alloc::{boxed::Box, vec::Vec};

src/raw/representation/inner_node_indirect.rs

@@ -11,7 +11,7 @@ use self::index::NonMaxIndex;
 use crate::{
     raw::{
         representation::assert_valid_range_bounds, Header, InnerNode, InnerNode16, InnerNodeCommon,
-        InnerNodeDirect, InnerNodeSorted, Node, NodeType, OpaqueNodePtr,
+        InnerNodeDirect, InnerNodeSorted, LeafNode, Node, NodeType, OpaqueNodePtr,
     },
     rust_nightly_apis::maybe_uninit_slice_assume_init_ref,
 };
@@ -444,6 +444,22 @@ unsafe impl<K, V, const PREFIX_LEN: usize, const SIZE: usize> InnerNodeCommon<K,
         }
         unreachable!("inner node must have non-zero number of children");
     }
+
+    fn any_child_prefer_leaf(&self) -> OpaqueNodePtr<K, V, PREFIX_LEN> {
+        let children = self.initialized_child_pointers();
+
+        // SAFETY: Since inner nodes must have at least two children, and the
+        // `InnerNodeIndirect` maintains a compact array of child pointers, this must be
+        // a valid access.
+        let first = unsafe { children.get_unchecked(0) };
+        if first.is::<LeafNode<K, V, PREFIX_LEN>>() {
+            return *first;
+        }
+
+        // SAFETY: Same as above, just for the second child
+        let second = unsafe { children.get_unchecked(1) };
+        *second
+    }
 }
 
 /// Node that references between 17 and 49 children.

src/raw/representation/inner_node_sorted.rs

@@ -14,7 +14,7 @@ use core::{
 use crate::{
     raw::{
         representation::assert_valid_range_bounds, Header, InnerNode, InnerNode48, InnerNodeCommon,
-        InnerNodeDirect, InnerNodeIndirect, Node, NodeType, OpaqueNodePtr,
+        InnerNodeDirect, InnerNodeIndirect, LeafNode, Node, NodeType, OpaqueNodePtr,
     },
     rust_nightly_apis::maybe_uninit_slice_assume_init_ref,
 };
@@ -547,6 +547,22 @@ unsafe impl<K, V, const PREFIX_LEN: usize, const SIZE: usize> InnerNodeCommon<K,
             )
         }
     }
+
+    fn any_child_prefer_leaf(&self) -> OpaqueNodePtr<K, V, PREFIX_LEN> {
+        let (_, children) = self.initialized_portion();
+        // SAFETY: Any Node4 must have at least 2 children, so the `keys` and
+        // `children` arrays must be non-empty
+        unsafe {
+            let first = *children.get_unchecked(0);
+            if first.is::<LeafNode<K, V, PREFIX_LEN>>() {
+                return first;
+            }
+
+            // Must have at least two children right
+            let second = *children.get_unchecked(1);
+            second
+        }
+    }
 }
 
 /// Node that references between 2 and 4 children