mikros/rust
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

BTreeMap: Add Cursor and CursorMut

Browse Source
This commit is contained in:
Amanieu d'Antras 2022-12-08 18:32:43 +00:00
parent eb70c82d29
commit 36831b3ccb
4 changed files with 839 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

730
library/alloc/src/collections/btree/map.rs
View File

@ -6,7 +6,7 @@ use core::hash::{Hash, Hasher};
use core::iter::{FromIterator, FusedIterator};
use core::marker::PhantomData;
use core::mem::{self, ManuallyDrop};
use core::ops::{Index, RangeBounds};
use core::ops::{Bound, Index, RangeBounds};
use core::ptr;
use crate::alloc::{Allocator, Global};
@ -15,7 +15,7 @@ use super::borrow::DormantMutRef;
use super::dedup_sorted_iter::DedupSortedIter;
use super::navigate::{LazyLeafRange, LeafRange};
use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};
use super::search::SearchResult::*;
use super::search::{SearchBound, SearchResult::*};
use super::set_val::SetValZST;
mod entry;
@ -2422,6 +2422,732 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
pub const fn is_empty(&self) -> bool {
self.len() == 0
}
/// Returns a [`Cursor`] pointing at the first element that is above the
/// given bound.
///
/// If no such element exists then a cursor pointing at the "ghost"
/// non-element is returned.
///
/// Passing [`Bound::Unbounded`] will return a cursor pointing at the first
/// element of the map.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// #![feature(btree_cursors)]
///
/// use std::collections::BTreeMap;
/// use std::ops::Bound;
///
/// let mut a = BTreeMap::new();
/// a.insert(1, "a");
/// a.insert(2, "b");
/// a.insert(3, "c");
/// a.insert(4, "c");
/// let cursor = a.lower_bound(Bound::Excluded(&2));
/// assert_eq!(cursor.key(), Some(&3));
/// ```
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn lower_bound<Q>(&self, bound: Bound<&Q>) -> Cursor<'_, K, V>
where
K: Borrow<Q> + Ord,
Q: Ord,
{
let root_node = match self.root.as_ref() {
None => return Cursor { current: None, root: None },
Some(root) => root.reborrow(),
};
let edge = root_node.lower_bound(SearchBound::from_range(bound));
Cursor { current: edge.next_kv().ok(), root: self.root.as_ref() }
}
/// Returns a [`CursorMut`] pointing at the first element that is above the
/// given bound.
///
/// If no such element exists then a cursor pointing at the "ghost"
/// non-element is returned.
///
/// Passing [`Bound::Unbounded`] will return a cursor pointing at the first
/// element of the map.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// #![feature(btree_cursors)]
///
/// use std::collections::BTreeMap;
/// use std::ops::Bound;
///
/// let mut a = BTreeMap::new();
/// a.insert(1, "a");
/// a.insert(2, "b");
/// a.insert(3, "c");
/// a.insert(4, "c");
/// let cursor = a.lower_bound_mut(Bound::Excluded(&2));
/// assert_eq!(cursor.key(), Some(&3));
/// ```
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn lower_bound_mut<Q>(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V, A>
where
K: Borrow<Q> + Ord,
Q: Ord,
{
let (root, dormant_root) = DormantMutRef::new(&mut self.root);
let root_node = match root.as_mut() {
None => {
return CursorMut {
current: None,
root: dormant_root,
length: &mut self.length,
alloc: &mut *self.alloc,
};
}
Some(root) => root.borrow_mut(),
};
let edge = root_node.lower_bound(SearchBound::from_range(bound));
CursorMut {
current: edge.next_kv().ok(),
root: dormant_root,
length: &mut self.length,
alloc: &mut *self.alloc,
}
}
/// Returns a [`Cursor`] pointing at the last element that is below the
/// given bound.
///
/// If no such element exists then a cursor pointing at the "ghost"
/// non-element is returned.
///
/// Passing [`Bound::Unbounded`] will return a cursor pointing at the last
/// element of the map.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// #![feature(btree_cursors)]
///
/// use std::collections::BTreeMap;
/// use std::ops::Bound;
///
/// let mut a = BTreeMap::new();
/// a.insert(1, "a");
/// a.insert(2, "b");
/// a.insert(3, "c");
/// a.insert(4, "c");
/// let cursor = a.upper_bound(Bound::Excluded(&3));
/// assert_eq!(cursor.key(), Some(&2));
/// ```
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn upper_bound<Q>(&self, bound: Bound<&Q>) -> Cursor<'_, K, V>
where
K: Borrow<Q> + Ord,
Q: Ord,
{
let root_node = match self.root.as_ref() {
None => return Cursor { current: None, root: None },
Some(root) => root.reborrow(),
};
let edge = root_node.upper_bound(SearchBound::from_range(bound));
Cursor { current: edge.next_back_kv().ok(), root: self.root.as_ref() }
}
/// Returns a [`CursorMut`] pointing at the last element that is below the
/// given bound.
///
/// If no such element exists then a cursor pointing at the "ghost"
/// non-element is returned.
///
/// Passing [`Bound::Unbounded`] will return a cursor pointing at the last
/// element of the map.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// #![feature(btree_cursors)]
///
/// use std::collections::BTreeMap;
/// use std::ops::Bound;
///
/// let mut a = BTreeMap::new();
/// a.insert(1, "a");
/// a.insert(2, "b");
/// a.insert(3, "c");
/// a.insert(4, "c");
/// let cursor = a.upper_bound_mut(Bound::Excluded(&3));
/// assert_eq!(cursor.key(), Some(&2));
/// ```
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn upper_bound_mut<Q>(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V, A>
where
K: Borrow<Q> + Ord,
Q: Ord,
{
let (root, dormant_root) = DormantMutRef::new(&mut self.root);
let root_node = match root.as_mut() {
None => {
return CursorMut {
current: None,
root: dormant_root,
length: &mut self.length,
alloc: &mut *self.alloc,
};
}
Some(root) => root.borrow_mut(),
};
let edge = root_node.upper_bound(SearchBound::from_range(bound));
CursorMut {
current: edge.next_back_kv().ok(),
root: dormant_root,
length: &mut self.length,
alloc: &mut *self.alloc,
}
}
}
/// A cursor over a `BTreeMap`.
///
/// A `Cursor` is like an iterator, except that it can freely seek back-and-forth.
///
/// Cursors always point to an element in the tree, and index in a logically circular way.
/// To accommodate this, there is a "ghost" non-element that yields `None` between the last and
/// first elements of the tree.
///
/// A `Cursor` is created with the [`BTreeMap::lower_bound`] and [`BTreeMap::upper_bound`] methods.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub struct Cursor<'a, K: 'a, V: 'a> {
current: Option<Handle<NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>, marker::KV>>,
root: Option<&'a node::Root<K, V>>,
}
#[unstable(feature = "btree_cursors", issue = "107540")]
impl<K, V> Clone for Cursor<'_, K, V> {
fn clone(&self) -> Self {
let Cursor { current, root } = *self;
Cursor { current, root }
}
}
#[unstable(feature = "btree_cursors", issue = "107540")]
impl<K: Debug, V: Debug> Debug for Cursor<'_, K, V> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("Cursor").field(&self.key_value()).finish()
}
}
/// A cursor over a `BTreeMap` with editing operations.
///
/// A `Cursor` is like an iterator, except that it can freely seek back-and-forth, and can
/// safely mutate the tree during iteration. This is because the lifetime of its yielded
/// references is tied to its own lifetime, instead of just the underlying tree. This means
/// cursors cannot yield multiple elements at once.
///
/// Cursors always point to an element in the tree, and index in a logically circular way.
/// To accommodate this, there is a "ghost" non-element that yields `None` between the last and
/// first elements of the tree.
///
/// A `Cursor` is created with the [`BTreeMap::lower_bound_mut`] and [`BTreeMap::upper_bound_mut`]
/// methods.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub struct CursorMut<
'a,
K: 'a,
V: 'a,
#[unstable(feature = "allocator_api", issue = "32838")] A = Global,
> {
current: Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>>,
root: DormantMutRef<'a, Option<node::Root<K, V>>>,
length: &'a mut usize,
alloc: &'a mut A,
}
#[unstable(feature = "btree_cursors", issue = "107540")]
impl<K: Debug, V: Debug, A> Debug for CursorMut<'_, K, V, A> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("CursorMut").field(&self.key_value()).finish()
}
}
impl<'a, K, V> Cursor<'a, K, V> {
/// Moves the cursor to the next element of the `BTreeMap`.
///
/// If the cursor is pointing to the "ghost" non-element then this will move it to
/// the first element of the `BTreeMap`. If it is pointing to the last
/// element of the `BTreeMap` then this will move it to the "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn move_next(&mut self) {
match self.current.take() {
None => {
self.current = self.root.and_then(|root| {
root.reborrow().first_leaf_edge().forget_node_type().right_kv().ok()
});
}
Some(current) => {
self.current = current.next_leaf_edge().next_kv().ok();
}
}
}
/// Moves the cursor to the previous element of the `BTreeMap`.
///
/// If the cursor is pointing to the "ghost" non-element then this will move it to
/// the last element of the `BTreeMap`. If it is pointing to the first
/// element of the `BTreeMap` then this will move it to the "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn move_prev(&mut self) {
match self.current.take() {
None => {
self.current = self.root.and_then(|root| {
root.reborrow().last_leaf_edge().forget_node_type().left_kv().ok()
});
}
Some(current) => {
self.current = current.next_back_leaf_edge().next_back_kv().ok();
}
}
}
/// Returns a reference to the key of the element that the cursor is
/// currently pointing to.
///
/// This returns `None` if the cursor is currently pointing to the
/// "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn key(&self) -> Option<&'a K> {
self.current.as_ref().map(|current| current.into_kv().0)
}
/// Returns a reference to the value of the element that the cursor is
/// currently pointing to.
///
/// This returns `None` if the cursor is currently pointing to the
/// "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn value(&self) -> Option<&'a V> {
self.current.as_ref().map(|current| current.into_kv().1)
}
/// Returns a reference to the key and value of the element that the cursor
/// is currently pointing to.
///
/// This returns `None` if the cursor is currently pointing to the
/// "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn key_value(&self) -> Option<(&'a K, &'a V)> {
self.current.as_ref().map(|current| current.into_kv())
}
/// Returns a reference to the next element.
///
/// If the cursor is pointing to the "ghost" non-element then this returns
/// the first element of the `BTreeMap`. If it is pointing to the last
/// element of the `BTreeMap` then this returns `None`.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn peek_next(&self) -> Option<(&'a K, &'a V)> {
let mut next = self.clone();
next.move_next();
next.current.as_ref().map(|current| current.into_kv())
}
/// Returns a reference to the previous element.
///
/// If the cursor is pointing to the "ghost" non-element then this returns
/// the last element of the `BTreeMap`. If it is pointing to the first
/// element of the `BTreeMap` then this returns `None`.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn peek_prev(&self) -> Option<(&'a K, &'a V)> {
let mut prev = self.clone();
prev.move_prev();
prev.current.as_ref().map(|current| current.into_kv())
}
}
impl<'a, K, V, A> CursorMut<'a, K, V, A> {
/// Moves the cursor to the next element of the `BTreeMap`.
///
/// If the cursor is pointing to the "ghost" non-element then this will move it to
/// the first element of the `BTreeMap`. If it is pointing to the last
/// element of the `BTreeMap` then this will move it to the "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn move_next(&mut self) {
match self.current.take() {
None => {
// SAFETY: The previous borrow of root has ended.
self.current = unsafe { self.root.reborrow() }.as_mut().and_then(|root| {
root.borrow_mut().first_leaf_edge().forget_node_type().right_kv().ok()
});
}
Some(current) => {
self.current = current.next_leaf_edge().next_kv().ok();
}
}
}
/// Moves the cursor to the previous element of the `BTreeMap`.
///
/// If the cursor is pointing to the "ghost" non-element then this will move it to
/// the last element of the `BTreeMap`. If it is pointing to the first
/// element of the `BTreeMap` then this will move it to the "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn move_prev(&mut self) {
match self.current.take() {
None => {
// SAFETY: The previous borrow of root has ended.
self.current = unsafe { self.root.reborrow() }.as_mut().and_then(|root| {
root.borrow_mut().last_leaf_edge().forget_node_type().left_kv().ok()
});
}
Some(current) => {
self.current = current.next_back_leaf_edge().next_back_kv().ok();
}
}
}
/// Returns a reference to the key of the element that the cursor is
/// currently pointing to.
///
/// This returns `None` if the cursor is currently pointing to the
/// "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn key(&self) -> Option<&K> {
self.current.as_ref().map(|current| current.reborrow().into_kv().0)
}
/// Returns a reference to the value of the element that the cursor is
/// currently pointing to.
///
/// This returns `None` if the cursor is currently pointing to the
/// "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn value(&self) -> Option<&V> {
self.current.as_ref().map(|current| current.reborrow().into_kv().1)
}
/// Returns a reference to the key and value of the element that the cursor
/// is currently pointing to.
///
/// This returns `None` if the cursor is currently pointing to the
/// "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn key_value(&self) -> Option<(&K, &V)> {
self.current.as_ref().map(|current| current.reborrow().into_kv())
}
/// Returns a mutable reference to the value of the element that the cursor
/// is currently pointing to.
///
/// This returns `None` if the cursor is currently pointing to the
/// "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn value_mut(&mut self) -> Option<&mut V> {
self.current.as_mut().map(|current| current.kv_mut().1)
}
/// Returns a reference to the key and mutable reference to the value of the
/// element that the cursor is currently pointing to.
///
/// This returns `None` if the cursor is currently pointing to the
/// "ghost" non-element.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn key_value_mut(&mut self) -> Option<(&K, &mut V)> {
self.current.as_mut().map(|current| {
let (k, v) = current.kv_mut();
(&*k, v)
})
}
/// Returns a mutable reference to the of the element that the cursor is
/// currently pointing to.
///
/// This returns `None` if the cursor is currently pointing to the
/// "ghost" non-element.
///
/// # Safety
///
/// This can be used to modify the key, but you must ensure that the
/// `BTreeMap` invariants are maintained. Specifically:
///
/// * The key must remain unique within the tree.
/// * The key must remain in sorted order with regards to other elements in
/// the tree.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub unsafe fn key_mut_unchecked(&mut self) -> Option<&mut K> {
self.current.as_mut().map(|current| current.kv_mut().0)
}
/// Returns a reference to the key and value of the next element.
///
/// If the cursor is pointing to the "ghost" non-element then this returns
/// the first element of the `BTreeMap`. If it is pointing to the last
/// element of the `BTreeMap` then this returns `None`.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn peek_next(&mut self) -> Option<(&K, &mut V)> {
let (k, v) = match self.current {
None => {
// SAFETY: The previous borrow of root has ended.
unsafe { self.root.reborrow() }
.as_mut()?
.borrow_mut()
.first_leaf_edge()
.next_kv()
.ok()?
.into_kv_valmut()
}
// SAFETY: We're not using this to mutate the tree.
Some(ref mut current) => {
unsafe { current.reborrow_mut() }.next_leaf_edge().next_kv().ok()?.into_kv_valmut()
}
};
Some((k, v))
}
/// Returns a reference to the key and value of the previous element.
///
/// If the cursor is pointing to the "ghost" non-element then this returns
/// the last element of the `BTreeMap`. If it is pointing to the first
/// element of the `BTreeMap` then this returns `None`.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn peek_prev(&mut self) -> Option<(&K, &mut V)> {
let (k, v) = match self.current.as_mut() {
None => {
// SAFETY: The previous borrow of root has ended.
unsafe { self.root.reborrow() }
.as_mut()?
.borrow_mut()
.first_leaf_edge()
.next_kv()
.ok()?
.into_kv_valmut()
}
Some(current) => {
// SAFETY: We're not using this to mutate the tree.
unsafe { current.reborrow_mut() }
.next_back_leaf_edge()
.next_back_kv()
.ok()?
.into_kv_valmut()
}
};
Some((k, v))
}
/// Returns a read-only cursor pointing to the current element.
///
/// The lifetime of the returned `Cursor` is bound to that of the
/// `CursorMut`, which means it cannot outlive the `CursorMut` and that the
/// `CursorMut` is frozen for the lifetime of the `Cursor`.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn as_cursor(&self) -> Cursor<'_, K, V> {
Cursor {
// SAFETY: The tree is immutable while the cursor exists.
root: unsafe { self.root.reborrow_shared().as_ref() },
current: self.current.as_ref().map(|current| current.reborrow()),
}
}
}
// Now the tree editing operations
impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> {
/// Inserts a new element into the `BTreeMap` after the current one.
///
/// If the cursor is pointing at the "ghost" non-element then the new element is
/// inserted at the front of the `BTreeMap`.
///
/// # Safety
///
/// You must ensure that the `BTreeMap` invariants are maintained.
/// Specifically:
///
/// * The key of the newly inserted element must be unique in the tree.
/// * All keys in the tree must remain in sorted order.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub unsafe fn insert_after_unchecked(&mut self, key: K, value: V) {
let edge = match self.current.take() {
None => {
// SAFETY: We have no other reference to the tree.
match unsafe { self.root.reborrow() } {
root @ None => {
// Tree is empty, allocate a new root.
let mut node = NodeRef::new_leaf(self.alloc.clone());
node.borrow_mut().push(key, value);
*root = Some(node.forget_type());
*self.length += 1;
return;
}
Some(root) => root.borrow_mut().first_leaf_edge(),
}
}
Some(current) => current.next_leaf_edge(),
};
let handle = edge.insert_recursing(key, value, self.alloc.clone(), |ins| {
drop(ins.left);
// SAFETY: The handle to the newly inserted value is always on a
// leaf node, so adding a new root node doesn't invalidate it.
let root = unsafe { self.root.reborrow().as_mut().unwrap() };
root.push_internal_level(self.alloc.clone()).push(ins.kv.0, ins.kv.1, ins.right)
});
self.current = handle.left_edge().next_back_kv().ok();
*self.length += 1;
}
/// Inserts a new element into the `BTreeMap` before the current one.
///
/// If the cursor is pointing at the "ghost" non-element then the new element is
/// inserted at the end of the `BTreeMap`.
///
/// # Safety
///
/// You must ensure that the `BTreeMap` invariants are maintained.
/// Specifically:
///
/// * The key of the newly inserted element must be unique in the tree.
/// * All keys in the tree must remain in sorted order.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub unsafe fn insert_before_unchecked(&mut self, key: K, value: V) {
let edge = match self.current.take() {
None => {
// SAFETY: We have no other reference to the tree.
match unsafe { self.root.reborrow() } {
root @ None => {
// Tree is empty, allocate a new root.
let mut node = NodeRef::new_leaf(self.alloc.clone());
node.borrow_mut().push(key, value);
*root = Some(node.forget_type());
*self.length += 1;
return;
}
Some(root) => root.borrow_mut().last_leaf_edge(),
}
}
Some(current) => current.next_back_leaf_edge(),
};
let handle = edge.insert_recursing(key, value, self.alloc.clone(), |ins| {
drop(ins.left);
// SAFETY: The handle to the newly inserted value is always on a
// leaf node, so adding a new root node doesn't invalidate it.
let root = unsafe { self.root.reborrow().as_mut().unwrap() };
root.push_internal_level(self.alloc.clone()).push(ins.kv.0, ins.kv.1, ins.right)
});
self.current = handle.right_edge().next_kv().ok();
*self.length += 1;
}
/// Inserts a new element into the `BTreeMap` after the current one.
///
/// If the cursor is pointing at the "ghost" non-element then the new element is
/// inserted at the front of the `BTreeMap`.
///
/// # Panics
///
/// This function panics if:
/// - the given key compares less than or equal to the current element (if
/// any).
/// - the given key compares greater than or equal to the next element (if
/// any).
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn insert_after(&mut self, key: K, value: V) {
if let Some(current) = self.key() {
if &key <= current {
panic!("key must be ordered above the current element");
}
}
if let Some((next, _)) = self.peek_prev() {
if &key >= next {
panic!("key must be ordered below the next element");
}
}
unsafe {
self.insert_after_unchecked(key, value);
}
}
/// Inserts a new element into the `BTreeMap` before the current one.
///
/// If the cursor is pointing at the "ghost" non-element then the new element is
/// inserted at the end of the `BTreeMap`.
///
/// # Panics
///
/// This function panics if:
/// - the given key compares greater than or equal to the current element
/// (if any).
/// - the given key compares less than or equal to the previous element (if
/// any).
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn insert_before(&mut self, key: K, value: V) {
if let Some(current) = self.key() {
if &key >= current {
panic!("key must be ordered below the current element");
}
}
if let Some((prev, _)) = self.peek_prev() {
if &key <= prev {
panic!("key must be ordered above the previous element");
}
}
unsafe {
self.insert_before_unchecked(key, value);
}
}
/// Removes the current element from the `BTreeMap`.
///
/// The element that was removed is returned, and the cursor is
/// moved to point to the next element in the `BTreeMap`.
///
/// If the cursor is currently pointing to the "ghost" non-element then no element
/// is removed and `None` is returned. The cursor is not moved in this case.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn remove_current(&mut self) -> Option<(K, V)> {
let current = self.current.take()?;
let mut emptied_internal_root = false;
let (kv, pos) =
current.remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone());
self.current = pos.next_kv().ok();
*self.length -= 1;
if emptied_internal_root {
// SAFETY: This is safe since current does not point within the now
// empty root node.
let root = unsafe { self.root.reborrow().as_mut().unwrap() };
root.pop_internal_level(self.alloc.clone());
}
Some(kv)
}
/// Removes the current element from the `BTreeMap`.
///
/// The element that was removed is returned, and the cursor is
/// moved to point to the previous element in the `BTreeMap`.
///
/// If the cursor is currently pointing to the "ghost" non-element then no element
/// is removed and `None` is returned. The cursor is not moved in this case.
#[unstable(feature = "btree_cursors", issue = "107540")]
pub fn remove_current_and_move_back(&mut self) -> Option<(K, V)> {
let current = self.current.take()?;
let mut emptied_internal_root = false;
let (kv, pos) =
current.remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone());
self.current = pos.next_back_kv().ok();
*self.length -= 1;
if emptied_internal_root {
// SAFETY: This is safe since current does not point within the now
// empty root node.
let root = unsafe { self.root.reborrow().as_mut().unwrap() };
root.pop_internal_level(self.alloc.clone());
}
Some(kv)
}
}
#[cfg(test)]

