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();
             }
         }
     }

src/liballoc/collections/btree/navigate.rs

@@ -3,151 +3,112 @@ use core::ptr;
 use super::node::{marker, ForceResult::*, Handle, NodeRef};
 use super::unwrap_unchecked;
 
-macro_rules! def_next {
-    { unsafe fn $name:ident : $next_kv:ident $next_edge:ident $initial_leaf_edge:ident } => {
-        /// Given a leaf edge handle into an immutable tree, returns a handle to the next
-        /// leaf edge and references to the key and value between these edges.
-        /// Unsafe because the caller must ensure that the given leaf edge has a successor.
-        unsafe fn $name <'a, K: 'a, V: 'a>(
-            leaf_edge: Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
-        ) -> (Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>, &'a K, &'a V) {
-            let mut cur_handle = match leaf_edge.$next_kv() {
-                Ok(leaf_kv) => {
-                    let (k, v) = leaf_kv.into_kv();
-                    let next_leaf_edge = leaf_kv.$next_edge();
-                    return (next_leaf_edge, k, v);
-                }
-                Err(last_edge) => {
-                    let next_level = last_edge.into_node().ascend().ok();
-                    unwrap_unchecked(next_level)
-                }
-            };
-
-            loop {
-                cur_handle = match cur_handle.$next_kv() {
-                    Ok(internal_kv) => {
-                        let (k, v) = internal_kv.into_kv();
-                        let next_internal_edge = internal_kv.$next_edge();
-                        let next_leaf_edge = next_internal_edge.descend().$initial_leaf_edge();
-                        return (next_leaf_edge, k, v);
-                    }
-                    Err(last_edge) => {
-                        let next_level = last_edge.into_node().ascend().ok();
-                        unwrap_unchecked(next_level)
-                    }
-                }
+impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
+    /// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
+    /// on the right side, which is either in the same leaf node or in an ancestor node.
+    /// If the leaf edge is the last one in the tree, returns [`Result::Err`] with the root node.
+    pub fn next_kv(
+        self,
+    ) -> Result<
+        Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
+        NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
+    > {
+        let mut edge = self.forget_node_type();
+        loop {
+            edge = match edge.right_kv() {
+                Ok(internal_kv) => return Ok(internal_kv),
+                Err(last_edge) => match last_edge.into_node().ascend() {
+                    Ok(parent_edge) => parent_edge.forget_node_type(),
+                    Err(root) => return Err(root.forget_type()),
+                },
             }
         }
-    };
-}
-
-macro_rules! def_next_mut {
-    { unsafe fn $name:ident : $next_kv:ident $next_edge:ident $initial_leaf_edge:ident } => {
-        /// Given a leaf edge handle into a mutable tree, returns handles to the next
-        /// leaf edge and to the KV between these edges.
-        /// Unsafe for two reasons:
-        /// - the caller must ensure that the given leaf edge has a successor;
-        /// - both returned handles represent mutable references into the same tree
-        ///   that can easily invalidate each other, even on immutable use.
-        unsafe fn $name <'a, K: 'a, V: 'a>(
-            leaf_edge: Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>,
-        ) -> (Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>,
-              Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>) {
-            let mut cur_handle = match leaf_edge.$next_kv() {
-                Ok(leaf_kv) => {
-                    let next_leaf_edge = ptr::read(&leaf_kv).$next_edge();
-                    return (next_leaf_edge, leaf_kv.forget_node_type());
-                }
-                Err(last_edge) => {
-                    let next_level = last_edge.into_node().ascend().ok();
-                    unwrap_unchecked(next_level)
-                }
-            };
+    }
 
-            loop {
-                cur_handle = match cur_handle.$next_kv() {
-                    Ok(internal_kv) => {
-                        let next_internal_edge = ptr::read(&internal_kv).$next_edge();
-                        let next_leaf_edge = next_internal_edge.descend().$initial_leaf_edge();
-                        return (next_leaf_edge, internal_kv.forget_node_type());
-                    }
-                    Err(last_edge) => {
-                        let next_level = last_edge.into_node().ascend().ok();
-                        unwrap_unchecked(next_level)
-                    }
-                }
+    /// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
+    /// on the left side, which is either in the same leaf node or in an ancestor node.
+    /// If the leaf edge is the first one in the tree, returns [`Result::Err`] with the root node.
+    pub fn next_back_kv(
+        self,
+    ) -> Result<
+        Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
+        NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
+    > {
+        let mut edge = self.forget_node_type();
+        loop {
+            edge = match edge.left_kv() {
+                Ok(internal_kv) => return Ok(internal_kv),
+                Err(last_edge) => match last_edge.into_node().ascend() {
+                    Ok(parent_edge) => parent_edge.forget_node_type(),
+                    Err(root) => return Err(root.forget_type()),
+                },
             }
         }
-    };
+    }
 }
 
-macro_rules! def_next_dealloc {
-    { unsafe fn $name:ident : $next_kv:ident $next_edge:ident $initial_leaf_edge:ident } => {
-        /// Given a leaf edge handle into an owned tree, returns a handle to the next
-        /// leaf edge and the key and value between these edges, while deallocating
-        /// any node left behind.
+macro_rules! def_next_kv_uncheched_dealloc {
+    { unsafe fn $name:ident : $adjacent_kv:ident } => {
+        /// Given a leaf edge handle into an owned tree, returns a handle to the next KV,
+        /// while deallocating any node left behind.
         /// Unsafe for two reasons:
-        /// - the caller must ensure that the given leaf edge has a successor;
-        /// - the node pointed at by the given handle, and its ancestors, may be deallocated,
+        /// - The caller must ensure that the leaf edge is not the last one in the tree.
+        /// - The node pointed at by the given handle, and its ancestors, may be deallocated,
         ///   while the reference to those nodes in the surviving ancestors is left dangling;
-        ///   thus using the returned handle is dangerous.
+        ///   thus using the returned handle to navigate further is dangerous.
         unsafe fn $name <K, V>(
             leaf_edge: Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
-        ) -> (Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>, K, V) {
-            let mut cur_handle = match leaf_edge.$next_kv() {
-                Ok(leaf_kv) => {
-                    let k = ptr::read(leaf_kv.reborrow().into_kv().0);
-                    let v = ptr::read(leaf_kv.reborrow().into_kv().1);
-                    let next_leaf_edge = leaf_kv.$next_edge();
-                    return (next_leaf_edge, k, v);
-                }
-                Err(last_edge) => {
-                    unwrap_unchecked(last_edge.into_node().deallocate_and_ascend())
-                }
-            };
-
+        ) -> Handle<NodeRef<marker::Owned, K, V, marker::LeafOrInternal>, marker::KV> {
+            let mut edge = leaf_edge.forget_node_type();
             loop {
-                cur_handle = match cur_handle.$next_kv() {
-                    Ok(internal_kv) => {
-                        let k = ptr::read(internal_kv.reborrow().into_kv().0);
-                        let v = ptr::read(internal_kv.reborrow().into_kv().1);
-                        let next_internal_edge = internal_kv.$next_edge();
-                        let next_leaf_edge = next_internal_edge.descend().$initial_leaf_edge();
-                        return (next_leaf_edge, k, v);
-                    }
+                edge = match edge.$adjacent_kv() {
+                    Ok(internal_kv) => return internal_kv,
                     Err(last_edge) => {
-                        unwrap_unchecked(last_edge.into_node().deallocate_and_ascend())
+                        let parent_edge = last_edge.into_node().deallocate_and_ascend();
+                        unwrap_unchecked(parent_edge).forget_node_type()
                     }
                 }
             }
         }
     };
 }
 
-def_next! {unsafe fn next_unchecked: right_kv right_edge first_leaf_edge}
-def_next! {unsafe fn next_back_unchecked: left_kv left_edge last_leaf_edge}
-def_next_mut! {unsafe fn next_unchecked_mut: right_kv right_edge first_leaf_edge}
-def_next_mut! {unsafe fn next_back_unchecked_mut: left_kv left_edge last_leaf_edge}
-def_next_dealloc! {unsafe fn next_unchecked_deallocating: right_kv right_edge first_leaf_edge}
-def_next_dealloc! {unsafe fn next_back_unchecked_deallocating: left_kv left_edge last_leaf_edge}
+def_next_kv_uncheched_dealloc! {unsafe fn next_kv_unchecked_dealloc: right_kv}
+def_next_kv_uncheched_dealloc! {unsafe fn next_back_kv_unchecked_dealloc: left_kv}
+
+/// This replaces the value behind the `v` unique reference by calling the
+/// relevant function.
+///
+/// Safety: The change closure must not panic.
+#[inline]
+unsafe fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
+    let value = ptr::read(v);
+    let (new_value, ret) = change(value);
+    ptr::write(v, new_value);
+    ret
+}
 
 impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge> {
     /// Moves the leaf edge handle to the next leaf edge and returns references to the
     /// key and value in between.
     /// Unsafe because the caller must ensure that the leaf edge is not the last one in the tree.
     pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
-        let (next_edge, k, v) = next_unchecked(*self);
-        *self = next_edge;
-        (k, v)
+        replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_kv();
+            let kv = unwrap_unchecked(kv.ok());
+            (kv.next_leaf_edge(), kv.into_kv())
+        })
     }
 
     /// Moves the leaf edge handle to the previous leaf edge and returns references to the
     /// key and value in between.
     /// Unsafe because the caller must ensure that the leaf edge is not the first one in the tree.
     pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
-        let (next_edge, k, v) = next_back_unchecked(*self);
-        *self = next_edge;
-        (k, v)
+        replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_back_kv();
+            let kv = unwrap_unchecked(kv.ok());
+            (kv.next_back_leaf_edge(), kv.into_kv())
+        })
     }
 }
 
