Skip to content

BTreeMap: convert search functions to methods #81159

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 1 commit into from
Jan 18, 2021
Merged
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
18 changes: 9 additions & 9 deletions library/alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
@@ -10,7 +10,7 @@ use core::ptr;

use super::borrow::DormantMutRef;
use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};
use super::search::{self, SearchResult::*};
use super::search::SearchResult::*;
use super::unwrap_unchecked;

mod entry;
@@ -230,7 +230,7 @@ where

fn get(&self, key: &Q) -> Option<&K> {
let root_node = self.root.as_ref()?.reborrow();
match search::search_tree(root_node, key) {
match root_node.search_tree(key) {
Found(handle) => Some(handle.into_kv().0),
GoDown(_) => None,
}
@@ -239,7 +239,7 @@ where
fn take(&mut self, key: &Q) -> Option<K> {
let (map, dormant_map) = DormantMutRef::new(self);
let root_node = map.root.as_mut()?.borrow_mut();
match search::search_tree(root_node, key) {
match root_node.search_tree(key) {
Found(handle) => {
Some(OccupiedEntry { handle, dormant_map, _marker: PhantomData }.remove_kv().0)
}
@@ -250,7 +250,7 @@ where
fn replace(&mut self, key: K) -> Option<K> {
let (map, dormant_map) = DormantMutRef::new(self);
let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut();
match search::search_tree::<marker::Mut<'_>, K, (), K>(root_node, &key) {
match root_node.search_tree::<K>(&key) {
Found(mut kv) => Some(mem::replace(kv.key_mut(), key)),
GoDown(handle) => {
VacantEntry { key, handle, dormant_map, _marker: PhantomData }.insert(());
@@ -526,7 +526,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
Q: Ord,
{
let root_node = self.root.as_ref()?.reborrow();
match search::search_tree(root_node, key) {
match root_node.search_tree(key) {
Found(handle) => Some(handle.into_kv().1),
GoDown(_) => None,
}
@@ -554,7 +554,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
Q: Ord,
{
let root_node = self.root.as_ref()?.reborrow();
match search::search_tree(root_node, k) {
match root_node.search_tree(k) {
Found(handle) => Some(handle.into_kv()),
GoDown(_) => None,
}
@@ -762,7 +762,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
Q: Ord,
{
let root_node = self.root.as_mut()?.borrow_mut();
match search::search_tree(root_node, key) {
match root_node.search_tree(key) {
Found(handle) => Some(handle.into_val_mut()),
GoDown(_) => None,
}
@@ -858,7 +858,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
{
let (map, dormant_map) = DormantMutRef::new(self);
let root_node = map.root.as_mut()?.borrow_mut();
match search::search_tree(root_node, key) {
match root_node.search_tree(key) {
Found(handle) => {
Some(OccupiedEntry { handle, dormant_map, _marker: PhantomData }.remove_entry())
}
@@ -1051,7 +1051,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
// FIXME(@porglezomp) Avoid allocating if we don't insert
let (map, dormant_map) = DormantMutRef::new(self);
let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut();
match search::search_tree(root_node, &key) {
match root_node.search_tree(&key) {
Found(handle) => Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData }),
GoDown(handle) => {
Vacant(VacantEntry { key, handle, dormant_map, _marker: PhantomData })
10 changes: 5 additions & 5 deletions library/alloc/src/collections/btree/navigate.rs
Original file line number Diff line number Diff line change
@@ -5,7 +5,7 @@ use core::ops::RangeBounds;
use core::ptr;

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

/// Finds the leaf edges delimiting a specified range in or underneath a node.
@@ -42,14 +42,14 @@ where

loop {
let front = match (min_found, range.start_bound()) {
(false, Included(key)) => match search::search_node(min_node, key) {
(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 search::search_node(min_node, key) {
(false, Excluded(key)) => match min_node.search_node(key) {
SearchResult::Found(kv) => {
min_found = true;
kv.right_edge()
@@ -62,14 +62,14 @@ where
};

let back = match (max_found, range.end_bound()) {
(false, Included(key)) => match search::search_node(max_node, key) {
(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 search::search_node(max_node, key) {
(false, Excluded(key)) => match max_node.search_node(key) {
SearchResult::Found(kv) => {
max_found = true;
kv.left_edge()
134 changes: 67 additions & 67 deletions library/alloc/src/collections/btree/search.rs
Original file line number Diff line number Diff line change
@@ -10,79 +10,79 @@ pub enum SearchResult<BorrowType, K, V, FoundType, GoDownType> {
GoDown(Handle<NodeRef<BorrowType, K, V, GoDownType>, marker::Edge>),
}

/// Looks up a given key in a (sub)tree headed by the given node, recursively.
/// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
/// returns a `GoDown` with the handle of the leaf edge where the key belongs.
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
pub fn search_tree<BorrowType, K, V, Q: ?Sized>(
mut node: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
key: &Q,
) -> SearchResult<BorrowType, K, V, marker::LeafOrInternal, marker::Leaf>
where
Q: Ord,
K: Borrow<Q>,
{
loop {
match search_node(node, key) {
Found(handle) => return Found(handle),
GoDown(handle) => match handle.force() {
Leaf(leaf) => return GoDown(leaf),
Internal(internal) => {
node = internal.descend();
continue;
}
},
pub enum IndexResult {
KV(usize),
Edge(usize),
}

impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
/// Looks up a given key in a (sub)tree headed by the node, recursively.
/// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
/// returns a `GoDown` with the handle of the leaf edge where the key belongs.
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
pub fn search_tree<Q: ?Sized>(
mut self,
key: &Q,
) -> SearchResult<BorrowType, K, V, marker::LeafOrInternal, marker::Leaf>
where
Q: Ord,
K: Borrow<Q>,
{
loop {
self = match self.search_node(key) {
Found(handle) => return Found(handle),
GoDown(handle) => match handle.force() {
Leaf(leaf) => return GoDown(leaf),
Internal(internal) => internal.descend(),
},
}
}
}
}

/// Looks up a given key in a given node, without recursion.
/// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
/// returns a `GoDown` with the handle of the edge where the key might be found
/// (if the node is internal) or where the key can be inserted.
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
pub fn search_node<BorrowType, K, V, Type, Q: ?Sized>(
node: NodeRef<BorrowType, K, V, Type>,
key: &Q,
) -> SearchResult<BorrowType, K, V, Type, Type>
where
Q: Ord,
K: Borrow<Q>,
{
match search_linear(&node, key) {
(idx, true) => Found(unsafe { Handle::new_kv(node, idx) }),
(idx, false) => GoDown(unsafe { Handle::new_edge(node, idx) }),
impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
/// Looks up a given key in the node, without recursion.
/// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
/// returns a `GoDown` with the handle of the edge where the key might be found
/// (if the node is internal) or where the key can be inserted.
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
pub fn search_node<Q: ?Sized>(self, key: &Q) -> SearchResult<BorrowType, K, V, Type, Type>
where
Q: Ord,
K: Borrow<Q>,
{
match self.find_index(key) {
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 and `true`, or the edge index where the key belongs and `false`.
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
node: &NodeRef<BorrowType, K, V, Type>,
key: &Q,
) -> (usize, bool)
where
Q: Ord,
K: Borrow<Q>,
{
// This function is defined over all borrow types (immutable, mutable, owned).
// Using `keys_at()` is fine here even if BorrowType is mutable, as all we return
// is an index -- not a reference.
let len = node.len();
for i in 0..len {
let k = unsafe { node.reborrow().key_at(i) };
match key.cmp(k.borrow()) {
Ordering::Greater => {}
Ordering::Equal => return (i, true),
Ordering::Less => return (i, false),
/// Returns either the KV index in the node at which the key (or an equivalent)
/// exists, or the edge index where the key belongs.
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
fn find_index<Q: ?Sized>(&self, key: &Q) -> IndexResult
where
Q: Ord,
K: Borrow<Q>,
{
// This function is defined over all borrow types (immutable, mutable, owned).
// Using `keys_at()` is fine here even if BorrowType is mutable, as all we return
// is an index -- not a reference.
let len = self.len();
for i in 0..len {
let k = unsafe { self.reborrow().key_at(i) };
match key.cmp(k.borrow()) {
Ordering::Greater => {}
Ordering::Equal => return IndexResult::KV(i),
Ordering::Less => return IndexResult::Edge(i),
}
}
IndexResult::Edge(len)
}
(len, false)
}
4 changes: 2 additions & 2 deletions library/alloc/src/collections/btree/split.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use super::map::MIN_LEN;
use super::node::{ForceResult::*, Root};
use super::search::{search_node, SearchResult::*};
use super::search::SearchResult::*;
use core::borrow::Borrow;

impl<K, V> Root<K, V> {
@@ -21,7 +21,7 @@ impl<K, V> Root<K, V> {
let mut right_node = right_root.borrow_mut();

loop {
let mut split_edge = match search_node(left_node, key) {
let mut split_edge = match left_node.search_node(key) {
// key is going to the right tree
Found(kv) => kv.left_edge(),
GoDown(edge) => edge,