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

Implement append for b-trees. #32466

Merged
merged 1 commit into from
Apr 23, 2016
Merged

Implement append for b-trees. #32466

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
189 changes: 165 additions & 24 deletions src/libcollections/btree/map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
use core::cmp::Ordering;
use core::fmt::Debug;
use core::hash::{Hash, Hasher};
use core::iter::FromIterator;
use core::iter::{FromIterator, Peekable};
use core::marker::PhantomData;
use core::ops::Index;
use core::{fmt, intrinsics, mem, ptr};
Expand Down Expand Up @@ -348,6 +348,12 @@ pub struct OccupiedEntry<'a, K: 'a, V: 'a> {
_marker: PhantomData<&'a mut (K, V)>,
}

// An iterator for merging two sorted sequences into one
struct MergeIter<K, V, I: Iterator<Item=(K, V)>> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As a generally useful concept, it seems weird that MergeIter would be in map.rs. That said, I can't really think of anywhere else that it might go, so I don't know if anything should be changed here.

left: Peekable<I>,
right: Peekable<I>,
}

impl<K: Ord, V> BTreeMap<K, V> {
/// Makes a new empty BTreeMap with a reasonable choice for B.
///
Expand Down Expand Up @@ -535,6 +541,62 @@ impl<K: Ord, V> BTreeMap<K, V> {
}
}

/// Moves all elements from `other` into `Self`, leaving `other` empty.
///
/// # Examples
///
/// ```
/// #![feature(btree_append)]
/// use std::collections::BTreeMap;
///
/// let mut a = BTreeMap::new();
/// a.insert(1, "a");
/// a.insert(2, "b");
/// a.insert(3, "c");
///
/// let mut b = BTreeMap::new();
/// b.insert(3, "d");
/// b.insert(4, "e");
/// b.insert(5, "f");
///
/// a.append(&mut b);
///
/// assert_eq!(a.len(), 5);
/// assert_eq!(b.len(), 0);
///
/// assert_eq!(a[&1], "a");
/// assert_eq!(a[&2], "b");
/// assert_eq!(a[&3], "d");
/// assert_eq!(a[&4], "e");
/// assert_eq!(a[&5], "f");
/// ```
#[unstable(feature = "btree_append", reason = "recently added as part of collections reform 2",
issue = "19986")]
pub fn append(&mut self, other: &mut Self) {
// Do we have to append anything at all?
if other.len() == 0 {
return;
}

// We can just swap `self` and `other` if `self` is empty.
if self.len() == 0 {
mem::swap(self, other);
return;
}

// First, we merge `self` and `other` into a sorted sequence in linear time.
let self_iter = mem::replace(self, BTreeMap::new()).into_iter();
let other_iter = mem::replace(other, BTreeMap::new()).into_iter();
let iter = MergeIter {
left: self_iter.peekable(),
right: other_iter.peekable(),
};

// 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, starting
/// at min, and ending at max. If min is `Unbounded`, then it will be treated as "negative
/// infinity", and if max is `Unbounded`, then it will be treated as "positive infinity".
Expand Down Expand Up @@ -724,6 +786,76 @@ impl<K: Ord, V> BTreeMap<K, V> {
})
}
}

fn from_sorted_iter<I: Iterator<Item=(K, V)>>(&mut self, iter: I) {
let mut cur_node = last_leaf_edge(self.root.as_mut()).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.
if cur_node.len() < node::CAPACITY {
cur_node.push(key, value);
} else {
// No space left, go up and push there.
let mut open_node;
let mut test_node = cur_node.forget_type();
loop {
match test_node.ascend() {
Ok(parent) => {
let parent = parent.into_node();
if parent.len() < node::CAPACITY {
// Found a node with space left, push here.
open_node = parent;
break;
} else {
// Go up again.
test_node = parent.forget_type();
}
},
Err(node) => {
// We are at the top, create a new root node and push there.
open_node = node.into_root_mut().push_level();
break;
},
}
}

// Push key-value pair and new right subtree.
let tree_height = open_node.height() - 1;
let mut right_tree = node::Root::new_leaf();
for _ in 0..tree_height {
right_tree.push_level();
}
open_node.push(key, value, right_tree);

// Go down to the right-most leaf again.
cur_node = last_leaf_edge(open_node.forget_type()).into_node();
}

self.length += 1;
}
}

