Rollup merge of #105641 - Amanieu:btree_cursor, r=m-ou-se

Implement cursors for BTreeMap

See the ACP for an overview of the API: https://github.com/rust-lang/libs-team/issues/141

The implementation is split into 2 commits:
- The first changes the internal insertion functions to return a handle to the newly inserted element. The lifetimes involved are a bit hairy since we need a mutable handle to both the `BTreeMap` itself (which holds the root) and the nodes allocated in memory. I have tested that this passes the standard library testsuite under miri.
- The second commit implements the cursor API itself. This is more straightforward to follow but still involves some unsafe code to deal with simultaneous mutable borrows of the tree root and the node that is currently being iterated.
This commit is contained in:
Matthias Krüger 2023-02-08 18:32:41 +01:00 committed by GitHub
commit 562581c2db
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 957 additions and 41 deletions

View File

@ -41,6 +41,28 @@ impl<'a, T> DormantMutRef<'a, T> {
// SAFETY: our own safety conditions imply this reference is again unique.
unsafe { &mut *self.ptr.as_ptr() }
}
/// Borrows a new mutable reference from the unique borrow initially captured.
///
/// # Safety
///
/// The reborrow must have ended, i.e., the reference returned by `new` and
/// all pointers and references derived from it, must not be used anymore.
pub unsafe fn reborrow(&mut self) -> &'a mut T {
// SAFETY: our own safety conditions imply this reference is again unique.
unsafe { &mut *self.ptr.as_ptr() }
}
/// Borrows a new shared reference from the unique borrow initially captured.
///
/// # Safety
///
/// The reborrow must have ended, i.e., the reference returned by `new` and
/// all pointers and references derived from it, must not be used anymore.
pub unsafe fn reborrow_shared(&self) -> &'a T {
// SAFETY: our own safety conditions imply this reference is again unique.
unsafe { &*self.ptr.as_ptr() }
}
}
#[cfg(test)]

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)]

View File

@ -347,7 +347,7 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
/// assert_eq!(map["poneyland"], 37);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn insert(self, value: V) -> &'a mut V {
pub fn insert(mut self, value: V) -> &'a mut V {
let out_ptr = match self.handle {
None => {
// SAFETY: There is no tree yet so no reference to it exists.
@ -358,25 +358,27 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
map.length = 1;
val_ptr
}
Some(handle) => match handle.insert_recursing(self.key, value, self.alloc.clone()) {
(None, val_ptr) => {
// SAFETY: We have consumed self.handle.
let map = unsafe { self.dormant_map.awaken() };
map.length += 1;
val_ptr
}
(Some(ins), val_ptr) => {
drop(ins.left);
// SAFETY: We have consumed self.handle and dropped the
// remaining reference to the tree, ins.left.
let map = unsafe { self.dormant_map.awaken() };
let root = map.root.as_mut().unwrap(); // same as ins.left
root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right);
map.length += 1;
val_ptr
}
},
Some(handle) => {
let new_handle =
handle.insert_recursing(self.key, value, self.alloc.clone(), |ins| {
drop(ins.left);
// SAFETY: Pushing a new root node doesn't invalidate
// handles to existing nodes.
let map = unsafe { self.dormant_map.reborrow() };
let root = map.root.as_mut().unwrap(); // same as ins.left
root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right)
});
// Get the pointer to the value
let val_ptr = new_handle.into_val_mut();
// SAFETY: We have consumed self.handle.
let map = unsafe { self.dormant_map.awaken() };
map.length += 1;
val_ptr
}
};
// Now that we have finished growing the tree using borrowed references,
// dereference the pointer to a part of it, that we picked up along the way.
unsafe { &mut *out_ptr }

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')]));
}

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

View File

