From 9b903bf49c9bf0f777fc3bc2b83007241084a254 Mon Sep 17 00:00:00 2001 From: Johannes Oertel Date: Thu, 21 May 2015 12:26:18 +0200 Subject: [PATCH 1/2] Implement append and split_off for BTreeMap and BTreeSet Changes the internal SearchStack API to return the key on removal as well. --- src/libcollections/btree/map.rs | 147 +++++++++++++++++++++++++--- src/libcollections/btree/node.rs | 24 +++++ src/libcollections/btree/set.rs | 68 +++++++++++++ src/libcollectionstest/btree/map.rs | 99 +++++++++++++++++++ src/libcollectionstest/btree/set.rs | 99 +++++++++++++++++++ 5 files changed, 426 insertions(+), 11 deletions(-) diff --git a/src/libcollections/btree/map.rs b/src/libcollections/btree/map.rs index 27b10213ecd7c..331bcba28ac44 100644 --- a/src/libcollections/btree/map.rs +++ b/src/libcollections/btree/map.rs @@ -348,7 +348,7 @@ impl BTreeMap { loop { let result = stack.with(move |pusher, node| { - // Same basic logic as found in `find`, but with PartialSearchStack mediating the + // Same basic logic as found in `get`, but with PartialSearchStack mediating the // actual nodes for us match Node::search(node, &key) { Found(mut handle) => { @@ -438,14 +438,14 @@ impl BTreeMap { /// ``` #[stable(feature = "rust1", since = "1.0.0")] pub fn remove(&mut self, key: &Q) -> Option where K: Borrow, Q: Ord { - // See `swap` for a more thorough description of the stuff going on in here + // See `insert` for a more thorough description of the stuff going on in here let mut stack = stack::PartialSearchStack::new(self); loop { let result = stack.with(move |pusher, node| { match Node::search(node, key) { Found(handle) => { // Perfect match. Terminate the stack here, and remove the entry - Finished(Some(pusher.seal(handle).remove())) + Finished(Some(pusher.seal(handle).remove().1)) }, GoDown(handle) => { // We need to keep searching, try to go down the next edge @@ -463,6 +463,131 @@ impl BTreeMap { } } } + + /// Moves all elements from `other` into `Self`, leaving `other` empty. + /// + /// # Examples + /// + /// ``` + /// # #![feature(btree_append_split_off)] + /// 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 = "append", + reason = "recently added as part of collections reform 2")] + pub fn append(&mut self, other: &mut Self) { + let b = other.b; + for (key, value) in mem::replace(other, BTreeMap::with_b(b)) { + self.insert(key, value); + } + } + + /// Splits the map into two at the given key, + /// retaining the first half in-place and returning the second one. + /// + /// # Examples + /// + /// ``` + /// # #![feature(btree_append_split_off)] + /// use std::collections::BTreeMap; + /// + /// a.insert(1, "a"); + /// a.insert(2, "b"); + /// a.insert(3, "c"); + /// a.insert(4, "d"); + /// a.insert(5, "e"); + /// + /// let b = a.split_off(3); + /// + /// assert_eq!(a.len(), 2); + /// assert_eq!(b.len(), 3); + /// + /// assert_eq!(a[&1], "a"); + /// assert_eq!(a[&2], "b"); + /// assert_eq!(b[&3], "c"); + /// assert_eq!(b[&4], "d"); + /// assert_eq!(b[&5], "e"); + /// ``` + #[unstable(feature = "split_off", + reason = "recently added as part of collections reform 2")] + pub fn split_off(&mut self, at: &Q) -> Self where K: Borrow, Q: Ord { + let mut other = BTreeMap::new(); + + if self.len() == 0 { + return other; + } + + // FIXME(RFC #811) We can't check for `at` pointing before the + // first element and then swap `self` and `other`, because + // `self` will still be borrowed immutably. + // `unwrap` won't panic because `self.len()` > 0. + let should_swap = at <= self.keys().next().unwrap().borrow(); + + // Does `at` point before the first element? + if should_swap { + mem::swap(self, &mut other); + return other; + } + // Does `at` point behind the last element? + // `unwrap` won't panic because `self.len()` > 0. + else if at > self.keys().rev().next().unwrap().borrow() { + return other; + } + + let mut remove_greater_or_equal = || { + let mut stack = stack::PartialSearchStack::new(self); + loop { + let result = stack.with(move |pusher, node| { + match Node::greater_or_equal(node, at) { + Found(handle) => { + // Found a matching key. Terminate the stack here, and remove the entry + Finished(Some(pusher.seal(handle).remove())) + }, + GoDown(handle) => { + // We need to keep searching, try to go down the next edge + match handle.force() { + // We're at a leaf; no matching key found + Leaf(_) => Finished(None), + Internal(internal_handle) => Continue(pusher.push(internal_handle)) + } + } + } + }); + match result { + Finished(ret) => return ret, + Continue(new_stack) => stack = new_stack + } + } + }; + + // Remove and move all elements greater than or equal to at + loop { + match remove_greater_or_equal() { + Some((key, value)) => other.insert(key, value), + None => return other, + }; + } + } } #[stable(feature = "rust1", since = "1.0.0")] @@ -693,16 +818,16 @@ mod stack { impl<'a, K, V> SearchStack<'a, K, V, handle::KV, handle::Leaf> { /// Removes the key and value in the top element of the stack, then handles underflows as /// described in BTree's pop function. - fn remove_leaf(mut self) -> V { + fn remove_leaf(mut self) -> (K, V) { self.map.length -= 1; // Remove the key-value pair from the leaf that this search stack points to. // Then, note if the leaf is underfull, and promptly forget the leaf and its ptr // to avoid ownership issues. - let (value, mut underflow) = unsafe { - let (_, value) = self.top.from_raw_mut().remove_as_leaf(); + let (key, value, mut underflow) = unsafe { + let (key, value) = self.top.from_raw_mut().remove_as_leaf(); let underflow = self.top.from_raw().node().is_underfull(); - (value, underflow) + (key, value, underflow) }; loop { @@ -717,7 +842,7 @@ mod stack { self.map.depth -= 1; self.map.root.hoist_lone_child(); } - return value; + return (key, value); } Some(mut handle) => { if underflow { @@ -728,7 +853,7 @@ mod stack { } } else { // All done! - return value; + return (key, value); } } } @@ -739,7 +864,7 @@ mod stack { impl<'a, K, V> SearchStack<'a, K, V, handle::KV, handle::LeafOrInternal> { /// Removes the key and value in the top element of the stack, then handles underflows as /// described in BTree's pop function. - pub fn remove(self) -> V { + pub fn remove(self) -> (K, V) { // Ensure that the search stack goes to a leaf. This is necessary to perform deletion // in a BTree. Note that this may put the tree in an inconsistent state (further // described in into_leaf's comments), but this is immediately fixed by the @@ -1220,7 +1345,7 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> { /// Takes the value of the entry out of the map, and returns it. #[stable(feature = "rust1", since = "1.0.0")] pub fn remove(self) -> V { - self.stack.remove() + self.stack.remove().1 } } diff --git a/src/libcollections/btree/node.rs b/src/libcollections/btree/node.rs index 4d76a986700a2..6882c2c0276a6 100644 --- a/src/libcollections/btree/node.rs +++ b/src/libcollections/btree/node.rs @@ -557,6 +557,15 @@ impl Node { (index, false) => GoDown(Handle { node: node, index: index, marker: PhantomData }), } } + + /// Searches for the first key greater than or equal to `key`. + pub fn greater_or_equal>>(node: NodeRef, key: &Q) + -> SearchResult where K: Borrow, Q: Ord { + match node.as_slices_internal().greater_or_equal_linear(key) { + (index, true) => Found(Handle { node: node, index: index, marker: PhantomData }), + (index, false) => GoDown(Handle { node: node, index: index, marker: PhantomData }), + } + } } // Public interface @@ -1609,3 +1618,18 @@ macro_rules! node_slice_impl { node_slice_impl!(NodeSlice, Traversal, as_slices_internal, index, iter); node_slice_impl!(MutNodeSlice, MutTraversal, as_slices_internal_mut, index_mut, iter_mut); + +impl<'a, K: Ord + 'a, V: 'a> NodeSlice<'a, K, V> { + /// Performs linear search in a slice to find the first element greater than + /// or equal to `key`. + fn greater_or_equal_linear(&self, key: &Q) -> (usize, bool) + where K: Borrow, Q: Ord { + for (i, k) in self.keys.iter().enumerate() { + match key.cmp(k.borrow()) { + Greater => {} + Less | Equal => return (i, true), + } + } + (self.keys.len(), false) + } +} diff --git a/src/libcollections/btree/set.rs b/src/libcollections/btree/set.rs index 7c4cda305adf2..53a6057ca6c30 100644 --- a/src/libcollections/btree/set.rs +++ b/src/libcollections/btree/set.rs @@ -453,6 +453,74 @@ impl BTreeSet { pub fn remove(&mut self, value: &Q) -> bool where T: Borrow, Q: Ord { self.map.remove(value).is_some() } + + /// Moves all elements from `other` into `Self`, leaving `other` empty. + /// + /// # Examples + /// + /// ``` + /// # #![feature(btree_append_split_off)] + /// use std::collections::BTreeMap; + /// + /// let mut a = BTreeSet::new(); + /// a.insert(1); + /// a.insert(2); + /// a.insert(3); + /// + /// let mut b = BTreeSet::new(); + /// b.insert(3); + /// b.insert(4); + /// b.insert(5); + /// + /// a.append(&mut b); + /// + /// assert_eq!(a.len(), 5); + /// assert_eq!(b.len(), 0); + /// + /// assert!(a.contains(&1)); + /// assert!(a.contains(&2)); + /// assert!(a.contains(&3)); + /// assert!(a.contains(&4)); + /// assert!(a.contains(&5)); + /// ``` + #[unstable(feature = "btree_append_split_off", + reason = "recently added as part of collections reform 2")] + pub fn append(&mut self, other: &mut Self) { + self.map.append(&mut other.map); + } + + /// Splits the set into two at the given key, + /// retaining the first half in-place and returning the second one. + /// + /// # Examples + /// + /// ``` + /// # #![feature(btree_append_split_off)] + /// use std::collections::BTreeMap; + /// + /// a.insert(1) + /// a.insert(2) + /// a.insert(3) + /// a.insert(4) + /// a.insert(5) + /// + /// let b = a.split_off(3); + /// + /// assert_eq!(a.len(), 2); + /// assert_eq!(b.len(), 3); + /// + /// assert!(a.contains(&1)); + /// assert!(a.contains(&2)); + /// assert!(b.contains(&3)); + /// assert!(b.contains(&4)); + /// assert!(b.contains(&5)); + /// ``` + #[unstable(feature = "btree_append_split_off", + reason = "recently added as part of collections reform 2")] + pub fn split_off(&mut self, at: &Q) -> Self + where T: Borrow, Q: Ord { + BTreeSet { map: self.map.split_off(at) } + } } #[stable(feature = "rust1", since = "1.0.0")] diff --git a/src/libcollectionstest/btree/map.rs b/src/libcollectionstest/btree/map.rs index 62b46433da957..a53722920bdac 100644 --- a/src/libcollectionstest/btree/map.rs +++ b/src/libcollectionstest/btree/map.rs @@ -294,6 +294,105 @@ fn test_extend_ref() { assert_eq!(a[&3], "three"); } +#[test] +fn test_append() { + 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"); // Overwrite element from a + 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"); +} + +#[test] +fn test_split_off() { + // Split empty map + let mut a: BTreeMap = BTreeMap::new(); + + let b = a.split_off(&2); + + assert_eq!(a.len(), 0); + assert_eq!(b.len(), 0); + + // Split before first element + let mut a = BTreeMap::new(); + a.insert(4, "a"); + a.insert(5, "b"); + a.insert(6, "c"); + + let b = a.split_off(&2); + + assert_eq!(a.len(), 0); + assert_eq!(b.len(), 3); + + assert_eq!(b[&4], "a"); + assert_eq!(b[&5], "b"); + assert_eq!(b[&6], "c"); + + // Split at first element + let mut a = BTreeMap::new(); + a.insert(4, "a"); + a.insert(5, "b"); + a.insert(6, "c"); + + let b = a.split_off(&4); + + assert_eq!(a.len(), 0); + assert_eq!(b.len(), 3); + + assert_eq!(b[&4], "a"); + assert_eq!(b[&5], "b"); + assert_eq!(b[&6], "c"); + + // Split behind last element + let mut a = BTreeMap::new(); + a.insert(1, "a"); + a.insert(2, "b"); + a.insert(3, "c"); + + let b = a.split_off(&4); + + assert_eq!(a.len(), 3); + assert_eq!(b.len(), 0); + + assert_eq!(a[&1], "a"); + assert_eq!(a[&2], "b"); + assert_eq!(a[&3], "c"); + + // Split at arbitrary position + let mut a = BTreeMap::new(); + a.insert(1, "a"); + a.insert(2, "b"); + a.insert(3, "c"); + a.insert(4, "d"); + a.insert(5, "e"); + + let b = a.split_off(&3); + + assert_eq!(a.len(), 2); + assert_eq!(b.len(), 3); + + assert_eq!(a[&1], "a"); + assert_eq!(a[&2], "b"); + assert_eq!(b[&3], "c"); + assert_eq!(b[&4], "d"); + assert_eq!(b[&5], "e"); +} + mod bench { use std::collections::BTreeMap; use std::__rand::{Rng, thread_rng}; diff --git a/src/libcollectionstest/btree/set.rs b/src/libcollectionstest/btree/set.rs index 56257233aa558..f330903999dd2 100644 --- a/src/libcollectionstest/btree/set.rs +++ b/src/libcollectionstest/btree/set.rs @@ -212,3 +212,102 @@ fn test_extend_ref() { assert!(a.contains(&5)); assert!(a.contains(&6)); } + +#[test] +fn test_append() { + let mut a = BTreeSet::new(); + a.insert(1); + a.insert(2); + a.insert(3); + + let mut b = BTreeSet::new(); + b.insert(3); + b.insert(4); + b.insert(5); + + a.append(&mut b); + + assert_eq!(a.len(), 5); + assert_eq!(b.len(), 0); + + assert_eq!(a.contains(&1), true); + assert_eq!(a.contains(&2), true); + assert_eq!(a.contains(&3), true); + assert_eq!(a.contains(&4), true); + assert_eq!(a.contains(&5), true); +} + +#[test] +fn test_split_off() { + // Split empty set + let mut a: BTreeSet = BTreeMap::new(); + + let b = a.split_off(2); + + assert_eq!(a.len(), 0); + assert_eq!(b.len(), 0); + + // Split before first element + let mut a = BTreeSet::new(); + a.insert(4); + a.insert(5); + a.insert(6); + + let b = a.split_off(2); + + assert_eq!(a.len(), 0); + assert_eq!(b.len(), 3); + + assert!(b.contains(&4)); + assert!(b.contains(&5)); + assert!(b.contains(&6)); + + // Split at first element + let mut a = BTreeSet::new(); + a.insert(4); + a.insert(5); + a.insert(6); + + let b = a.split_off(4); + + assert_eq!(a.len(), 0); + assert_eq!(b.len(), 3); + + assert_eq!(b.contains(&4)); + assert_eq!(b.contains(&5)); + assert_eq!(b.contains(&6)); + + // Split behind last element + let mut a = BTreeSet::new(); + a.insert(1); + a.insert(2); + a.insert(3); + + let b = a.split_off(4); + + assert_eq!(a.len(), 3); + assert_eq!(b.len(), 0); + + assert_eq!(a.contains(&1)); + assert_eq!(a.contains(&2)); + assert_eq!(a.contains(&3)); + + // Split at arbitrary position + let mut a = BTreeSet::new(); + a.insert(1); + a.insert(2); + a.insert(3); + a.insert(4); + a.insert(5); + + let b = a.split_off(3); + + assert_eq!(a.len(), 2); + assert_eq!(b.len(), 3); + + assert_eq!(a.contains(&1)); + assert_eq!(a.contains(&2)); + assert_eq!(b.contains(&3)); + assert_eq!(b.contains(&4)); + assert_eq!(b.contains(&5)); +} From 5f1a116442fc326237c5c24dc4016317cde5c553 Mon Sep 17 00:00:00 2001 From: Johannes Oertel Date: Sat, 27 Jun 2015 15:49:47 +0200 Subject: [PATCH 2/2] Add linear time implementation of `append`. --- src/libcollections/btree/map.rs | 89 ++++++++++++++- src/libcollections/btree/node.rs | 169 ++++++++++++++++++++++++++++ src/libcollectionstest/btree/map.rs | 56 ++++++--- 3 files changed, 293 insertions(+), 21 deletions(-) diff --git a/src/libcollections/btree/map.rs b/src/libcollections/btree/map.rs index 331bcba28ac44..12cd9cdfd840f 100644 --- a/src/libcollections/btree/map.rs +++ b/src/libcollections/btree/map.rs @@ -148,6 +148,13 @@ pub struct OccupiedEntry<'a, K:'a, V:'a> { stack: stack::SearchStack<'a, K, V, node::handle::KV, node::handle::LeafOrInternal>, } +struct MergeIter> { + left: I, + right: I, + left_cur: Option<(K, V)>, + right_cur: Option<(K, V)>, +} + impl BTreeMap { /// Makes a new empty BTreeMap with a reasonable choice for B. #[stable(feature = "rust1", since = "1.0.0")] @@ -496,10 +503,36 @@ impl BTreeMap { #[unstable(feature = "append", reason = "recently added as part of collections reform 2")] pub fn append(&mut self, other: &mut Self) { - let b = other.b; - for (key, value) in mem::replace(other, BTreeMap::with_b(b)) { - self.insert(key, value); + // Do we have to append anything at all? + if other.len() == 0 { + return; } + + // If the values of `b` of `self` and `other` are equal, we can just swap them if `self` is + // empty. + if self.len() == 0 && self.b == other.b { + mem::swap(self, other); + } + + // First, we merge `self` and `other` into a sorted sequence in linear time. + let self_b = self.b; + let other_b = other.b; + let mut self_iter = mem::replace(self, BTreeMap::with_b(self_b)).into_iter(); + let mut other_iter = mem::replace(other, BTreeMap::with_b(other_b)).into_iter(); + let self_cur = self_iter.next(); + let other_cur = other_iter.next(); + + // Second, we build a tree from the sorted sequence in linear time. + let (length, depth, root) = Node::from_sorted_iter(MergeIter { + left: self_iter, + right: other_iter, + left_cur: self_cur, + right_cur: other_cur, + }, self_b); + + self.length = length; + self.depth = depth; + self.root = root.unwrap(); // `unwrap` won't panic because length can't be zero. } /// Splits the map into two at the given key, @@ -644,6 +677,56 @@ impl<'a, K, V> IntoIterator for &'a mut BTreeMap { } } +// Helper enum for MergeIter +enum MergeResult { + Left, + Right, + Both, + None, +} + +impl> Iterator for MergeIter { + type Item = (K, V); + + fn next(&mut self) -> Option<(K, V)> { + let res = match (&self.left_cur, &self.right_cur) { + (&Some((ref left_key, _)), &Some((ref right_key, _))) => { + match left_key.cmp(right_key) { + Ordering::Less => MergeResult::Left, + Ordering::Equal => MergeResult::Both, + Ordering::Greater => MergeResult::Right, + } + }, + (&Some(_), &None) => MergeResult::Left, + (&None, &Some(_)) => MergeResult::Right, + (&None, &None) => MergeResult::None, + }; + + // Check which elements comes first and only advance the corresponding iterator. + // If two keys are equal, take the value from `right`. + match res { + MergeResult::Left => { + let ret = self.left_cur.take(); + self.left_cur = self.left.next(); + ret + }, + MergeResult::Right => { + let ret = self.right_cur.take(); + self.right_cur = self.right.next(); + ret + }, + MergeResult::Both => { + let ret = self.right_cur.take(); + self.left_cur = self.left.next(); + self.right_cur = self.right.next(); + ret + }, + MergeResult::None => None, + } + } +} + + /// A helper enum useful for deciding whether to continue a loop since we can't /// return from a closure enum Continuation { diff --git a/src/libcollections/btree/node.rs b/src/libcollections/btree/node.rs index 6882c2c0276a6..52d3463f24695 100644 --- a/src/libcollections/btree/node.rs +++ b/src/libcollections/btree/node.rs @@ -27,6 +27,7 @@ use core::ptr::Unique; use core::{slice, mem, ptr, cmp, raw}; use alloc::heap::{self, EMPTY}; +use vec::Vec; use borrow::Borrow; /// Represents the result of an Insertion: either the item fit, or the node had to split @@ -904,6 +905,17 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<&'a mut Node, handle::KV, NodeType marker: PhantomData, } } + + /// Convert this handle into one pointing at the edge immediately to the right of the key/value + /// pair pointed-to by this handle. This is useful because it returns a reference with larger + /// lifetime than `right_edge`. + pub fn into_right_edge(self) -> Handle<&'a mut Node, handle::Edge, NodeType> { + Handle { + node: &mut *self.node, + index: self.index + 1, + marker: PhantomData, + } + } } impl<'a, K: 'a, V: 'a, NodeRef: Deref> + 'a, NodeType> Handle Node { } } +impl Node { + pub fn from_sorted_iter(iter: I, b: usize) -> (usize, usize, Option>) + where I: Iterator { + let capacity = capacity_from_b(b); + let minimum = min_load_from_capacity(capacity); + + // Holds the current level. + let mut num_level = 0; + // Needed to count the number of key-value pairs in `iter`. + let mut length = 0; + // `levels` contains the current node on every level, going from the leaves level to the + // root level. + let mut levels: Vec>> = Vec::new(); + + // Iterate through all key-value pairs, pushing them into nodes of appropriate size at the + // right level. + for (key, value) in iter { + // Always go down to a leaf after inserting an element into an internal node. + if num_level > 0 { + num_level = 0; + } + + loop { + // If we are in an internal node, extract node from the level below to insert it as + // edge on the level above; `unsafe` is needed for unchecked access. + let new_edge = unsafe { + if num_level > 0 { + levels.get_unchecked_mut(num_level - 1).take() + } else { + None + } + }; + + // Get current node on current level. + // If we are past the top-most level, insert a new level. + if num_level == levels.len() { + levels.push(None); + } + // If there is no node on this level, create a new node. `unsafe` + // is needed for unchecked access. + let level = unsafe { levels.get_unchecked_mut(num_level) }; + if level.is_none() { + *level = if num_level == 0 { + Some(Node::new_leaf(capacity)) + } else { + // `unsafe` is needed for `new_internal`. + unsafe { + Some(Node::new_internal(capacity)) + } + }; + } + let node = level.as_mut().unwrap(); + + // Insert edge from the level below; `unsafe` is needed for `push_edge`. + if let Some(edge) = new_edge { + unsafe { + node.push_edge(edge); + } + } + + // If node is already full, we have to go up one level before we can insert the + // key-value pair. + if !node.is_full() { + // Insert key-value pair into node; `unsafe` is needed for `push_kv`. + unsafe { + node.push_kv(key, value); + } + break; + } + num_level += 1; + } + + length += 1; + } + + // Fix "right edge" of the tree. + if levels.len() > 1 { + + num_level = 0; + while num_level < levels.len() - 1 { + // Extract node from this level or create a new one if there isn't any. `unsafe` is + // needed for unchecked access and `new_internal`. + let edge = unsafe { + match levels.get_unchecked_mut(num_level).take() { + Some(n) => n, + None => { + if num_level == 0 { + Node::new_leaf(capacity) + } else { + Node::new_internal(capacity) + } + }, + } + }; + + // Go to the level above. + num_level += 1; + + // Get node on this level, create one if there isn't any; `unsafe` is needed for + // unchecked access. + let level = unsafe { levels.get_unchecked_mut(num_level) }; + if level.is_none() { + // `unsafe` is needed for `new_internal`. + unsafe { + *level = Some(Node::new_internal(capacity)); + } + } + let mut node = level.as_mut().unwrap(); + + // Insert `edge` as new edge in `node`; `unsafe` is needed for `push_edge`. + unsafe { + node.push_edge(edge); + } + } + + // Start at the root and steal to fix underfull nodes on the "right edge" of the tree. + let root_index = levels.len() - 1; + let mut node = unsafe { levels.get_unchecked_mut(root_index).as_mut().unwrap() }; + + loop { + let mut temp_node = node; + let index = temp_node.len() - 1; + let mut handle = match temp_node.kv_handle(index).force() { + ForceResult::Internal(h) => h, + ForceResult::Leaf(_) => break, + }; + + // Check if we need to steal, i.e. is the length of the right edge less than + // `minimum`? + let right_len = handle.right_edge().node().len(); + if right_len < minimum { + // Steal! + let num_steals = minimum - right_len; + for _ in 0..num_steals { + // `unsafe` is needed for stealing. + unsafe { + handle.steal_rightward(); + } + } + } + + // Go down the right edge. + node = handle.into_right_edge().into_edge_mut(); + } + } + + // Get root node from `levels`. + let root = match levels.pop() { + Some(option) => option, + _ => None, + }; + + // Return (length, depth, root_node). + (length, levels.len(), root) + } +} + // Private implementation details impl Node { /// Node is full, so split it into two nodes, and yield the middle-most key-value pair diff --git a/src/libcollectionstest/btree/map.rs b/src/libcollectionstest/btree/map.rs index a53722920bdac..10da08f953ce0 100644 --- a/src/libcollectionstest/btree/map.rs +++ b/src/libcollectionstest/btree/map.rs @@ -294,30 +294,50 @@ fn test_extend_ref() { assert_eq!(a[&3], "three"); } -#[test] -fn test_append() { - let mut a = BTreeMap::new(); - a.insert(1, "a"); - a.insert(2, "b"); - a.insert(3, "c"); +macro_rules! create_append_test { + ($name:ident, $len:expr) => { + #[test] + fn $name() { + let mut a = BTreeMap::with_b(6); + for i in 0..8 { + a.insert(i, i); + } - let mut b = BTreeMap::new(); - b.insert(3, "d"); // Overwrite element from a - b.insert(4, "e"); - b.insert(5, "f"); + let mut b = BTreeMap::with_b(6); + for i in 5..$len { + b.insert(i, 2*i); + } - a.append(&mut b); + a.append(&mut b); - assert_eq!(a.len(), 5); - assert_eq!(b.len(), 0); + assert_eq!(a.len(), $len); + 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"); + for i in 0..$len { + if i < 5 { + assert_eq!(a[&i], i); + } else { + assert_eq!(a[&i], 2*i); + } + } + } + }; } +// These are mostly for testing the algorithm that "fixes" the right edge after insertion. +// Single node. +create_append_test!(test_append_9, 9); +// Two leafs that don't need fixing. +create_append_test!(test_append_17, 17); +// Two leafs where the second one ends up underfull and needs stealing at the end. +create_append_test!(test_append_14, 14); +// Two leafs where the first one isn't full; finish insertion at root. +create_append_test!(test_append_12, 12); +// Three levels; finish insertion at root. +create_append_test!(test_append_144, 144); +// Three levels; finish insertion at leaf without a node on the second level. +create_append_test!(test_append_145, 145); + #[test] fn test_split_off() { // Split empty map