fn fix_right_edge(&mut self) {
// Handle underfull nodes, start from the top.
let mut cur_node = self.root.as_mut();
while let Internal(internal) = cur_node.force() {
// Check if right-most child is underfull.
let mut last_edge = internal.last_edge();
let right_child_len = last_edge.reborrow().descend().len();
if right_child_len < node::CAPACITY / 2 {
// We need to steal.
let mut last_kv = match last_edge.left_kv() {
Ok(left) => left,
Err(_) => unreachable!(),
};
last_kv.bulk_steal_left(node::CAPACITY/2 - right_child_len);
last_edge = last_kv.right_edge();
}

// Go further down.
cur_node = last_edge.descend();
}
}
}

impl<'a, K: 'a, V: 'a> IntoIterator for &'a BTreeMap<K, V> {
Expand Down Expand Up @@ -1690,32 +1822,41 @@ fn handle_underfull_node<'a, K, V>(node: NodeRef<marker::Mut<'a>,
};

if handle.can_merge() {
return Merged(handle.merge().into_node());
Merged(handle.merge().into_node())
} else {
unsafe {
let (k, v, edge) = if is_left {
handle.reborrow_mut().left_edge().descend().pop()
} else {
handle.reborrow_mut().right_edge().descend().pop_front()
};
if is_left {
handle.steal_left();
} else {
handle.steal_right();
}
Stole(handle.into_node())
}
}

let k = mem::replace(handle.reborrow_mut().into_kv_mut().0, k);
let v = mem::replace(handle.reborrow_mut().into_kv_mut().1, v);
impl<K: Ord, V, I: Iterator<Item=(K, V)>> Iterator for MergeIter<K, V, I> {
type Item = (K, V);

// FIXME: reuse cur_node?
if is_left {
match handle.reborrow_mut().right_edge().descend().force() {
Leaf(mut leaf) => leaf.push_front(k, v),
Internal(mut internal) => internal.push_front(k, v, edge.unwrap())
}
} else {
match handle.reborrow_mut().left_edge().descend().force() {
Leaf(mut leaf) => leaf.push(k, v),
Internal(mut internal) => internal.push(k, v, edge.unwrap())
}
}
}
fn next(&mut self) -> Option<(K, V)> {
let res = match (self.left.peek(), self.right.peek()) {
(Some(&(ref left_key, _)), Some(&(ref right_key, _))) => left_key.cmp(right_key),
(Some(_), None) => Ordering::Less,
(None, Some(_)) => Ordering::Greater,
(None, None) => return None,
};

return Stole(handle.into_node());
// Check which elements comes first and only advance the corresponding iterator.
// If two keys are equal, take the value from `right`.
match res {
Ordering::Less => {
self.left.next()
},
Ordering::Greater => {
self.right.next()
},
Ordering::Equal => {
self.left.next();
self.right.next()
},
}
}
}
139 changes: 139 additions & 0 deletions src/libcollections/btree/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,12 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
self.as_leaf().len as usize
}

/// Returns the height of this node in the whole tree. Zero height denotes the
/// leaf level.
pub fn height(&self) -> usize {
self.height
}

/// Removes any static information about whether this node is a `Leaf` or an
/// `Internal` node.
pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
Expand Down Expand Up @@ -1233,6 +1239,139 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
Handle::new_edge(self.node, self.idx)
}
}

/// This removes a key/value pair from the left child and replaces it with the key/value pair
/// pointed to by this handle while pushing the old key/value pair of this handle into the right
/// child.
pub fn steal_left(&mut self) {
unsafe {
let (k, v, edge) = self.reborrow_mut().left_edge().descend().pop();

let k = mem::replace(self.reborrow_mut().into_kv_mut().0, k);
let v = mem::replace(self.reborrow_mut().into_kv_mut().1, v);

match self.reborrow_mut().right_edge().descend().force() {
ForceResult::Leaf(mut leaf) => leaf.push_front(k, v),
ForceResult::Internal(mut internal) => internal.push_front(k, v, edge.unwrap())
}
}
}