@ -442,6 +442,24 @@ impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
// SAFETY: we have exclusive access to the entire node.
unsafe { &mut *ptr }
}
/// Returns a dormant copy of this node with its lifetime erased which can
/// be reawakened later.
pub fn dormant(&self) -> NodeRef<marker::DormantMut, K, V, Type> {
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
}
}
impl<K, V, Type> NodeRef<marker::DormantMut, K, V, Type> {
/// Revert to the unique borrow initially captured.
///
/// # Safety
///
/// The reborrow must have ended, i.e., the reference returned by `new` and
/// all pointers and references derived from it, must not be used anymore.
pub unsafe fn awaken<'a>(self) -> NodeRef<marker::Mut<'a>, K, V, Type> {
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
}
}
impl<K, V, Type> NodeRef<marker::Dying, K, V, Type> {
@ -798,6 +816,25 @@ impl<'a, K, V, NodeType, HandleType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeT
// We can't use Handle::new_kv or Handle::new_edge because we don't know our type
Handle { node: unsafe { self.node.reborrow_mut() }, idx: self.idx, _marker: PhantomData }
}
/// Returns a dormant copy of this handle which can be reawakened later.
///
/// 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 }
}
}
impl<K, V, NodeType, HandleType> Handle<NodeRef<marker::DormantMut, K, V, NodeType>, HandleType> {
/// Revert to the unique borrow initially captured.
///
/// # Safety
///
/// The reborrow must have ended, i.e., the reference returned by `new` and
/// all pointers and references derived from it, must not be used anymore.
pub unsafe fn awaken<'a>(self) -> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, HandleType> {
Handle { node: unsafe { self.node.awaken() }, idx: self.idx, _marker: PhantomData }
}
}
impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
@ -851,9 +888,11 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
/// Inserts a new key-value pair between the key-value pairs to the right and left of
/// this edge. This method assumes that there is enough space in the node for the new
/// pair to fit.
///
/// The returned pointer points to the inserted value.
fn insert_fit(&mut self, key: K, val: V) -> *mut V {
unsafe fn insert_fit(
mut self,
key: K,
val: V,
) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
debug_assert!(self.node.len() < CAPACITY);
let new_len = self.node.len() + 1;
@ -862,7 +901,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
slice_insert(self.node.val_area_mut(..new_len), self.idx, val);
*self.node.len_mut() = new_len as u16;
self.node.val_area_mut(self.idx).assume_init_mut()
Handle::new_kv(self.node, self.idx)
}
}
}
@ -871,21 +910,26 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
/// Inserts a new key-value pair between the key-value pairs to the right and left of
/// this edge. This method splits the node if there isn't enough room.
///
/// The returned pointer points to the inserted value.
/// Returns a dormant handle to the inserted node which can be reawakened
/// once splitting is complete.
fn insert<A: Allocator + Clone>(
mut self,
self,
key: K,
val: V,
alloc: A,
) -> (Option<SplitResult<'a, K, V, marker::Leaf>>, *mut V) {
) -> (
Option<SplitResult<'a, K, V, marker::Leaf>>,
Handle<NodeRef<marker::DormantMut, K, V, marker::Leaf>, marker::KV>,
) {
if self.node.len() < CAPACITY {
let val_ptr = self.insert_fit(key, val);
(None, val_ptr)
// SAFETY: There is enough space in the node for insertion.
let handle = unsafe { self.insert_fit(key, val) };
(None, handle.dormant())
} else {
let (middle_kv_idx, insertion) = splitpoint(self.idx);
let middle = unsafe { Handle::new_kv(self.node, middle_kv_idx) };
let mut result = middle.split(alloc);
let mut insertion_edge = match insertion {
let insertion_edge = match insertion {
LeftOrRight::Left(insert_idx) => unsafe {
Handle::new_edge(result.left.reborrow_mut(), insert_idx)
},
@ -893,8 +937,10 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
Handle::new_edge(result.right.borrow_mut(), insert_idx)
},
};
let val_ptr = insertion_edge.insert_fit(key, val);
(Some(result), val_ptr)
// SAFETY: We just split the node, so there is enough space for
// insertion.
let handle = unsafe { insertion_edge.insert_fit(key, val).dormant() };
(Some(result), handle)
}
}
}
@ -976,21 +1022,31 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
key: K,
value: V,
alloc: A,
) -> (Option<SplitResult<'a, K, V, marker::LeafOrInternal>>, *mut V) {
let (mut split, val_ptr) = match self.insert(key, value, alloc.clone()) {
(None, val_ptr) => return (None, val_ptr),
(Some(split), val_ptr) => (split.forget_node_type(), val_ptr),
split_root: impl FnOnce(SplitResult<'a, K, V, marker::LeafOrInternal>),
) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
let (mut split, handle) = match self.insert(key, value, alloc.clone()) {
// SAFETY: we have finished splitting and can now re-awaken the
// handle to the inserted element.
(None, handle) => return unsafe { handle.awaken() },
(Some(split), handle) => (split.forget_node_type(), handle),
};
loop {
split = match split.left.ascend() {
Ok(parent) => {
match parent.insert(split.kv.0, split.kv.1, split.right, alloc.clone()) {
None => return (None, val_ptr),
// SAFETY: we have finished splitting and can now re-awaken the
// handle to the inserted element.
None => return unsafe { handle.awaken() },
Some(split) => split.forget_node_type(),
}
}
Err(root) => return (Some(SplitResult { left: root, ..split }), val_ptr),
Err(root) => {
split_root(SplitResult { left: root, ..split });
// SAFETY: we have finished splitting and can now re-awaken the
// handle to the inserted element.
return unsafe { handle.awaken() };
}
};
}
}
@ -1043,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> {
@ -1667,6 +1731,7 @@ pub mod marker {
pub enum Owned {}
pub enum Dying {}
pub enum DormantMut {}
pub struct Immut<'a>(PhantomData<&'a ()>);
pub struct Mut<'a>(PhantomData<&'a mut ()>);
pub struct ValMut<'a>(PhantomData<&'a mut ()>);
@ -1688,6 +1753,7 @@ pub mod marker {
impl<'a> BorrowType for Immut<'a> {}
impl<'a> BorrowType for Mut<'a> {}
impl<'a> BorrowType for ValMut<'a> {}
impl BorrowType for DormantMut {}
pub enum KV {}
pub enum Edge {}