Skip to content
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.

Sign up for GitHub

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

BTreeMap: better distinguish the root holder from the root node #75329

Merged
merged 1 commit into from
Aug 11, 2020
Merged

BTreeMap: better distinguish the root holder from the root node #75329

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
73 changes: 40 additions & 33 deletions library/alloc/src/collections/btree/map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,14 @@

use core::borrow::Borrow;
use core::cmp::Ordering;
use core::fmt::Debug;
use core::fmt::{self, Debug};
use core::hash::{Hash, Hasher};
use core::iter::{FromIterator, FusedIterator, Peekable};
use core::marker::PhantomData;
use core::mem::{self, ManuallyDrop};
use core::ops::Bound::{Excluded, Included, Unbounded};
use core::ops::{Index, RangeBounds};
use core::{fmt, ptr};
use core::ptr;

use super::node::{self, marker, ForceResult::*, Handle, InsertResult::*, NodeRef};
use super::search::{self, SearchResult::*};
Expand Down Expand Up @@ -154,7 +154,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {

{
let root = out_tree.root.as_mut().unwrap(); // unwrap succeeds because we just wrapped
let mut out_node = match root.as_mut().force() {
let mut out_node = match root.node_as_mut().force() {
Leaf(leaf) => leaf,
Internal(_) => unreachable!(),
};
Expand Down Expand Up @@ -210,7 +210,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
// Ord` constraint, which this method lacks.
BTreeMap { root: None, length: 0 }
} else {
clone_subtree(self.root.as_ref().unwrap().as_ref()) // unwrap succeeds because not empty
clone_subtree(self.root.as_ref().unwrap().node_as_ref()) // unwrap succeeds because not empty
}
}
}
Expand All @@ -223,14 +223,16 @@ where
type Key = K;