/// This removes a key/value pair from the right child and replaces it with the key/value pair
/// pointed to by this handle while pushing the old key/value pair of this handle into the left
/// child.
pub fn steal_right(&mut self) {
unsafe {
let (k, v, edge) = self.reborrow_mut().right_edge().descend().pop_front();

let k = mem::replace(self.reborrow_mut().into_kv_mut().0, k);
let v = mem::replace(self.reborrow_mut().into_kv_mut().1, v);

match self.reborrow_mut().left_edge().descend().force() {
ForceResult::Leaf(mut leaf) => leaf.push(k, v),
ForceResult::Internal(mut internal) => internal.push(k, v, edge.unwrap())
}
}
}

/// This does stealing similar to `steal_left` but steals multiple elements at once.
pub fn bulk_steal_left(&mut self, n: usize) {
unsafe {
// Get raw pointers to left child's keys, values and edges.
let (left_len, left_k, left_v, left_e) = {
let mut left = self.reborrow_mut().left_edge().descend();

(left.len(),
left.keys_mut().as_mut_ptr(),
left.vals_mut().as_mut_ptr(),
match left.force() {
ForceResult::Leaf(_) => None,
ForceResult::Internal(mut i) => Some(i.as_internal_mut().edges.as_mut_ptr()),
})
};

// Get raw pointers to right child's keys, values and edges.
let (right_len, right_k, right_v, right_e) = {
let mut right = self.reborrow_mut().right_edge().descend();

(right.len(),
right.keys_mut().as_mut_ptr(),
right.vals_mut().as_mut_ptr(),
match right.force() {
ForceResult::Leaf(_) => None,
ForceResult::Internal(mut i) => Some(i.as_internal_mut().edges.as_mut_ptr()),
})
};

// Get raw pointers to parent's key and value.
let (parent_k, parent_v) = {
let kv = self.reborrow_mut().into_kv_mut();
(kv.0 as *mut K, kv.1 as *mut V)
};

// Make sure that we may steal safely.
debug_assert!(right_len + n <= CAPACITY);
debug_assert!(left_len >= n);

// Make room for stolen elements in right child.
ptr::copy(right_k,
right_k.offset(n as isize),
right_len);
ptr::copy(right_v,
right_v.offset(n as isize),
right_len);
if let Some(edges) = right_e {
ptr::copy(edges,
edges.offset(n as isize),
right_len+1);
}

// Move elements from the left child to the right one.
let left_ind = (left_len - n) as isize;
ptr::copy_nonoverlapping(left_k.offset(left_ind + 1),
right_k,
n - 1);
ptr::copy_nonoverlapping(left_v.offset(left_ind + 1),
right_v,
n - 1);
match (left_e, right_e) {
(Some(left), Some(right)) => {
ptr::copy_nonoverlapping(left.offset(left_ind + 1),
right,
n);
},
(Some(_), None) => unreachable!(),
(None, Some(_)) => unreachable!(),
(None, None) => {},
}

// Copy parent key/value pair to right child.
ptr::copy_nonoverlapping(parent_k,
right_k.offset(n as isize - 1),
1);
ptr::copy_nonoverlapping(parent_v,
right_v.offset(n as isize - 1),
1);
// Copy left-most stolen pair to parent.
ptr::copy_nonoverlapping(left_k.offset(left_ind),
parent_k,
1);
ptr::copy_nonoverlapping(left_v.offset(left_ind),
parent_v,
1);

// Fix lengths of left and right child and parent pointers in children of the right
// child.
self.reborrow_mut().left_edge().descend().as_leaf_mut().len -= n as u16;
let mut right = self.reborrow_mut().right_edge().descend();
right.as_leaf_mut().len += n as u16;
if let ForceResult::Internal(mut node) = right.force() {
for i in 0..(right_len+n+1) {
Handle::new_edge(node.reborrow_mut(), i as usize).correct_parent_link();
}
}
}
}
}

impl<BorrowType, K, V, HandleType>
Expand Down
Loading