Skip to content

BTree: no longer search arrays twice to check Ord #83267

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 4, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions library/alloc/src/collections/btree/map/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -776,7 +776,6 @@ fn test_range_backwards_4() {
}

#[test]
#[should_panic]
fn test_range_finding_ill_order_in_map() {
let mut map = BTreeMap::new();
map.insert(Cyclic3::B, ());
Expand All @@ -789,7 +788,6 @@ fn test_range_finding_ill_order_in_map() {
}

#[test]
#[should_panic]
fn test_range_finding_ill_order_in_range_ord() {
// Has proper order the first time asked, then flips around.
struct EvilTwin(i32);
Expand Down
15 changes: 11 additions & 4 deletions library/alloc/src/collections/btree/navigate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,18 @@ impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {

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.
///
/// If such distinct edges exist, returns them in ascending order, meaning
/// that a non-zero number of calls to `next_unchecked` on the `front` of
/// the result and/or calls to `next_back_unchecked` on the `back` of the
/// result will eventually reach the same edge.
///
/// If there are no such edges, i.e., if the tree contains no key within
/// the range, returns a pair of empty options.
///
/// # Safety
/// Unless `BorrowType` is `Immut`, do not use the duplicate handles to
/// visit the same KV twice.
/// Unless `BorrowType` is `Immut`, do not use the handles to visit the same
/// KV twice.
unsafe fn find_leaf_edges_spanning_range<Q: ?Sized, R>(
self,
range: R,
Expand Down
53 changes: 32 additions & 21 deletions library/alloc/src/collections/btree/search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -68,13 +68,16 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
/// of the range is different from the edge matching the upper bound, i.e.,
/// the nearest node that has at least one key contained in the range.
///
/// If found, returns an `Ok` with that node, the pair of edge indices in it
/// delimiting the range, and the corresponding pair of bounds for
/// continuing the search in the child nodes, in case the node is internal.
/// If found, returns an `Ok` with that node, the strictly ascending pair of
/// edge indices in the node delimiting the range, and the corresponding
/// pair of bounds for continuing the search in the child nodes, in case
/// the node is internal.
///
/// If not found, returns an `Err` with the leaf edge matching the entire
/// range.
///
/// As a diagnostic service, panics if the range specifies impossible bounds.
///
/// The result is meaningful only if the tree is ordered by key.
pub fn search_tree_for_bifurcation<'r, Q: ?Sized, R>(
mut self,
Expand Down Expand Up @@ -113,10 +116,8 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
let mut upper_bound = SearchBound::from_range(end);
loop {
let (lower_edge_idx, lower_child_bound) = self.find_lower_bound_index(lower_bound);
let (upper_edge_idx, upper_child_bound) = self.find_upper_bound_index(upper_bound);
if lower_edge_idx > upper_edge_idx {
panic!("Ord is ill-defined in BTreeMap range")
}
let (upper_edge_idx, upper_child_bound) =
unsafe { self.find_upper_bound_index(upper_bound, lower_edge_idx) };
if lower_edge_idx < upper_edge_idx {
return Ok((
self,
Expand All @@ -126,6 +127,7 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
upper_child_bound,
));
}
debug_assert_eq!(lower_edge_idx, upper_edge_idx);
let common_edge = unsafe { Handle::new_edge(self, lower_edge_idx) };
match common_edge.force() {
Leaf(common_edge) => return Err(common_edge),
Expand Down Expand Up @@ -165,7 +167,7 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
Q: ?Sized + Ord,
K: Borrow<Q>,
{
let (edge_idx, bound) = self.find_upper_bound_index(bound);
let (edge_idx, bound) = unsafe { self.find_upper_bound_index(bound, 0) };
let edge = unsafe { Handle::new_edge(self, edge_idx) };
(edge, bound)
}
Expand All @@ -184,29 +186,33 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
Q: Ord,
K: Borrow<Q>,
{
match self.find_key_index(key) {
match unsafe { self.find_key_index(key, 0) } {
IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }),
IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }),
}
}

/// Returns either the KV index in the node at which the key (or an equivalent)
/// exists, or the edge index where the key belongs.
/// exists, or the edge index where the key belongs, starting from a particular index.
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
fn find_key_index<Q: ?Sized>(&self, key: &Q) -> IndexResult
///
/// # Safety
/// `start_index` must be a valid edge index for the node.
unsafe fn find_key_index<Q: ?Sized>(&self, key: &Q, start_index: usize) -> IndexResult
where
Q: Ord,
K: Borrow<Q>,
{
let node = self.reborrow();
let keys = node.keys();
for (i, k) in keys.iter().enumerate() {
debug_assert!(start_index <= keys.len());
for (offset, k) in unsafe { keys.get_unchecked(start_index..) }.iter().enumerate() {
match key.cmp(k.borrow()) {
Ordering::Greater => {}
Ordering::Equal => return IndexResult::KV(i),
Ordering::Less => return IndexResult::Edge(i),
Ordering::Equal => return IndexResult::KV(start_index + offset),
Ordering::Less => return IndexResult::Edge(start_index + offset),
}
}
IndexResult::Edge(keys.len())
Expand All @@ -226,11 +232,11 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
K: Borrow<Q>,
{
match bound {
Included(key) => match self.find_key_index(key) {
Included(key) => match unsafe { self.find_key_index(key, 0) } {
IndexResult::KV(idx) => (idx, AllExcluded),
IndexResult::Edge(idx) => (idx, bound),
},
Excluded(key) => match self.find_key_index(key) {
Excluded(key) => match unsafe { self.find_key_index(key, 0) } {
IndexResult::KV(idx) => (idx + 1, AllIncluded),
IndexResult::Edge(idx) => (idx, bound),
},
Expand All @@ -239,26 +245,31 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
}
}

/// Clone of `find_lower_bound_index` for the upper bound.
fn find_upper_bound_index<'r, Q>(
/// Mirror image of `find_lower_bound_index` for the upper bound,
/// with an additional parameter to skip part of the key array.
///
/// # Safety
/// `start_index` must be a valid edge index for the node.
unsafe fn find_upper_bound_index<'r, Q>(
&self,
bound: SearchBound<&'r Q>,
start_index: usize,
) -> (usize, SearchBound<&'r Q>)
where
Q: ?Sized + Ord,
K: Borrow<Q>,
{
match bound {
Included(key) => match self.find_key_index(key) {
Included(key) => match unsafe { self.find_key_index(key, start_index) } {
IndexResult::KV(idx) => (idx + 1, AllExcluded),
IndexResult::Edge(idx) => (idx, bound),
},
Excluded(key) => match self.find_key_index(key) {
Excluded(key) => match unsafe { self.find_key_index(key, start_index) } {
IndexResult::KV(idx) => (idx, AllIncluded),
IndexResult::Edge(idx) => (idx, bound),
},
AllIncluded => (self.len(), AllIncluded),
AllExcluded => (0, AllExcluded),
AllExcluded => (start_index, AllExcluded),
}
}
}