BTreeMap navigation done safer & faster

src/liballoc/benches/btree/set.rs

@@ -50,43 +50,112 @@ macro_rules! set_bench {
     };
 }
 
-const BUILD_SET_SIZE: usize = 100;
+#[bench]
+pub fn clone_100(b: &mut Bencher) {
+    let src = pos(100);
+    b.iter(|| src.clone())
+}
 
 #[bench]
-pub fn build_and_clear(b: &mut Bencher) {
-    b.iter(|| pos(BUILD_SET_SIZE).clear())
+pub fn clone_100_and_clear(b: &mut Bencher) {
+    let src = pos(100);
+    b.iter(|| src.clone().clear())
 }
 
 #[bench]
-pub fn build_and_drop(b: &mut Bencher) {
-    b.iter(|| pos(BUILD_SET_SIZE))
+pub fn clone_100_and_into_iter(b: &mut Bencher) {
+    let src = pos(100);
+    b.iter(|| src.clone().into_iter().count())
 }
 
 #[bench]
-pub fn build_and_into_iter(b: &mut Bencher) {
-    b.iter(|| pos(BUILD_SET_SIZE).into_iter().count())
+pub fn clone_100_and_pop_all(b: &mut Bencher) {
+    let src = pos(100);
+    b.iter(|| {
+        let mut set = src.clone();
+        while set.pop_first().is_some() {}
+        set
+    });
 }
 
 #[bench]
-pub fn build_and_pop_all(b: &mut Bencher) {
+pub fn clone_100_and_remove_all(b: &mut Bencher) {
+    let src = pos(100);
     b.iter(|| {
-        let mut s = pos(BUILD_SET_SIZE);
-        while s.pop_first().is_some() {}
-        s
+        let mut set = src.clone();
+        while let Some(elt) = set.iter().copied().next() {
+            set.remove(&elt);
+        }
+        set
     });
 }
 
 #[bench]
-pub fn build_and_remove_all(b: &mut Bencher) {
+pub fn clone_100_and_remove_half(b: &mut Bencher) {
+    let src = pos(100);
     b.iter(|| {
-        let mut s = pos(BUILD_SET_SIZE);
-        while let Some(elt) = s.iter().copied().next() {
-            s.remove(&elt);
+        let mut set = src.clone();
+        for i in (2..=100 as i32).step_by(2) {
+            set.remove(&i);
         }
-        s
+        assert_eq!(set.len(), 100 / 2);
+        set
+    })
+}
+
+#[bench]
+pub fn clone_10k(b: &mut Bencher) {
+    let src = pos(10_000);
+    b.iter(|| src.clone())
+}
+
+#[bench]
+pub fn clone_10k_and_clear(b: &mut Bencher) {
+    let src = pos(10_000);
+    b.iter(|| src.clone().clear())
+}
+
+#[bench]
+pub fn clone_10k_and_into_iter(b: &mut Bencher) {
+    let src = pos(10_000);
+    b.iter(|| src.clone().into_iter().count())
+}
+
+#[bench]
+pub fn clone_10k_and_pop_all(b: &mut Bencher) {
+    let src = pos(10_000);
+    b.iter(|| {
+        let mut set = src.clone();
+        while set.pop_first().is_some() {}
+        set
+    });
+}
+
+#[bench]
+pub fn clone_10k_and_remove_all(b: &mut Bencher) {
+    let src = pos(10_000);
+    b.iter(|| {
+        let mut set = src.clone();
+        while let Some(elt) = set.iter().copied().next() {
+            set.remove(&elt);
+        }
+        set
     });
 }
 
+#[bench]
+pub fn clone_10k_and_remove_half(b: &mut Bencher) {
+    let src = pos(10_000);
+    b.iter(|| {
+        let mut set = src.clone();
+        for i in (2..=10_000 as i32).step_by(2) {
+            set.remove(&i);
+        }
+        assert_eq!(set.len(), 10_000 / 2);
+        set
+    })
+}
+
 set_bench! {intersection_100_neg_vs_100_pos, intersection, count, [neg(100), pos(100)]}
 set_bench! {intersection_100_neg_vs_10k_pos, intersection, count, [neg(100), pos(10_000)]}
 set_bench! {intersection_100_pos_vs_100_neg, intersection, count, [pos(100), neg(100)]}

src/liballoc/collections/btree/map.rs

@@ -675,13 +675,15 @@ impl<K: Ord, V> BTreeMap<K, V> {
         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 @@ impl<K: Ord, V> BTreeMap<K, V> {
         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
         }
     }
 
@@ -1469,16 +1473,13 @@ impl<K, V> Drop for IntoIter<K, V> {
     fn drop(&mut self) {
         self.for_each(drop);
         unsafe {
-            let leaf_node = ptr::read(&self.front).into_node();
-            if leaf_node.is_shared_root() {
+            let mut node = ptr::read(&self.front).into_node().forget_type();
+            if node.is_shared_root() {
                 return;
             }
 
-            if let Some(first_parent) = leaf_node.deallocate_and_ascend() {
-                let mut cur_internal_node = first_parent.into_node();
-                while let Some(parent) = cur_internal_node.deallocate_and_ascend() {
-                    cur_internal_node = parent.into_node()
-                }
+            while let Some(parent) = node.deallocate_and_ascend() {
+                node = parent.into_node().forget_type();
             }
         }
     }