Skip to content

Commit

Permalink
Rollup merge of rust-lang#74391 - ssomers:btree_refactor, r=Mark-Simu…
Browse files Browse the repository at this point in the history 
…lacrum

BtreeMap: superficially refactor root access

Remove or comment every unwrap in BTreeMap's main code and more.

r? @Mark-Simulacrum
  • Loading branch information
@Manishearth
Manishearth committed Jul 16, 2020
2 parents 79a9089 + b82d332 commit a1936ba
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 92 deletions.
175 changes: 83 additions & 92 deletions src/liballoc/collections/btree/map.rs
View file
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
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

0 comments on commit a1936ba

Please sign in to comment.