Fix and test implementation of BTreeMap's first_entry, last_entry, pop_first, pop_last
diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
index 8eabc17..c1778f2 100644
--- a/src/liballoc/collections/btree/map.rs
+++ b/src/liballoc/collections/btree/map.rs
@@ -675,13 +675,15 @@
T: Ord,
K: Borrow<T>,
{
- match self.length {
- 0 => None,
- _ => Some(OccupiedEntry {
- handle: self.root.as_mut().first_kv(),
+ let front = self.root.as_mut().first_leaf_edge();
+ if let Ok(kv) = front.right_kv() {
+ Some(OccupiedEntry {
+ handle: kv.forget_node_type(),
length: &mut self.length,
_marker: PhantomData,
- }),
+ })
+ } else {
+ None
}
}
@@ -736,13 +738,15 @@
T: Ord,
K: Borrow<T>,
{
- match self.length {
- 0 => None,
- _ => Some(OccupiedEntry {
- handle: self.root.as_mut().last_kv(),
+ let back = self.root.as_mut().last_leaf_edge();
+ if let Ok(kv) = back.left_kv() {
+ Some(OccupiedEntry {
+ handle: kv.forget_node_type(),
length: &mut self.length,
_marker: PhantomData,
- }),
+ })
+ } else {
+ None
}
}
diff --git a/src/liballoc/tests/btree/map.rs b/src/liballoc/tests/btree/map.rs
index 0d00950..0a26d7b 100644
--- a/src/liballoc/tests/btree/map.rs
+++ b/src/liballoc/tests/btree/map.rs
@@ -23,6 +23,11 @@
assert_eq!(map.len(), i + 1);
}
+ assert_eq!(map.first_key_value(), Some((&0, &0)));
+ assert_eq!(map.last_key_value(), Some((&(size - 1), &(10 * (size - 1)))));
+ assert_eq!(map.first_entry().unwrap().key(), &0);
+ assert_eq!(map.last_entry().unwrap().key(), &(size - 1));
+
for i in 0..size {
assert_eq!(map.get(&i).unwrap(), &(i * 10));
}
diff --git a/src/liballoc/tests/btree/set.rs b/src/liballoc/tests/btree/set.rs
index 265ef75..1a2b62d 100644
--- a/src/liballoc/tests/btree/set.rs
+++ b/src/liballoc/tests/btree/set.rs
@@ -487,21 +487,26 @@
a.insert(2);
assert_eq!(a.first(), Some(&1));
assert_eq!(a.last(), Some(&2));
- a.insert(3);
+ for i in 3..=12 {
+ a.insert(i);
+ }
assert_eq!(a.first(), Some(&1));
- assert_eq!(a.last(), Some(&3));
-
- assert_eq!(a.len(), 3);
+ assert_eq!(a.last(), Some(&12));
assert_eq!(a.pop_first(), Some(1));
- assert_eq!(a.len(), 2);
- assert_eq!(a.pop_last(), Some(3));
- assert_eq!(a.len(), 1);
+ assert_eq!(a.pop_last(), Some(12));
assert_eq!(a.pop_first(), Some(2));
- assert_eq!(a.len(), 0);
- assert_eq!(a.pop_last(), None);
- assert_eq!(a.len(), 0);
+ assert_eq!(a.pop_last(), Some(11));
+ assert_eq!(a.pop_first(), Some(3));
+ assert_eq!(a.pop_last(), Some(10));
+ assert_eq!(a.pop_first(), Some(4));
+ assert_eq!(a.pop_first(), Some(5));
+ assert_eq!(a.pop_first(), Some(6));
+ assert_eq!(a.pop_first(), Some(7));
+ assert_eq!(a.pop_first(), Some(8));
+ assert_eq!(a.clone().pop_last(), Some(9));
+ assert_eq!(a.pop_first(), Some(9));
assert_eq!(a.pop_first(), None);
- assert_eq!(a.len(), 0);
+ assert_eq!(a.pop_last(), None);
}
fn rand_data(len: usize) -> Vec<u32> {