@@ -158,8 +119,11 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
     /// - The caller must ensure that the leaf edge is not the last one in the tree.
     /// - Using the updated handle may well invalidate the returned references.
     pub unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
-        let (next_edge, kv) = next_unchecked_mut(ptr::read(self));
-        *self = next_edge;
+        let kv = replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_kv();
+            let kv = unwrap_unchecked(kv.ok());
+            (ptr::read(&kv).next_leaf_edge(), kv)
+        });
         // Doing the descend (and perhaps another move) invalidates the references
         // returned by `into_kv_mut`, so we have to do this last.
         kv.into_kv_mut()
@@ -171,8 +135,11 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
     /// - The caller must ensure that the leaf edge is not the first one in the tree.
     /// - Using the updated handle may well invalidate the returned references.
     pub unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
-        let (next_edge, kv) = next_back_unchecked_mut(ptr::read(self));
-        *self = next_edge;
+        let kv = replace(self, |leaf_edge| {
+            let kv = leaf_edge.next_back_kv();
+            let kv = unwrap_unchecked(kv.ok());
+            (ptr::read(&kv).next_back_leaf_edge(), kv)
+        });
         // Doing the descend (and perhaps another move) invalidates the references
         // returned by `into_kv_mut`, so we have to do this last.
         kv.into_kv_mut()