49
library/alloc/src/collections/btree/map/tests.rs
View File

@ -2336,3 +2336,52 @@ fn from_array() {
let unordered_duplicates = BTreeMap::from([(3, 4), (1, 2), (1, 2)]);
assert_eq!(map, unordered_duplicates);
}
#[test]
fn test_cursor() {
let map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]);
let mut cur = map.lower_bound(Bound::Unbounded);
assert_eq!(cur.key(), Some(&1));
cur.move_next();
assert_eq!(cur.key(), Some(&2));
assert_eq!(cur.peek_next(), Some((&3, &'c')));
cur.move_prev();
assert_eq!(cur.key(), Some(&1));
assert_eq!(cur.peek_prev(), None);
let mut cur = map.upper_bound(Bound::Excluded(&1));
assert_eq!(cur.key(), None);
cur.move_next();
assert_eq!(cur.key(), Some(&1));
cur.move_prev();
assert_eq!(cur.key(), None);
assert_eq!(cur.peek_prev(), Some((&3, &'c')));
}
#[test]
fn test_cursor_mut() {
let mut map = BTreeMap::from([(1, 'a'), (3, 'c'), (5, 'e')]);
let mut cur = map.lower_bound_mut(Bound::Excluded(&3));
assert_eq!(cur.key(), Some(&5));
cur.insert_before(4, 'd');
assert_eq!(cur.key(), Some(&5));
assert_eq!(cur.peek_prev(), Some((&4, &mut 'd')));
cur.move_next();
assert_eq!(cur.key(), None);
cur.insert_before(6, 'f');
assert_eq!(cur.key(), None);
assert_eq!(cur.remove_current(), None);
assert_eq!(cur.key(), None);
cur.insert_after(0, '?');
assert_eq!(cur.key(), None);
assert_eq!(map, BTreeMap::from([(0, '?'), (1, 'a'), (3, 'c'), (4, 'd'), (5, 'e'), (6, 'f')]));
let mut cur = map.upper_bound_mut(Bound::Included(&5));
assert_eq!(cur.key(), Some(&5));
assert_eq!(cur.remove_current(), Some((5, 'e')));
assert_eq!(cur.key(), Some(&6));
assert_eq!(cur.remove_current_and_move_back(), Some((6, 'f')));
assert_eq!(cur.key(), Some(&4));
assert_eq!(map, BTreeMap::from([(0, '?'), (1, 'a'), (3, 'c'), (4, 'd')]));
}

