Rewrite the BTreeMap cursor API using gaps

Tracking issue: #107540

Currently, a `Cursor` points to a single element in the tree, and allows
moving to the next or previous element while mutating the tree. However
this was found to be confusing and hard to use.

This PR completely refactors cursors to instead point to a gap between
two elements in the tree. This eliminates the need for a "ghost" element
that exists after the last element and before the first one.
Additionally, `upper_bound` and `lower_bound` now have a much clearer
meaning.

The ability to mutate keys is also factored out into a separate
`CursorMutKey` type which is unsafe to create. This makes the API easier
to use since it avoids duplicated versions of each method with and
without key mutation.

API summary:

```rust
impl<K, V> BTreeMap<K, V> {
    fn lower_bound<Q>(&self, bound: Bound<&Q>) -> Cursor<'_, K, V>
    where
        K: Borrow<Q> + Ord,
        Q: Ord;
    fn lower_bound_mut<Q>(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V>
    where
        K: Borrow<Q> + Ord,
        Q: Ord;
    fn upper_bound<Q>(&self, bound: Bound<&Q>) -> Cursor<'_, K, V>
    where
        K: Borrow<Q> + Ord,
        Q: Ord;
    fn upper_bound_mut<Q>(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V>
    where
        K: Borrow<Q> + Ord,
        Q: Ord;
}

struct Cursor<'a, K: 'a, V: 'a>;

impl<'a, K, V> Cursor<'a, K, V> {
    fn next(&mut self) -> Option<(&'a K, &'a V)>;
    fn prev(&mut self) -> Option<(&'a K, &'a V)>;
    fn peek_next(&self) -> Option<(&'a K, &'a V)>;
    fn peek_prev(&self) -> Option<(&'a K, &'a V)>;
}

struct CursorMut<'a, K: 'a, V: 'a>;

impl<'a, K, V> CursorMut<'a, K, V> {
    fn next(&mut self) -> Option<(&K, &mut V)>;
    fn prev(&mut self) -> Option<(&K, &mut V)>;
    fn peek_next(&mut self) -> Option<(&K, &mut V)>;
    fn peek_prev(&mut self) -> Option<(&K, &mut V)>;

    unsafe fn insert_after_unchecked(&mut self, key: K, value: V);
    unsafe fn insert_before_unchecked(&mut self, key: K, value: V);
    fn insert_after(&mut self, key: K, value: V);
    fn insert_before(&mut self, key: K, value: V);
    fn remove_next(&mut self) -> Option<(K, V)>;
    fn remove_prev(&mut self) -> Option<(K, V)>;

    fn as_cursor(&self) -> Cursor<'_, K, V>;

    unsafe fn with_mutable_key(self) -> CursorMutKey<'a, K, V, A>;
}

struct CursorMutKey<'a, K: 'a, V: 'a>;

impl<'a, K, V> CursorMut<'a, K, V> {
    fn next(&mut self) -> Option<(&mut K, &mut V)>;
    fn prev(&mut self) -> Option<(&mut K, &mut V)>;
    fn peek_next(&mut self) -> Option<(&mut K, &mut V)>;
    fn peek_prev(&mut self) -> Option<(&mut K, &mut V)>;

    unsafe fn insert_after_unchecked(&mut self, key: K, value: V);
    unsafe fn insert_before_unchecked(&mut self, key: K, value: V);
    fn insert_after(&mut self, key: K, value: V);
    fn insert_before(&mut self, key: K, value: V);
    fn remove_next(&mut self) -> Option<(K, V)>;
    fn remove_prev(&mut self) -> Option<(K, V)>;

    fn as_cursor(&self) -> Cursor<'_, K, V>;

    unsafe fn with_mutable_key(self) -> CursorMutKey<'a, K, V, A>;
}
```
This commit is contained in:
Amanieu d'Antras 2023-06-26 14:26:43 +01:00
parent 4f3da903a4
commit 8ee9693177
3 changed files with 539 additions and 394 deletions

File diff suppressed because it is too large Load Diff

View File