@@ -192,9 +159,12 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
     ///   if the two preconditions above hold.
     /// - Using the updated handle may well invalidate the returned references.
     pub unsafe fn next_unchecked(&mut self) -> (K, V) {
-        let (next_edge, k, v) = next_unchecked_deallocating(ptr::read(self));
-        *self = next_edge;
-        (k, v)
+        replace(self, |leaf_edge| {
+            let kv = next_kv_unchecked_dealloc(leaf_edge);
+            let k = ptr::read(kv.reborrow().into_kv().0);
+            let v = ptr::read(kv.reborrow().into_kv().1);
+            (kv.next_leaf_edge(), (k, v))
+        })
     }
 
     /// Moves the leaf edge handle to the previous leaf edge and returns the key
@@ -209,9 +179,12 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
     ///   if the two preconditions above hold.
     /// - Using the updated handle may well invalidate the returned references.
     pub unsafe fn next_back_unchecked(&mut self) -> (K, V) {
-        let (next_edge, k, v) = next_back_unchecked_deallocating(ptr::read(self));
-        *self = next_edge;
-        (k, v)
+        replace(self, |leaf_edge| {
+            let kv = next_back_kv_unchecked_dealloc(leaf_edge);
+            let k = ptr::read(kv.reborrow().into_kv().0);
+            let v = ptr::read(kv.reborrow().into_kv().1);
+            (kv.next_back_leaf_edge(), (k, v))
+        })
     }
 }
 
