Skip to content

Commit

Permalink
Auto merge of #81094 - ssomers:btree_drainy_refactor_3, r=Mark-Simula…
Browse files Browse the repository at this point in the history 
…crum

BTreeMap: split up range_search into two stages

`range_search` expects the caller to pass the same root twice and starts searching a node for both bounds of a range. It's not very clear that in the early iterations, it searches twice in the same node. This PR splits that search up in an initial `find_leaf_edges_spanning_range` that postpones aliasing until the last second, and a second phase for continuing the search for the range in the each subtree independently (`find_lower_bound_edge` & `find_upper_bound_edge`), which greatly helps for use in #81075. It also moves those functions over to the search module.

r? `@Mark-Simulacrum`
  • Loading branch information
@bors
bors committed Mar 1, 2021
2 parents 43fc1b3 + deebb63 commit 3b150b7
Show file tree
Hide file tree
Showing 5 changed files with 228 additions and 118 deletions.
142 changes: 45 additions & 97 deletions library/alloc/src/collections/btree/navigate.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,8 @@
use core::borrow::Borrow;
use core::cmp::Ordering;
use core::ops::Bound::{Excluded, Included, Unbounded};
use core::ops::RangeBounds;
use core::ptr;

use super::node::{marker, ForceResult::*, Handle, NodeRef};
use super::search::SearchResult;

pub struct LeafRange<BorrowType, K, V> {
pub front: Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
Expand All @@ -30,100 +27,50 @@ impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {
}
}

