diff --git a/src/raw/representation.rs b/src/raw/representation.rs index 519c4a44..accf9b25 100644 --- a/src/raw/representation.rs +++ b/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: 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: /// because if we had, no children this current node should have been /// deleted. fn max(&self) -> (u8, OpaqueNodePtr); + + /// 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 { + // 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) { } } +/// 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( + root: &N, +) -> NodePtr> +where + N: InnerNodeCommon, +{ + 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}; diff --git a/src/raw/representation/inner_node_indirect.rs b/src/raw/representation/inner_node_indirect.rs index 01fee42b..43d8917a 100644 --- a/src/raw/representation/inner_node_indirect.rs +++ b/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 InnerNodeCommon OpaqueNodePtr { + 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::>() { + 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. diff --git a/src/raw/representation/inner_node_sorted.rs b/src/raw/representation/inner_node_sorted.rs index 8f12522d..9c72b929 100644 --- a/src/raw/representation/inner_node_sorted.rs +++ b/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 InnerNodeCommon OpaqueNodePtr { + 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::>() { + return first; + } + + // Must have at least two children right + let second = *children.get_unchecked(1); + second + } + } } /// Node that references between 2 and 4 children