@@ -242,3 +215,29 @@ impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
         }
     }
 }
+
+impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
+    /// Returns the leaf edge closest to a KV for forward navigation.
+    pub fn next_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
+        match self.force() {
+            Leaf(leaf_kv) => leaf_kv.right_edge(),
+            Internal(internal_kv) => {
+                let next_internal_edge = internal_kv.right_edge();
+                next_internal_edge.descend().first_leaf_edge()
+            }
+        }
+    }
+
+    /// Returns the leaf edge closest to a KV for backward navigation.
+    pub fn next_back_leaf_edge(
+        self,
+    ) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
+        match self.force() {
+            Leaf(leaf_kv) => leaf_kv.left_edge(),
+            Internal(internal_kv) => {
+                let next_internal_edge = internal_kv.left_edge();
+                next_internal_edge.descend().last_leaf_edge()
+            }
+        }
+    }
+}

src/liballoc/collections/btree/node.rs

@@ -451,31 +451,25 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
     }
 }
 
-impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
+impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
     /// Similar to `ascend`, gets a reference to a node's parent node, but also
     /// deallocate the current node in the process. This is unsafe because the
     /// current node will still be accessible despite being deallocated.
     pub unsafe fn deallocate_and_ascend(
         self,
     ) -> Option<Handle<NodeRef<marker::Owned, K, V, marker::Internal>, marker::Edge>> {
         assert!(!self.is_shared_root());
+        let height = self.height;
         let node = self.node;
         let ret = self.ascend().ok();
-        Global.dealloc(node.cast(), Layout::new::<LeafNode<K, V>>());
-        ret
-    }
-}
-
-impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
-    /// Similar to `ascend`, gets a reference to a node's parent node, but also
-    /// deallocate the current node in the process. This is unsafe because the
-    /// current node will still be accessible despite being deallocated.
-    pub unsafe fn deallocate_and_ascend(
-        self,
-    ) -> Option<Handle<NodeRef<marker::Owned, K, V, marker::Internal>, marker::Edge>> {
-        let node = self.node;
-        let ret = self.ascend().ok();
-        Global.dealloc(node.cast(), Layout::new::<InternalNode<K, V>>());
+        Global.dealloc(
+            node.cast(),
+            if height > 0 {
+                Layout::new::<InternalNode<K, V>>()
+            } else {
+                Layout::new::<LeafNode<K, V>>()
+            },
+        );
         ret
     }
 }
@@ -1418,15 +1412,23 @@ unsafe fn move_edges<K, V>(
     dest.correct_childrens_parent_links(dest_offset, dest_offset + count);
 }
 
-impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::KV> {
+impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
     pub fn forget_node_type(
         self,
-    ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
-        unsafe { Handle::new_kv(self.node.forget_type(), self.idx) }
+    ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
+        unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
+    }
+}
+
+impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
+    pub fn forget_node_type(
+        self,
+    ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
+        unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
     }
 }
 
-impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV> {
+impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::KV> {
     pub fn forget_node_type(
         self,
     ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {

src/liballoc/tests/btree/map.rs

@@ -23,6 +23,11 @@ fn test_basic_large() {
         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));
     }

src/liballoc/tests/btree/set.rs

@@ -487,21 +487,26 @@ fn test_first_last() {
     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> {