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

Jump to bottom

BtreeMap: superficially refactor root access #74391

Merged
merged 2 commits into from
Jul 16, 2020
Merged

BtreeMap: superficially refactor root access #74391

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ca253ca
Clean up or comment every unwrap in BTreeMap's main code.
ssomers Jul 13, 2020
b82d332
Separate off BTreeMap support functions and loose their irrelevant bo…
ssomers Jul 14, 2020
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
175 changes: 83 additions & 92 deletions src/liballoc/collections/btree/map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
let mut out_tree = BTreeMap { root: Some(node::Root::new_leaf()), length: 0 };

{
let root = out_tree.root.as_mut().unwrap();
let root = out_tree.root.as_mut().unwrap(); // unwrap succeeds because we just wrapped
let mut out_node = match root.as_mut().force() {
Leaf(leaf) => leaf,
Internal(_) => unreachable!(),
Expand All @@ -171,14 +171,10 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
}
Internal(internal) => {
let mut out_tree = clone_subtree(internal.first_edge().descend());
out_tree.ensure_root_is_owned();

{
// Ideally we'd use the return of ensure_root_is_owned
// instead of re-unwrapping here but unfortunately that
// borrows all of out_tree and we need access to the
// length below.
let mut out_node = out_tree.root.as_mut().unwrap().push_level();
let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);
let mut out_node = out_root.push_level();
let mut in_edge = internal.first_edge();
while let Ok(kv) = in_edge.right_kv() {
let (k, v) = kv.into_kv();
Expand Down Expand Up @@ -212,7 +208,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())
clone_subtree(self.root.as_ref().unwrap().as_ref()) // unwrap succeeds because not empty
}
}
}
Expand Down Expand Up @@ -243,8 +239,8 @@ where
}

fn replace(&mut self, key: K) -> Option<K> {
self.ensure_root_is_owned();
match search::search_tree::<marker::Mut<'_>, K, (), K>(self.root.as_mut()?.as_mut(), &key) {
let root = Self::ensure_is_owned(&mut self.root);
match search::search_tree::<marker::Mut<'_>, K, (), K>(root.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 @@ -943,7 +939,6 @@ impl<K: Ord, V> BTreeMap<K, V> {

// Second, we build a tree from the sorted sequence in linear time.
self.from_sorted_iter(iter);
self.fix_right_edge();
}

/// Constructs a double-ended iterator over a sub-range of elements in the map.
Expand Down Expand Up @@ -1058,8 +1053,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
#[stable(feature = "rust1", since = "1.0.0")]
pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
// FIXME(@porglezomp) Avoid allocating if we don't insert
self.ensure_root_is_owned();
match search::search_tree(self.root.as_mut().unwrap().as_mut(), &key) {
let root = Self::ensure_is_owned(&mut self.root);
match search::search_tree(root.as_mut(), &key) {
Found(handle) => {
Occupied(OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData })
}
Expand All @@ -1070,8 +1065,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
}

fn from_sorted_iter<I: Iterator<Item = (K, V)>>(&mut self, iter: I) {
self.ensure_root_is_owned();
let mut cur_node = self.root.as_mut().unwrap().as_mut().last_leaf_edge().into_node();
let root = Self::ensure_is_owned(&mut self.root);
let mut cur_node = root.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 @@ -1116,11 +1111,12 @@ impl<K: Ord, V> BTreeMap<K, V> {

self.length += 1;
}
Self::fix_right_edge(root)
}