fn get(&self, key: &Q) -> Option<&K> {
match search::search_tree(self.root.as_ref()?.as_ref(), key) {
let root_node = self.root.as_ref()?.node_as_ref();
match search::search_tree(root_node, key) {
Found(handle) => Some(handle.into_kv().0),
GoDown(_) => None,
}
}

fn take(&mut self, key: &Q) -> Option<K> {
match search::search_tree(self.root.as_mut()?.as_mut(), key) {
let root_node = self.root.as_mut()?.node_as_mut();
match search::search_tree(root_node, key) {
Found(handle) => Some(
OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData }
.remove_kv()
Expand All @@ -242,7 +244,7 @@ where

fn replace(&mut self, key: K) -> Option<K> {
let root = Self::ensure_is_owned(&mut self.root);
match search::search_tree::<marker::Mut<'_>, K, (), K>(root.as_mut(), &key) {
match search::search_tree::<marker::Mut<'_>, K, (), K>(root.node_as_mut(), &key) {
Found(handle) => Some(mem::replace(handle.into_kv_mut().0, key)),
GoDown(handle) => {
VacantEntry { key, handle, length: &mut self.length, _marker: PhantomData }
Expand Down Expand Up @@ -565,7 +567,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
K: Borrow<Q>,
Q: Ord,
{
match search::search_tree(self.root.as_ref()?.as_ref(), key) {
let root_node = self.root.as_ref()?.node_as_ref();
match search::search_tree(root_node, key) {
Found(handle) => Some(handle.into_kv().1),
GoDown(_) => None,
}
Expand All @@ -592,7 +595,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
K: Borrow<Q>,
Q: Ord,
{
match search::search_tree(self.root.as_ref()?.as_ref(), k) {
let root_node = self.root.as_ref()?.node_as_ref();
match search::search_tree(root_node, k) {
Found(handle) => Some(handle.into_kv()),
GoDown(_) => None,
}
Expand All @@ -617,8 +621,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
/// ```
#[unstable(feature = "map_first_last", issue = "62924")]
pub fn first_key_value(&self) -> Option<(&K, &V)> {
let front = self.root.as_ref()?.as_ref().first_leaf_edge();
front.right_kv().ok().map(Handle::into_kv)
let root_node = self.root.as_ref()?.node_as_ref();
root_node.first_leaf_edge().right_kv().ok().map(Handle::into_kv)
}

/// Returns the first entry in the map for in-place manipulation.
Expand All @@ -643,8 +647,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
/// ```
#[unstable(feature = "map_first_last", issue = "62924")]
pub fn first_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>> {
let front = self.root.as_mut()?.as_mut().first_leaf_edge();
let kv = front.right_kv().ok()?;
let root_node = self.root.as_mut()?.node_as_mut();
let kv = root_node.first_leaf_edge().right_kv().ok()?;
Some(OccupiedEntry {
handle: kv.forget_node_type(),
length: &mut self.length,
Expand Down Expand Up @@ -694,8 +698,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
/// ```
#[unstable(feature = "map_first_last", issue = "62924")]
pub fn last_key_value(&self) -> Option<(&K, &V)> {
let back = self.root.as_ref()?.as_ref().last_leaf_edge();
back.left_kv().ok().map(Handle::into_kv)
let root_node = self.root.as_ref()?.node_as_ref();
root_node.last_leaf_edge().left_kv().ok().map(Handle::into_kv)
}

/// Returns the last entry in the map for in-place manipulation.
Expand All @@ -720,8 +724,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
/// ```
#[unstable(feature = "map_first_last", issue = "62924")]
pub fn last_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>> {
let back = self.root.as_mut()?.as_mut().last_leaf_edge();
let kv = back.left_kv().ok()?;
let root_node = self.root.as_mut()?.node_as_mut();
let kv = root_node.last_leaf_edge().left_kv().ok()?;
Some(OccupiedEntry {
handle: kv.forget_node_type(),
length: &mut self.length,
Expand Down Expand Up @@ -805,7 +809,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
K: Borrow<Q>,
Q: Ord,
{
match search::search_tree(self.root.as_mut()?.as_mut(), key) {
let root_node = self.root.as_mut()?.node_as_mut();
match search::search_tree(root_node, key) {
Found(handle) => Some(handle.into_kv_mut().1),
GoDown(_) => None,
}
Expand Down Expand Up @@ -899,7 +904,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
K: Borrow<Q>,
Q: Ord,
{
match search::search_tree(self.root.as_mut()?.as_mut(), key) {
let root_node = self.root.as_mut()?.node_as_mut();
match search::search_tree(root_node, key) {
Found(handle) => Some(
OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData }
.remove_entry(),
Expand Down Expand Up @@ -995,7 +1001,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
R: RangeBounds<T>,
{
if let Some(root) = &self.root {
let (f, b) = range_search(root.as_ref(), range);
let (f, b) = range_search(root.node_as_ref(), range);

Range { front: Some(f), back: Some(b) }
} else {
Expand Down Expand Up @@ -1041,7 +1047,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
R: RangeBounds<T>,
{
if let Some(root) = &mut self.root {
let (f, b) = range_search(root.as_mut(), range);
let (f, b) = range_search(root.node_as_mut(), range);

RangeMut { front: Some(f), back: Some(b), _marker: PhantomData }
} else {
Expand Down Expand Up @@ -1071,7 +1077,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
// FIXME(@porglezomp) Avoid allocating if we don't insert
let root = Self::ensure_is_owned(&mut self.root);
match search::search_tree(root.as_mut(), &key) {
match search::search_tree(root.node_as_mut(), &key) {
Found(handle) => {
Occupied(OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData })
}
Expand All @@ -1083,7 +1089,7 @@ impl<K: Ord, V> BTreeMap<K, V> {

fn from_sorted_iter<I: Iterator<Item = (K, V)>>(&mut self, iter: I) {
let root = Self::ensure_is_owned(&mut self.root);
let mut cur_node = root.as_mut().last_leaf_edge().into_node();
let mut cur_node = root.node_as_mut().last_leaf_edge().into_node();
// Iterate through all key-value pairs, pushing them into nodes at the right level.
for (key, value) in iter {
// Try to push key-value pair into the current leaf node.
Expand Down Expand Up @@ -1133,7 +1139,7 @@ impl<K: Ord, V> BTreeMap<K, V> {

fn fix_right_edge(root: &mut node::Root<K, V>) {
// Handle underfull nodes, start from the top.
let mut cur_node = root.as_mut();
let mut cur_node = root.node_as_mut();
while let Internal(internal) = cur_node.force() {
// Check if right-most child is underfull.
let mut last_edge = internal.last_edge();
Expand Down Expand Up @@ -1201,8 +1207,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
}

{
let mut left_node = left_root.as_mut();
let mut right_node = right_root.as_mut();
let mut left_node = left_root.node_as_mut();
let mut right_node = right_root.node_as_mut();

loop {
let mut split_edge = match search::search_node(left_node, key) {
Expand Down Expand Up @@ -1280,7 +1286,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
DrainFilter { pred, inner: self.drain_filter_inner() }
}
pub(super) fn drain_filter_inner(&mut self) -> DrainFilterInner<'_, K, V> {
let front = self.root.as_mut().map(|r| r.as_mut().first_leaf_edge());
let root_node = self.root.as_mut().map(|r| r.node_as_mut());
let front = root_node.map(|rn| rn.first_leaf_edge());
DrainFilterInner {
length: &mut self.length,
cur_leaf_edge: front,
Expand Down Expand Up @@ -1315,7 +1322,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
res
}

self.length = dfs(self.root.as_ref().unwrap().as_ref());
self.length = dfs(self.root.as_ref().unwrap().node_as_ref());
}

/// Creates a consuming iterator visiting all the keys, in sorted order.
Expand Down Expand Up @@ -2251,7 +2258,7 @@ impl<K, V> BTreeMap<K, V> {
#[stable(feature = "rust1", since = "1.0.0")]
pub fn iter(&self) -> Iter<'_, K, V> {
if let Some(root) = &self.root {
let (f, b) = full_range_search(root.as_ref());
let (f, b) = full_range_search(root.node_as_ref());

Iter { range: Range { front: Some(f), back: Some(b) }, length: self.length }
} else {
Expand Down Expand Up @@ -2283,7 +2290,7 @@ impl<K, V> BTreeMap<K, V> {
#[stable(feature = "rust1", since = "1.0.0")]
pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
if let Some(root) = &mut self.root {
let (f, b) = full_range_search(root.as_mut());
let (f, b) = full_range_search(root.node_as_mut());

IterMut {
range: RangeMut { front: Some(f), back: Some(b), _marker: PhantomData },
Expand Down Expand Up @@ -2895,7 +2902,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInter
impl<K, V> node::Root<K, V> {
/// Removes empty levels on the top, but keep an empty leaf if the entire tree is empty.
fn fix_top(&mut self) {
while self.height() > 0 && self.as_ref().len() == 0 {
while self.height() > 0 && self.node_as_ref().len() == 0 {
self.pop_internal_level();
}
}
Expand All @@ -2904,7 +2911,7 @@ impl<K, V> node::Root<K, V> {
self.fix_top();

{
let mut cur_node = self.as_mut();
let mut cur_node = self.node_as_mut();

while let Internal(node) = cur_node.force() {
let mut last_kv = node.last_kv();
Expand All @@ -2930,7 +2937,7 @@ impl<K, V> node::Root<K, V> {
self.fix_top();

{
let mut cur_node = self.as_mut();
let mut cur_node = self.node_as_mut();

while let Internal(node) = cur_node.force() {
let mut first_kv = node.first_kv();
Expand Down
20 changes: 11 additions & 9 deletions library/alloc/src/collections/btree/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,8 @@ impl<K, V> Root<K, V> {
Root { node: BoxedNode::from_leaf(Box::new(unsafe { LeafNode::new() })), height: 0 }
}

pub fn as_ref(&self) -> NodeRef<marker::Immut<'_>, K, V, marker::LeafOrInternal> {
/// Borrows and returns an immutable reference to the node owned by the root.
pub fn node_as_ref(&self) -> NodeRef<marker::Immut<'_>, K, V, marker::LeafOrInternal> {
NodeRef {
height: self.height,
node: self.node.as_ptr(),
Expand All @@ -172,7 +173,8 @@ impl<K, V> Root<K, V> {
}
}

pub fn as_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, marker::LeafOrInternal> {
/// Borrows and returns a mutable reference to the node owned by the root.
pub fn node_as_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, marker::LeafOrInternal> {
NodeRef {
height: self.height,
node: self.node.as_ptr(),
Expand Down Expand Up @@ -226,12 +228,12 @@ impl<K, V> Root<K, V> {

self.node = unsafe {
BoxedNode::from_ptr(
self.as_mut().cast_unchecked::<marker::Internal>().first_edge().descend().node,
self.node_as_mut().cast_unchecked::<marker::Internal>().first_edge().descend().node,
)
};
self.height -= 1;
unsafe {
(*self.as_mut().as_leaf_mut()).parent = ptr::null();
(*self.node_as_mut().as_leaf_mut()).parent = ptr::null();
}

unsafe {
Expand Down Expand Up @@ -616,7 +618,7 @@ impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
let edge =
ptr::read(internal.as_internal().edges.get_unchecked(idx + 1).as_ptr());
let mut new_root = Root { node: edge, height: internal.height - 1 };
(*new_root.as_mut().as_leaf_mut()).parent = ptr::null();
(*new_root.node_as_mut().as_leaf_mut()).parent = ptr::null();
Some(new_root)
}
};
Expand Down Expand Up @@ -648,7 +650,7 @@ impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
);

let mut new_root = Root { node: edge, height: internal.height - 1 };
(*new_root.as_mut().as_leaf_mut()).parent = ptr::null();
(*new_root.node_as_mut().as_leaf_mut()).parent = ptr::null();

for i in 0..old_len {
Handle::new_edge(internal.reborrow_mut(), i).correct_parent_link();
Expand Down Expand Up @@ -868,7 +870,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
} else {
unsafe {
Handle::new_edge(
right.as_mut().cast_unchecked::<marker::Leaf>(),
right.node_as_mut().cast_unchecked::<marker::Leaf>(),
self.idx - (B + 1),
)
.insert_fit(key, val)
Expand Down Expand Up @@ -943,7 +945,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
} else {
unsafe {
Handle::new_edge(
right.as_mut().cast_unchecked::<marker::Internal>(),
right.node_as_mut().cast_unchecked::<marker::Internal>(),
self.idx - (B + 1),
)
.insert_fit(key, val, edge);
Expand Down Expand Up @@ -1117,7 +1119,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
let mut new_root = Root { node: BoxedNode::from_internal(new_node), height };

for i in 0..(new_len + 1) {
Handle::new_edge(new_root.as_mut().cast_unchecked(), i).correct_parent_link();
Handle::new_edge(new_root.node_as_mut().cast_unchecked(), i).correct_parent_link();
}

(self.node, k, v, new_root)
Expand Down