/// Finds the leaf edges delimiting a specified range in or underneath a node.
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
fn range_search<BorrowType: marker::BorrowType, K, V, Q, R>(
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
range: R,
) -> LeafRange<BorrowType, K, V>
where
Q: ?Sized + Ord,
K: Borrow<Q>,
R: RangeBounds<Q>,
{
// WARNING: Inlining these variables would be unsound (#81138)
// We assume the bounds reported by `range` remain the same, but
// an adversarial implementation could change between calls
let start = range.start_bound();
let end = range.end_bound();
match (start, end) {
(Excluded(s), Excluded(e)) if s == e => {
panic!("range start and end are equal and excluded in BTreeMap")
}
(Included(s) | Excluded(s), Included(e) | Excluded(e)) if s > e => {
panic!("range start is greater than range end in BTreeMap")
}
_ => {}
};

let mut min_node = root1;
let mut max_node = root2;
let mut min_found = false;
let mut max_found = false;

loop {
// Using `range` again would be unsound (#81138)
let front = match (min_found, start) {
(false, Included(key)) => match min_node.search_node(key) {
SearchResult::Found(kv) => {
min_found = true;
kv.left_edge()
}
SearchResult::GoDown(edge) => edge,
},
(false, Excluded(key)) => match min_node.search_node(key) {
SearchResult::Found(kv) => {
min_found = true;
kv.right_edge()
}
SearchResult::GoDown(edge) => edge,
},
(true, Included(_)) => min_node.last_edge(),
(true, Excluded(_)) => min_node.first_edge(),
(_, Unbounded) => min_node.first_edge(),
};

// Using `range` again would be unsound (#81138)
let back = match (max_found, end) {
(false, Included(key)) => match max_node.search_node(key) {
SearchResult::Found(kv) => {
max_found = true;
kv.right_edge()
}
SearchResult::GoDown(edge) => edge,
},
(false, Excluded(key)) => match max_node.search_node(key) {
SearchResult::Found(kv) => {
max_found = true;
kv.left_edge()
impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
/// Finds the distinct leaf edges delimiting a specified range in a tree.
/// Returns either a pair of different handles into the same tree or a pair
/// of empty options.
/// # Safety
/// Unless `BorrowType` is `Immut`, do not use the duplicate handles to
/// visit the same KV twice.
unsafe fn find_leaf_edges_spanning_range<Q: ?Sized, R>(
self,
range: R,
) -> LeafRange<BorrowType, K, V>
where
Q: Ord,
K: Borrow<Q>,
R: RangeBounds<Q>,
{
match self.search_tree_for_bifurcation(&range) {
Err(_) => LeafRange::none(),
Ok((
node,
lower_edge_idx,
upper_edge_idx,
mut lower_child_bound,
mut upper_child_bound,
)) => {
let mut lower_edge = unsafe { Handle::new_edge(ptr::read(&node), lower_edge_idx) };
let mut upper_edge = unsafe { Handle::new_edge(node, upper_edge_idx) };
loop {
match (lower_edge.force(), upper_edge.force()) {
(Leaf(f), Leaf(b)) => return LeafRange { front: Some(f), back: Some(b) },
(Internal(f), Internal(b)) => {
(lower_edge, lower_child_bound) =
f.descend().find_lower_bound_edge(lower_child_bound);
(upper_edge, upper_child_bound) =
b.descend().find_upper_bound_edge(upper_child_bound);
}
_ => unreachable!("BTreeMap has different depths"),
}
}
SearchResult::GoDown(edge) => edge,
},
(true, Included(_)) => max_node.first_edge(),
(true, Excluded(_)) => max_node.last_edge(),
(_, Unbounded) => max_node.last_edge(),
};

if front.partial_cmp(&back) == Some(Ordering::Greater) {
panic!("Ord is ill-defined in BTreeMap range");
}
match (front.force(), back.force()) {
(Leaf(f), Leaf(b)) => {
return LeafRange { front: Some(f), back: Some(b) };
}
(Internal(min_int), Internal(max_int)) => {
min_node = min_int.descend();
max_node = max_int.descend();
}
_ => unreachable!("BTreeMap has different depths"),
};
}
}
}

/// Equivalent to `range_search(root1, root2, ..)` but without the `Ord` bound.
/// Equivalent to `(root1.first_leaf_edge(), root2.last_leaf_edge())` but more efficient.
fn full_range<BorrowType: marker::BorrowType, K, V>(
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
Expand Down Expand Up @@ -158,7 +105,8 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>
K: Borrow<Q>,
R: RangeBounds<Q>,
{
range_search(self, self, range)
// SAFETY: our borrow type is immutable.
unsafe { self.find_leaf_edges_spanning_range(range) }
}

/// Finds the pair of leaf edges delimiting an entire tree.
Expand All @@ -174,16 +122,16 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::ValMut<'a>, K, V, marker::LeafOrInternal>
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
///
/// # Safety
/// Do not use the duplicate handles to visit the same KV twice.
pub fn range_search<Q, R>(self, range: R) -> LeafRange<marker::ValMut<'a>, K, V>
where
Q: ?Sized + Ord,
K: Borrow<Q>,
R: RangeBounds<Q>,
{
// 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) };
range_search(self, self2, range)
unsafe { self.find_leaf_edges_spanning_range(range) }
}

/// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
Expand Down
10 changes: 0 additions & 10 deletions library/alloc/src/collections/btree/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@
// since leaf edges are empty and need no data representation. In an internal node,
// an edge both identifies a position and contains a pointer to a child node.

use core::cmp::Ordering;
use core::marker::PhantomData;
use core::mem::{self, MaybeUninit};
use core::ptr::{self, NonNull};
Expand Down Expand Up @@ -742,15 +741,6 @@ impl<BorrowType, K, V, NodeType, HandleType> PartialEq
}
}

impl<BorrowType, K, V, NodeType, HandleType> PartialOrd
for Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType>
{
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
let Self { node, idx, _marker } = self;
if node.eq(&other.node) { Some(idx.cmp(&other.idx)) } else { None }
}
}

impl<BorrowType, K, V, NodeType, HandleType>
Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType>
{
Expand Down
10 changes: 1 addition & 9 deletions library/alloc/src/collections/btree/node/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ use super::super::navigate;
use super::*;
use crate::fmt::Debug;
use crate::string::String;
use core::cmp::Ordering::*;

impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal> {
// Asserts that the back pointer in each reachable node points to its parent.
Expand Down Expand Up @@ -67,7 +66,7 @@ fn test_splitpoint() {
}

#[test]
fn test_partial_cmp_eq() {
fn test_partial_eq() {
let mut root1 = NodeRef::new_leaf();
root1.borrow_mut().push(1, ());
let mut root1 = NodeRef::new_internal(root1.forget_type()).forget_type();
Expand All @@ -87,13 +86,6 @@ fn test_partial_cmp_eq() {
assert!(top_edge_1 == top_edge_1);
assert!(top_edge_1 != top_edge_2);

assert_eq!(leaf_edge_1a.partial_cmp(&leaf_edge_1a), Some(Equal));
assert_eq!(leaf_edge_1a.partial_cmp(&leaf_edge_1b), Some(Less));
assert_eq!(leaf_edge_1a.partial_cmp(&top_edge_1), None);
assert_eq!(leaf_edge_1a.partial_cmp(&top_edge_2), None);
assert_eq!(top_edge_1.partial_cmp(&top_edge_1), Some(Equal));
assert_eq!(top_edge_1.partial_cmp(&top_edge_2), None);

root1.pop_internal_level();
unsafe { root1.into_dying().deallocate_and_ascend() };
unsafe { root2.into_dying().deallocate_and_ascend() };
Expand Down
Loading

0 comments on commit 3b150b7

Please sign in to comment.