55
library/alloc/src/collections/btree/navigate.rs
View File

@ -4,6 +4,7 @@ use core::ops::RangeBounds;
use core::ptr;
use super::node::{marker, ForceResult::*, Handle, NodeRef};
use super::search::SearchBound;
use crate::alloc::Allocator;
// `front` and `back` are always both `None` or both `Some`.
@ -386,7 +387,7 @@ impl<BorrowType: marker::BorrowType, K, V>
/// 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.
fn next_back_kv(
pub fn next_back_kv(
self,
) -> Result<
Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
@ -707,7 +708,9 @@ impl<BorrowType: marker::BorrowType, K, V>
}
/// Returns the leaf edge closest to a KV for backward navigation.
fn next_back_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
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) => {
@ -717,3 +720,51 @@ impl<BorrowType: marker::BorrowType, K, V>
}
}
}
impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
/// Returns the leaf edge corresponding to the first point at which the
/// given bound is true.
pub fn lower_bound<Q: ?Sized>(
self,
mut bound: SearchBound<&Q>,
) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>
where
Q: Ord,
K: Borrow<Q>,
{
let mut node = self;
loop {
let (edge, new_bound) = node.find_lower_bound_edge(bound);
match edge.force() {
Leaf(edge) => return edge,
Internal(edge) => {
node = edge.descend();
bound = new_bound;
}
}
}
}
/// Returns the leaf edge corresponding to the last point at which the
/// given bound is true.
pub fn upper_bound<Q: ?Sized>(
self,
mut bound: SearchBound<&Q>,
) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>
where
Q: Ord,
K: Borrow<Q>,
{
let mut node = self;
loop {
let (edge, new_bound) = node.find_upper_bound_edge(bound);
match edge.force() {
Leaf(edge) => return edge,
Internal(edge) => {
node = edge.descend();
bound = new_bound;
}
}
}
}
}

10
library/alloc/src/collections/btree/node.rs
View File

@ -819,7 +819,7 @@ impl<'a, K, V, NodeType, HandleType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeT
/// Returns a dormant copy of this handle which can be reawakened later.
///
/// See [`DormantMutRef`] for more details.
/// See `DormantMutRef` for more details.
pub fn dormant(&self) -> Handle<NodeRef<marker::DormantMut, K, V, NodeType>, HandleType> {
Handle { node: self.node.dormant(), idx: self.idx, _marker: PhantomData }
}
@ -1099,6 +1099,14 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>
let leaf = self.node.into_leaf_mut();
unsafe { leaf.vals.get_unchecked_mut(self.idx).assume_init_mut() }
}
pub fn into_kv_valmut(self) -> (&'a K, &'a mut V) {
debug_assert!(self.idx < self.node.len());
let leaf = self.node.into_leaf_mut();
let k = unsafe { leaf.keys.get_unchecked(self.idx).assume_init_ref() };
let v = unsafe { leaf.vals.get_unchecked_mut(self.idx).assume_init_mut() };
(k, v)
}
}
impl<'a, K, V, NodeType> Handle<NodeRef<marker::ValMut<'a>, K, V, NodeType>, marker::KV> {