@ -2354,48 +2354,95 @@ 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_next(), Some((&1, &'a')));
assert_eq!(cur.peek_prev(), None);
assert_eq!(cur.prev(), None);
assert_eq!(cur.next(), Some((&1, &'a')));
assert_eq!(cur.next(), Some((&2, &'b')));
assert_eq!(cur.peek_next(), Some((&3, &'c')));
assert_eq!(cur.prev(), Some((&2, &'b')));
assert_eq!(cur.peek_prev(), Some((&1, &'a')));
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')));
assert_eq!(cur.peek_prev(), None);
assert_eq!(cur.next(), Some((&1, &'a')));
assert_eq!(cur.prev(), Some((&1, &'a')));
}
#[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));
assert_eq!(cur.peek_next(), Some((&5, &mut 'e')));
assert_eq!(cur.peek_prev(), Some((&3, &mut 'c')));
cur.insert_before(4, 'd');
assert_eq!(cur.key(), Some(&5));
assert_eq!(cur.peek_next(), Some((&5, &mut 'e')));
assert_eq!(cur.peek_prev(), Some((&4, &mut 'd')));
cur.move_next();
assert_eq!(cur.key(), None);
assert_eq!(cur.next(), Some((&5, &mut 'e')));
assert_eq!(cur.peek_next(), None);
assert_eq!(cur.peek_prev(), Some((&5, &mut 'e')));
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')]));
assert_eq!(cur.peek_next(), None);
assert_eq!(cur.peek_prev(), Some((&6, &mut 'f')));
assert_eq!(cur.remove_prev(), Some((6, 'f')));
assert_eq!(cur.remove_prev(), Some((5, 'e')));
assert_eq!(cur.remove_next(), None);
assert_eq!(map, BTreeMap::from([(1, 'a'), (3, 'c'), (4, 'd')]));
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')]));
assert_eq!(cur.peek_next(), None);
assert_eq!(cur.prev(), Some((&4, &mut 'd')));
assert_eq!(cur.peek_next(), Some((&4, &mut 'd')));
assert_eq!(cur.peek_prev(), Some((&3, &mut 'c')));
assert_eq!(cur.remove_next(), Some((4, 'd')));
assert_eq!(map, BTreeMap::from([(1, 'a'), (3, 'c')]));
}
#[test]
fn test_cursor_mut_key() {
let mut map = BTreeMap::from([(1, 'a'), (3, 'c'), (5, 'e')]);
let mut cur = unsafe { map.lower_bound_mut(Bound::Excluded(&3)).with_mutable_key() };
assert_eq!(cur.peek_next(), Some((&mut 5, &mut 'e')));
assert_eq!(cur.peek_prev(), Some((&mut 3, &mut 'c')));
cur.insert_before(4, 'd');
assert_eq!(cur.peek_next(), Some((&mut 5, &mut 'e')));
assert_eq!(cur.peek_prev(), Some((&mut 4, &mut 'd')));
assert_eq!(cur.next(), Some((&mut 5, &mut 'e')));
assert_eq!(cur.peek_next(), None);
assert_eq!(cur.peek_prev(), Some((&mut 5, &mut 'e')));
cur.insert_before(6, 'f');
assert_eq!(cur.peek_next(), None);
assert_eq!(cur.peek_prev(), Some((&mut 6, &mut 'f')));
assert_eq!(cur.remove_prev(), Some((6, 'f')));
assert_eq!(cur.remove_prev(), Some((5, 'e')));
assert_eq!(cur.remove_next(), None);
assert_eq!(map, BTreeMap::from([(1, 'a'), (3, 'c'), (4, 'd')]));
let mut cur = unsafe { map.upper_bound_mut(Bound::Included(&5)).with_mutable_key() };
assert_eq!(cur.peek_next(), None);
assert_eq!(cur.prev(), Some((&mut 4, &mut 'd')));
assert_eq!(cur.peek_next(), Some((&mut 4, &mut 'd')));
assert_eq!(cur.peek_prev(), Some((&mut 3, &mut 'c')));
assert_eq!(cur.remove_next(), Some((4, 'd')));
assert_eq!(map, BTreeMap::from([(1, 'a'), (3, 'c')]));
}
#[test]
fn test_cursor_empty() {
let mut map = BTreeMap::new();
let mut cur = map.lower_bound_mut(Bound::Excluded(&3));
assert_eq!(cur.peek_next(), None);
assert_eq!(cur.peek_prev(), None);
cur.insert_after(0, 0);
assert_eq!(cur.peek_next(), Some((&0, &mut 0)));
assert_eq!(cur.peek_prev(), None);
assert_eq!(map, BTreeMap::from([(0, 0)]));
}
#[should_panic(expected = "key must be ordered above the previous element")]
@ -2414,7 +2461,7 @@ fn test_cursor_mut_insert_before_2() {
cur.insert_before(1, 'd');
}
#[should_panic(expected = "key must be ordered below the current element")]
#[should_panic(expected = "key must be ordered above the previous element")]
#[test]
fn test_cursor_mut_insert_before_3() {
let mut map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]);
@ -2422,7 +2469,7 @@ fn test_cursor_mut_insert_before_3() {
cur.insert_before(2, 'd');
}
#[should_panic(expected = "key must be ordered below the current element")]
#[should_panic(expected = "key must be ordered below the next element")]
#[test]
fn test_cursor_mut_insert_before_4() {
let mut map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]);
@ -2430,7 +2477,7 @@ fn test_cursor_mut_insert_before_4() {
cur.insert_before(3, 'd');
}
#[should_panic(expected = "key must be ordered above the current element")]
#[should_panic(expected = "key must be ordered above the previous element")]
#[test]
fn test_cursor_mut_insert_after_1() {
let mut map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]);
@ -2438,7 +2485,7 @@ fn test_cursor_mut_insert_after_1() {
cur.insert_after(1, 'd');
}
#[should_panic(expected = "key must be ordered above the current element")]
#[should_panic(expected = "key must be ordered above the previous element")]
#[test]
fn test_cursor_mut_insert_after_2() {
let mut map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]);
@ -2467,14 +2514,14 @@ fn cursor_peek_prev_agrees_with_cursor_mut() {
let mut map = BTreeMap::from([(1, 1), (2, 2), (3, 3)]);
let cursor = map.lower_bound(Bound::Excluded(&3));
assert!(cursor.key().is_none());
assert!(cursor.peek_next().is_none());
let prev = cursor.peek_prev();
assert_matches!(prev, Some((&3, _)));
// Shadow names so the two parts of this test match.
let mut cursor = map.lower_bound_mut(Bound::Excluded(&3));
assert!(cursor.key().is_none());
assert!(cursor.peek_next().is_none());
let prev = cursor.peek_prev();
assert_matches!(prev, Some((&3, _)));

View File

@ -648,17 +648,36 @@ impl<K, V, Type> NodeRef<marker::Owned, K, V, Type> {
impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> {
/// Adds a key-value pair to the end of the node, and returns
/// the mutable reference of the inserted value.
pub fn push(&mut self, key: K, val: V) -> &mut V {
/// a handle to the inserted value.
///
/// # Safety
///
/// The returned handle has an unbound lifetime.
pub unsafe fn push_with_handle<'b>(
&mut self,
key: K,
val: V,
) -> Handle<NodeRef<marker::Mut<'b>, K, V, marker::Leaf>, marker::KV> {
let len = self.len_mut();
let idx = usize::from(*len);
assert!(idx < CAPACITY);
*len += 1;
unsafe {
self.key_area_mut(idx).write(key);
self.val_area_mut(idx).write(val)
self.val_area_mut(idx).write(val);
Handle::new_kv(
NodeRef { height: self.height, node: self.node, _marker: PhantomData },
idx,
)
}
}
/// Adds a key-value pair to the end of the node, and returns
/// the mutable reference of the inserted value.
pub fn push(&mut self, key: K, val: V) -> *mut V {
// SAFETY: The unbound handle is no longer accessible.
unsafe { self.push_with_handle(key, val).into_val_mut() }
}
}
impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
@ -1100,10 +1119,10 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>
unsafe { leaf.vals.get_unchecked_mut(self.idx).assume_init_mut() }
}
pub fn into_kv_valmut(self) -> (&'a K, &'a mut V) {
pub fn into_kv_mut(self) -> (&'a mut 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 k = unsafe { leaf.keys.get_unchecked_mut(self.idx).assume_init_mut() };
let v = unsafe { leaf.vals.get_unchecked_mut(self.idx).assume_init_mut() };
(k, v)
}