fn fix_right_edge(&mut self) {
fn fix_right_edge(root: &mut node::Root<K, V>) {
// Handle underfull nodes, start from the top.
let mut cur_node = self.root.as_mut().unwrap().as_mut();
let mut cur_node = root.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 @@ -1179,16 +1175,17 @@ impl<K: Ord, V> BTreeMap<K, V> {
}

let total_num = self.len();
let left_root = self.root.as_mut().unwrap(); // unwrap succeeds because not empty

let mut right = Self::new();
let right_root = right.ensure_root_is_owned();
for _ in 0..(self.root.as_ref().unwrap().as_ref().height()) {
let right_root = Self::ensure_is_owned(&mut right.root);
for _ in 0..left_root.height() {
right_root.push_level();
}

{
let mut left_node = self.root.as_mut().unwrap().as_mut();
let mut right_node = right.root.as_mut().unwrap().as_mut();
let mut left_node = left_root.as_mut();
let mut right_node = right_root.as_mut();

loop {
let mut split_edge = match search::search_node(left_node, key) {
Expand All @@ -1214,12 +1211,10 @@ impl<K: Ord, V> BTreeMap<K, V> {
}
}

self.fix_right_border();
right.fix_left_border();
left_root.fix_right_border();
right_root.fix_left_border();

if self.root.as_ref().unwrap().as_ref().height()
< right.root.as_ref().unwrap().as_ref().height()
{
if left_root.height() < right_root.height() {
self.recalc_length();
right.length = total_num - self.len();
} else {
Expand Down Expand Up @@ -1301,69 +1296,6 @@ impl<K: Ord, V> BTreeMap<K, V> {

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

/// Removes empty levels on the top.
fn fix_top(&mut self) {
loop {
{
let node = self.root.as_ref().unwrap().as_ref();
if node.height() == 0 || node.len() > 0 {
break;
}
}
self.root.as_mut().unwrap().pop_level();
}
}

fn fix_right_border(&mut self) {
self.fix_top();

{
let mut cur_node = self.root.as_mut().unwrap().as_mut();

while let Internal(node) = cur_node.force() {
let mut last_kv = node.last_kv();

if last_kv.can_merge() {
cur_node = last_kv.merge().descend();
} else {
let right_len = last_kv.reborrow().right_edge().descend().len();
// `MINLEN + 1` to avoid readjust if merge happens on the next level.
if right_len < node::MIN_LEN + 1 {
last_kv.bulk_steal_left(node::MIN_LEN + 1 - right_len);
}
cur_node = last_kv.right_edge().descend();
}
}
}

self.fix_top();
}

/// The symmetric clone of `fix_right_border`.
fn fix_left_border(&mut self) {
self.fix_top();

{
let mut cur_node = self.root.as_mut().unwrap().as_mut();

while let Internal(node) = cur_node.force() {
let mut first_kv = node.first_kv();

if first_kv.can_merge() {
cur_node = first_kv.merge().descend();
} else {
let left_len = first_kv.reborrow().left_edge().descend().len();
if left_len < node::MIN_LEN + 1 {
first_kv.bulk_steal_right(node::MIN_LEN + 1 - left_len);
}
cur_node = first_kv.left_edge().descend();
}
}
}

self.fix_top();
}
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down Expand Up @@ -2321,9 +2253,9 @@ impl<K, V> BTreeMap<K, V> {
}

/// If the root node is the empty (non-allocated) root node, allocate our
/// own node.
fn ensure_root_is_owned(&mut self) -> &mut node::Root<K, V> {
self.root.get_or_insert_with(node::Root::new_leaf)
/// own node. Is an associated function to avoid borrowing the entire BTreeMap.
fn ensure_is_owned(root: &mut Option<node::Root<K, V>>) -> &mut node::Root<K, V> {
root.get_or_insert_with(node::Root::new_leaf)
}
}

Expand Down Expand Up @@ -2825,6 +2757,65 @@ 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 {
self.pop_level();
}
}

fn fix_right_border(&mut self) {
self.fix_top();

{
let mut cur_node = self.as_mut();

while let Internal(node) = cur_node.force() {
let mut last_kv = node.last_kv();

if last_kv.can_merge() {
cur_node = last_kv.merge().descend();
} else {
let right_len = last_kv.reborrow().right_edge().descend().len();
// `MINLEN + 1` to avoid readjust if merge happens on the next level.
if right_len < node::MIN_LEN + 1 {
last_kv.bulk_steal_left(node::MIN_LEN + 1 - right_len);
}
cur_node = last_kv.right_edge().descend();
}
}
}

self.fix_top();
}

/// The symmetric clone of `fix_right_border`.
fn fix_left_border(&mut self) {
self.fix_top();

{
let mut cur_node = self.as_mut();

while let Internal(node) = cur_node.force() {
let mut first_kv = node.first_kv();

if first_kv.can_merge() {
cur_node = first_kv.merge().descend();
} else {
let left_len = first_kv.reborrow().left_edge().descend().len();
if left_len < node::MIN_LEN + 1 {
first_kv.bulk_steal_right(node::MIN_LEN + 1 - left_len);
}
cur_node = first_kv.left_edge().descend();
}
}
}

self.fix_top();
}
}

enum UnderflowResult<'a, K, V> {
AtRoot,
Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
Expand Down
5 changes: 5 additions & 0 deletions src/liballoc/collections/btree/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,11 @@ unsafe impl<K: Sync, V: Sync> Sync for Root<K, V> {}
unsafe impl<K: Send, V: Send> Send for Root<K, V> {}

impl<K, V> Root<K, V> {
/// Returns the number of levels below the root.
pub fn height(&self) -> usize {
self.height
}

/// Returns a new owned tree, with its own root node that is initially empty.
pub fn new_leaf() -> Self {
Root { node: BoxedNode::from_leaf(Box::new(unsafe { LeafNode::new() })), height: 0 }
Expand Down