diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs index 17b4689e4a0ee..b57a9d53392b3 100644 --- a/library/alloc/src/collections/btree/navigate.rs +++ b/library/alloc/src/collections/btree/navigate.rs @@ -104,34 +104,39 @@ where } } -/// Equivalent to `range_search(k, v, ..)` but without the `Ord` bound. -fn full_range( - root1: NodeRef, - root2: NodeRef, -) -> ( - Handle, marker::Edge>, - Handle, marker::Edge>, -) { - let mut min_node = root1; - let mut max_node = root2; - loop { - let front = min_node.first_edge(); - let back = max_node.last_edge(); - match (front.force(), back.force()) { - (Leaf(f), Leaf(b)) => { - return (f, b); +/// Equivalent to `range_search(self, self, ..)` but without the `Ord` bound. +/// Equivalent to `(self.first_leaf_edge(), self.last_leaf_edge())` but +/// duplicating `self` and more efficient. +/// # Safety +/// - Do not expose for `BorrowType` `Mut`, because one handle could destroy the +/// edge or node referred to by the other. +/// - Unless `BorrowType` is `Immut`, do not use the two handles to visit the +/// same KV twice. +impl NodeRef { + unsafe fn find_leaf_edges_spanning_tree( + self, + ) -> ( + Handle, marker::Edge>, + Handle, marker::Edge>, + ) { + match self.force() { + Leaf(root1) => { + let root2 = unsafe { ptr::read(&root1) }; + return (root1.first_edge(), root2.last_edge()); } - (Internal(min_int), Internal(max_int)) => { - min_node = min_int.descend(); - max_node = max_int.descend(); + Internal(root1) => { + let root2 = unsafe { ptr::read(&root1) }; + // Read both ends of the edges array jointly. + let first_child = root1.first_edge().descend(); + let last_child = root2.last_edge().descend(); + (first_child.first_leaf_edge(), last_child.last_leaf_edge()) } - _ => unreachable!("BTreeMap has different depths"), - }; + } } } impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> { - /// Creates a pair of leaf edges delimiting a specified range in or underneath a node. + /// Finds the pair of leaf edges delimiting a specified range in the tree. /// /// The result is meaningful only if the tree is ordered by key, like the tree /// in a `BTreeMap` is. @@ -150,14 +155,15 @@ impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> range_search(self, self, range) } - /// Returns (self.first_leaf_edge(), self.last_leaf_edge()), but more efficiently. + /// Finds the pair of leaf edges delimiting the entire tree. pub fn full_range( self, ) -> ( Handle, K, V, marker::Leaf>, marker::Edge>, Handle, K, V, marker::Leaf>, marker::Edge>, ) { - full_range(self, self) + // SAFETY: our borrow type is `Immut`. + unsafe { self.find_leaf_edges_spanning_tree() } } } @@ -195,10 +201,9 @@ impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> Handle, K, V, marker::Leaf>, marker::Edge>, Handle, K, V, marker::Leaf>, marker::Edge>, ) { - // We duplicate the root NodeRef here -- we will never visit the same KV - // twice, and never end up with overlapping value references. - let self2 = unsafe { ptr::read(&self) }; - full_range(self, self2) + // SAFETY: we will never visit the same KV twice, + // and never end up with overlapping value references. + unsafe { self.find_leaf_edges_spanning_tree() } } } @@ -212,10 +217,9 @@ impl NodeRef { Handle, marker::Edge>, Handle, marker::Edge>, ) { - // We duplicate the root NodeRef here -- we will never access it in a way - // that overlaps references obtained from the root. - let self2 = unsafe { ptr::read(&self) }; - full_range(self, self2) + // SAFETY: we will never visit the same KV twice, + // because we only visit any KV to drop it. + unsafe { self.find_leaf_edges_spanning_tree() } } }