From 8ee96931776c9d2912b379868ae3f7351d67281c Mon Sep 17 00:00:00 2001 From: Amanieu d'Antras Date: Mon, 26 Jun 2023 14:26:43 +0100 Subject: [PATCH 1/3] 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 BTreeMap { fn lower_bound(&self, bound: Bound<&Q>) -> Cursor<'_, K, V> where K: Borrow + Ord, Q: Ord; fn lower_bound_mut(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V> where K: Borrow + Ord, Q: Ord; fn upper_bound(&self, bound: Bound<&Q>) -> Cursor<'_, K, V> where K: Borrow + Ord, Q: Ord; fn upper_bound_mut(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V> where K: Borrow + 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>; } ``` --- library/alloc/src/collections/btree/map.rs | 749 ++++++++++-------- .../alloc/src/collections/btree/map/tests.rs | 115 ++- library/alloc/src/collections/btree/node.rs | 29 +- 3 files changed, 519 insertions(+), 374 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 4bdd9639557d4..b36298cde3675 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -2521,14 +2521,10 @@ impl BTreeMap { self.len() == 0 } - /// Returns a [`Cursor`] pointing at the first element that is above the - /// given bound. + /// Returns a [`Cursor`] pointing to the first gap 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. + /// Passing [`Bound::Unbounded`] will return a cursor pointing to the start + /// of the map. /// /// # Examples /// @@ -2542,14 +2538,16 @@ impl BTreeMap { /// a.insert(1, "a"); /// a.insert(2, "b"); /// a.insert(3, "c"); - /// a.insert(4, "c"); + /// a.insert(4, "d"); /// let cursor = a.lower_bound(Bound::Included(&2)); - /// assert_eq!(cursor.key(), Some(&2)); + /// assert_eq!(cursor.peek_prev(), Some((&1, &"a"))); + /// assert_eq!(cursor.peek_next(), Some((&2, &"b"))); /// let cursor = a.lower_bound(Bound::Excluded(&2)); - /// assert_eq!(cursor.key(), Some(&3)); + /// assert_eq!(cursor.peek_prev(), Some((&2, &"b"))); + /// assert_eq!(cursor.peek_next(), Some((&3, &"c"))); /// ``` #[unstable(feature = "btree_cursors", issue = "107540")] - pub fn lower_bound(&self, bound: Bound<&Q>) -> Cursor<'_, K, V> + pub fn lower_bound(&self, bound: Bound<&Q>) -> Cursor<'_, K, V> where K: Borrow + Ord, Q: Ord, @@ -2559,17 +2557,14 @@ impl BTreeMap { 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() } + Cursor { current: Some(edge), root: self.root.as_ref() } } - /// Returns a [`CursorMut`] pointing at the first element that is above the - /// given bound. + /// Returns a [`CursorMut`] pointing to the first gap 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. + /// Passing [`Bound::Unbounded`] will return a cursor pointing to the start + /// of the map. /// /// # Examples /// @@ -2583,14 +2578,16 @@ impl BTreeMap { /// a.insert(1, "a"); /// a.insert(2, "b"); /// a.insert(3, "c"); - /// a.insert(4, "c"); - /// let cursor = a.lower_bound_mut(Bound::Included(&2)); - /// assert_eq!(cursor.key(), Some(&2)); - /// let cursor = a.lower_bound_mut(Bound::Excluded(&2)); - /// assert_eq!(cursor.key(), Some(&3)); + /// a.insert(4, "d"); + /// let mut cursor = a.lower_bound_mut(Bound::Included(&2)); + /// assert_eq!(cursor.peek_prev(), Some((&1, &mut "a"))); + /// assert_eq!(cursor.peek_next(), Some((&2, &mut "b"))); + /// let mut cursor = a.lower_bound_mut(Bound::Excluded(&2)); + /// assert_eq!(cursor.peek_prev(), Some((&2, &mut "b"))); + /// assert_eq!(cursor.peek_next(), Some((&3, &mut "c"))); /// ``` #[unstable(feature = "btree_cursors", issue = "107540")] - pub fn lower_bound_mut(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V, A> + pub fn lower_bound_mut(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V, A> where K: Borrow + Ord, Q: Ord, @@ -2599,31 +2596,31 @@ impl BTreeMap { let root_node = match root.as_mut() { None => { return CursorMut { - current: None, - root: dormant_root, - length: &mut self.length, - alloc: &mut *self.alloc, + inner: CursorMutKey { + 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, + inner: CursorMutKey { + current: Some(edge), + 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. + /// Returns a [`Cursor`] pointing at the last gap below the given bound. /// - /// Passing [`Bound::Unbounded`] will return a cursor pointing at the last - /// element of the map. + /// Passing [`Bound::Unbounded`] will return a cursor pointing to the end + /// of the map. /// /// # Examples /// @@ -2637,14 +2634,16 @@ impl BTreeMap { /// a.insert(1, "a"); /// a.insert(2, "b"); /// a.insert(3, "c"); - /// a.insert(4, "c"); + /// a.insert(4, "d"); /// let cursor = a.upper_bound(Bound::Included(&3)); - /// assert_eq!(cursor.key(), Some(&3)); + /// assert_eq!(cursor.peek_prev(), Some((&3, &"c"))); + /// assert_eq!(cursor.peek_next(), Some((&4, &"d"))); /// let cursor = a.upper_bound(Bound::Excluded(&3)); - /// assert_eq!(cursor.key(), Some(&2)); + /// assert_eq!(cursor.peek_prev(), Some((&2, &"b"))); + /// assert_eq!(cursor.peek_next(), Some((&3, &"c"))); /// ``` #[unstable(feature = "btree_cursors", issue = "107540")] - pub fn upper_bound(&self, bound: Bound<&Q>) -> Cursor<'_, K, V> + pub fn upper_bound(&self, bound: Bound<&Q>) -> Cursor<'_, K, V> where K: Borrow + Ord, Q: Ord, @@ -2654,17 +2653,13 @@ impl BTreeMap { 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() } + Cursor { current: Some(edge), 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. + /// Returns a [`CursorMut`] pointing at the last gap below the given bound. /// - /// Passing [`Bound::Unbounded`] will return a cursor pointing at the last - /// element of the map. + /// Passing [`Bound::Unbounded`] will return a cursor pointing to the end + /// of the map. /// /// # Examples /// @@ -2678,14 +2673,16 @@ impl BTreeMap { /// a.insert(1, "a"); /// a.insert(2, "b"); /// a.insert(3, "c"); - /// a.insert(4, "c"); - /// let cursor = a.upper_bound_mut(Bound::Included(&3)); - /// assert_eq!(cursor.key(), Some(&3)); - /// let cursor = a.upper_bound_mut(Bound::Excluded(&3)); - /// assert_eq!(cursor.key(), Some(&2)); + /// a.insert(4, "d"); + /// let mut cursor = a.upper_bound_mut(Bound::Included(&3)); + /// assert_eq!(cursor.peek_prev(), Some((&3, &mut "c"))); + /// assert_eq!(cursor.peek_next(), Some((&4, &mut "d"))); + /// let mut cursor = a.upper_bound_mut(Bound::Excluded(&3)); + /// assert_eq!(cursor.peek_prev(), Some((&2, &mut "b"))); + /// assert_eq!(cursor.peek_next(), Some((&3, &mut "c"))); /// ``` #[unstable(feature = "btree_cursors", issue = "107540")] - pub fn upper_bound_mut(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V, A> + pub fn upper_bound_mut(&mut self, bound: Bound<&Q>) -> CursorMut<'_, K, V, A> where K: Borrow + Ord, Q: Ord, @@ -2694,20 +2691,24 @@ impl BTreeMap { let root_node = match root.as_mut() { None => { return CursorMut { - current: None, - root: dormant_root, - length: &mut self.length, - alloc: &mut *self.alloc, + inner: CursorMutKey { + 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, + inner: CursorMutKey { + current: Some(edge), + root: dormant_root, + length: &mut self.length, + alloc: &mut *self.alloc, + }, } } } @@ -2716,14 +2717,14 @@ impl 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. +/// Cursors always point to a gao between two elements in the map, and can +/// operate on the two immediately adjacent elements. /// /// 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, K, V, marker::LeafOrInternal>, marker::KV>>, + // If current is None then it means the tree has not been allocated yet. + current: Option, K, V, marker::Leaf>, marker::Edge>>, root: Option<&'a node::Root>, } @@ -2738,22 +2739,21 @@ impl Clone for Cursor<'_, K, V> { #[unstable(feature = "btree_cursors", issue = "107540")] impl Debug for Cursor<'_, K, V> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_tuple("Cursor").field(&self.key_value()).finish() + f.write_str("Cursor") } } /// 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 +/// safely mutate the map during iteration. This is because the lifetime of its yielded +/// references is tied to its own lifetime, instead of just the underlying map. 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. +/// Cursors always point to a gao between two elements in the map, and can +/// operate on the two immediately adjacent elements. /// -/// A `Cursor` is created with the [`BTreeMap::lower_bound_mut`] and [`BTreeMap::upper_bound_mut`] +/// A `CursorMut` is created with the [`BTreeMap::lower_bound_mut`] and [`BTreeMap::upper_bound_mut`] /// methods. #[unstable(feature = "btree_cursors", issue = "107540")] pub struct CursorMut< @@ -2762,287 +2762,271 @@ pub struct CursorMut< V: 'a, #[unstable(feature = "allocator_api", issue = "32838")] A = Global, > { - current: Option, K, V, marker::LeafOrInternal>, marker::KV>>, + inner: CursorMutKey<'a, K, V, A>, +} + +#[unstable(feature = "btree_cursors", issue = "107540")] +impl Debug for CursorMut<'_, K, V, A> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("CursorMut") + } +} + +/// A cursor over a `BTreeMap` with editing operations, and which allows +/// mutating the key of elements. +/// +/// A `Cursor` is like an iterator, except that it can freely seek back-and-forth, and can +/// safely mutate the map during iteration. This is because the lifetime of its yielded +/// references is tied to its own lifetime, instead of just the underlying map. This means +/// cursors cannot yield multiple elements at once. +/// +/// Cursors always point to a gao between two elements in the map, and can +/// operate on the two immediately adjacent elements. +/// +/// A `CursorMutKey` is created from a [`CursorMut`] with the +/// [`CursorMut::with_mutable_key`] method. +/// +/// # Safety +/// +/// Since this cursor allows mutating keys, 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 struct CursorMutKey< + 'a, + K: 'a, + V: 'a, + #[unstable(feature = "allocator_api", issue = "32838")] A = Global, +> { + // If current is None then it means the tree has not been allocated yet. + current: Option, K, V, marker::Leaf>, marker::Edge>>, root: DormantMutRef<'a, Option>>, length: &'a mut usize, alloc: &'a mut A, } #[unstable(feature = "btree_cursors", issue = "107540")] -impl Debug for CursorMut<'_, K, V, A> { +impl Debug for CursorMutKey<'_, K, V, A> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_tuple("CursorMut").field(&self.key_value()).finish() + f.write_str("CursorMutKey") } } impl<'a, K, V> Cursor<'a, K, V> { - /// Moves the cursor to the next element of the `BTreeMap`. + /// Advances the cursor to the next gap, returning the key and value of the + /// element that it moved over. /// - /// 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. + /// If the cursor is already at the end of the map then `None` is returned + /// and the cursor is not moved. #[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() - }); + pub fn next(&mut self) -> Option<(&'a K, &'a V)> { + let current = self.current.take()?; + match current.next_kv() { + Ok(kv) => { + let result = kv.into_kv(); + self.current = Some(kv.next_leaf_edge()); + Some(result) } - Some(current) => { - self.current = current.next_leaf_edge().next_kv().ok(); + Err(root) => { + self.current = Some(root.last_leaf_edge()); + None } } } - /// Moves the cursor to the previous element of the `BTreeMap`. + /// Advances the cursor to the previous gap, returning the key and value of + /// the element that it moved over. /// - /// 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. + /// If the cursor is already at the start of the map then `None` is returned + /// and the cursor is not moved. #[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() - }); + pub fn prev(&mut self) -> Option<(&'a K, &'a V)> { + let current = self.current.take()?; + match current.next_back_kv() { + Ok(kv) => { + let result = kv.into_kv(); + self.current = Some(kv.next_back_leaf_edge()); + Some(result) } - Some(current) => { - self.current = current.next_back_leaf_edge().next_back_kv().ok(); + Err(root) => { + self.current = Some(root.first_leaf_edge()); + None } } } - /// 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. + /// Returns a reference to the the key and value of the next element without + /// moving the cursor. /// - /// 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`. + /// If the cursor is at the end of the map then `None` is returned #[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()) + self.clone().next() } - /// Returns a reference to the previous element. + /// Returns a reference to the the key and value of the previous element + /// without moving the cursor. /// - /// 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`. + /// If the cursor is at the start of the map then `None` is returned. #[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()) + self.clone().prev() } } impl<'a, K, V, A> CursorMut<'a, K, V, A> { - /// Moves the cursor to the next element of the `BTreeMap`. + /// Advances the cursor to the next gap, returning the key and value of the + /// element that it moved over. /// - /// 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. + /// If the cursor is already at the end of the map then `None` is returned + /// and the cursor is not moved. #[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(); - } - } + pub fn next(&mut self) -> Option<(&K, &mut V)> { + let (k, v) = self.inner.next()?; + Some((&*k, v)) } - /// Moves the cursor to the previous element of the `BTreeMap`. + /// Advances the cursor to the previous gap, returning the key and value of + /// the element that it moved over. /// - /// 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. + /// If the cursor is already at the start of the map then `None` is returned + /// and the cursor is not moved. #[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(); - } - } + pub fn prev(&mut self) -> Option<(&K, &mut V)> { + let (k, v) = self.inner.prev()?; + Some((&*k, v)) } - /// Returns a reference to the key of the element that the cursor is - /// currently pointing to. + /// Returns a reference to the the key and value of the next element without + /// moving the cursor. /// - /// This returns `None` if the cursor is currently pointing to the - /// "ghost" non-element. + /// If the cursor is at the end of the map then `None` is returned #[unstable(feature = "btree_cursors", issue = "107540")] - pub fn key(&self) -> Option<&K> { - self.current.as_ref().map(|current| current.reborrow().into_kv().0) + pub fn peek_next(&mut self) -> Option<(&K, &mut V)> { + let (k, v) = self.inner.peek_next()?; + Some((&*k, v)) } - /// Returns a reference to the value of the element that the cursor is - /// currently pointing to. + /// Returns a reference to the the key and value of the previous element + /// without moving the cursor. /// - /// This returns `None` if the cursor is currently pointing to the - /// "ghost" non-element. + /// If the cursor is at the start of the map then `None` is returned. #[unstable(feature = "btree_cursors", issue = "107540")] - pub fn value(&self) -> Option<&V> { - self.current.as_ref().map(|current| current.reborrow().into_kv().1) + pub fn peek_prev(&mut self) -> Option<(&K, &mut V)> { + let (k, v) = self.inner.peek_prev()?; + Some((&*k, v)) } - /// Returns a reference to the key and value of the element that the cursor - /// is currently pointing to. + /// Returns a read-only cursor pointing to the same location as the + /// `CursorMut`. /// - /// This returns `None` if the cursor is currently pointing to the - /// "ghost" non-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 key_value(&self) -> Option<(&K, &V)> { - self.current.as_ref().map(|current| current.reborrow().into_kv()) + pub fn as_cursor(&self) -> Cursor<'_, K, V> { + self.inner.as_cursor() } - /// Returns a mutable reference to the value of the element that the cursor - /// is currently pointing to. + /// Converts the cursor into a [`CursorMutKey`], which allows mutating + /// the key of elements in the tree. /// - /// This returns `None` if the cursor is currently pointing to the - /// "ghost" non-element. + /// # Safety + /// + /// Since this cursor allows mutating keys, 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 fn value_mut(&mut self) -> Option<&mut V> { - self.current.as_mut().map(|current| current.kv_mut().1) + pub unsafe fn with_mutable_key(self) -> CursorMutKey<'a, K, V, A> { + self.inner } +} - /// Returns a reference to the key and mutable reference to the value of the - /// element that the cursor is currently pointing to. +impl<'a, K, V, A> CursorMutKey<'a, K, V, A> { + /// Advances the cursor to the next gap, returning the key and value of the + /// element that it moved over. /// - /// This returns `None` if the cursor is currently pointing to the - /// "ghost" non-element. + /// If the cursor is already at the end of the map then `None` is returned + /// and the cursor is not moved. #[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) - }) + pub fn next(&mut self) -> Option<(&mut K, &mut V)> { + let current = self.current.take()?; + match current.next_kv() { + Ok(mut kv) => { + // SAFETY: The key/value pointers remain valid even after the + // cursor is moved forward. The lifetimes then prevent any + // further access to the cursor. + let (k, v) = unsafe { kv.reborrow_mut().into_kv_mut() }; + let (k, v) = (k as *mut _, v as *mut _); + self.current = Some(kv.next_leaf_edge()); + Some(unsafe { (&mut *k, &mut *v) }) + } + Err(root) => { + self.current = Some(root.last_leaf_edge()); + None + } + } } - /// Returns a mutable 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. - /// - /// # Safety - /// - /// This can be used to modify the key, but you must ensure that the - /// `BTreeMap` invariants are maintained. Specifically: + /// Advances the cursor to the previous gap, returning the key and value of + /// the element that it moved over. /// - /// * The key must remain unique within the tree. - /// * The key must remain in sorted order with regards to other elements in - /// the tree. + /// If the cursor is already at the start of the map then `None` is returned + /// and the cursor is not moved. #[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) + pub fn prev(&mut self) -> Option<(&mut K, &mut V)> { + let current = self.current.take()?; + match current.next_back_kv() { + Ok(mut kv) => { + // SAFETY: The key/value pointers remain valid even after the + // cursor is moved forward. The lifetimes then prevent any + // further access to the cursor. + let (k, v) = unsafe { kv.reborrow_mut().into_kv_mut() }; + let (k, v) = (k as *mut _, v as *mut _); + self.current = Some(kv.next_back_leaf_edge()); + Some(unsafe { (&mut *k, &mut *v) }) + } + Err(root) => { + self.current = Some(root.first_leaf_edge()); + None + } + } } - /// Returns a reference to the key and value of the next element. + /// Returns a reference to the the key and value of the next element without + /// moving the cursor. /// - /// 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`. + /// If the cursor is at the end of the map then `None` is returned #[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)) + pub fn peek_next(&mut self) -> Option<(&mut K, &mut V)> { + let current = self.current.as_mut()?; + // SAFETY: We're not using this to mutate the tree. + let kv = unsafe { current.reborrow_mut() }.next_kv().ok()?.into_kv_mut(); + Some(kv) } - /// Returns a reference to the key and value of the previous element. + /// Returns a reference to the the key and value of the previous element + /// without moving the cursor. /// - /// 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`. + /// If the cursor is at the start of the map then `None` is returned. #[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() - .last_leaf_edge() - .next_back_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)) + pub fn peek_prev(&mut self) -> Option<(&mut K, &mut V)> { + let current = self.current.as_mut()?; + // SAFETY: We're not using this to mutate the tree. + let kv = unsafe { current.reborrow_mut() }.next_back_kv().ok()?.into_kv_mut(); + Some(kv) } - /// Returns a read-only cursor pointing to the current element. + /// Returns a read-only cursor pointing to the same location as the + /// `CursorMutKey`. /// /// 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`. + /// `CursorMutKey`, which means it cannot outlive the `CursorMutKey` and that the + /// `CursorMutKey` is frozen for the lifetime of the `Cursor`. #[unstable(feature = "btree_cursors", issue = "107540")] pub fn as_cursor(&self) -> Cursor<'_, K, V> { Cursor { @@ -3054,11 +3038,12 @@ impl<'a, K, V, A> CursorMut<'a, K, V, A> { } // 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. +impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> { + /// Inserts a new element into the `BTreeMap` in the gap that the + /// `CursorMutKey` is currently pointing to. /// - /// If the cursor is pointing at the "ghost" non-element then the new element is - /// inserted at the front of the `BTreeMap`. + /// After the insertion the cursor will be pointing at the gap before the + /// newly inserted element. /// /// # Safety /// @@ -3071,20 +3056,19 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { pub unsafe fn insert_after_unchecked(&mut self, key: K, value: V) { let edge = match self.current.take() { None => { + // Tree is empty, allocate a new root. // 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(), - } + let root = unsafe { self.root.reborrow() }; + debug_assert!(root.is_none()); + let mut node = NodeRef::new_leaf(self.alloc.clone()); + // SAFETY: We don't touch the root while the handle is alive. + let handle = unsafe { node.borrow_mut().push_with_handle(key, value) }; + *root = Some(node.forget_type()); + *self.length += 1; + self.current = Some(handle.left_edge()); + return; } - Some(current) => current.next_leaf_edge(), + Some(current) => current, }; let handle = edge.insert_recursing(key, value, self.alloc.clone(), |ins| { @@ -3094,14 +3078,15 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { 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.current = Some(handle.left_edge()); *self.length += 1; } - /// Inserts a new element into the `BTreeMap` before the current one. + /// Inserts a new element into the `BTreeMap` in the gap that the + /// `CursorMutKey` is currently pointing to. /// - /// If the cursor is pointing at the "ghost" non-element then the new element is - /// inserted at the end of the `BTreeMap`. + /// After the insertion the cursor will be pointing at the gap after the + /// newly inserted element. /// /// # Safety /// @@ -3119,15 +3104,17 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { root @ None => { // Tree is empty, allocate a new root. let mut node = NodeRef::new_leaf(self.alloc.clone()); - node.borrow_mut().push(key, value); + // SAFETY: We don't touch the root while the handle is alive. + let handle = unsafe { node.borrow_mut().push_with_handle(key, value) }; *root = Some(node.forget_type()); *self.length += 1; + self.current = Some(handle.right_edge()); return; } Some(root) => root.borrow_mut().last_leaf_edge(), } } - Some(current) => current.next_back_leaf_edge(), + Some(current) => current, }; let handle = edge.insert_recursing(key, value, self.alloc.clone(), |ins| { @@ -3137,14 +3124,15 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { 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.current = Some(handle.right_edge()); *self.length += 1; } - /// Inserts a new element into the `BTreeMap` after the current one. + /// Inserts a new element into the `BTreeMap` in the gap that the + /// `CursorMutKey` is currently pointing to. /// - /// If the cursor is pointing at the "ghost" non-element then the new element is - /// inserted at the front of the `BTreeMap`. + /// After the insertion the cursor will be pointing at the gap before the + /// newly inserted element. /// /// # Panics /// @@ -3155,9 +3143,9 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { /// 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((prev, _)) = self.peek_prev() { + if &key <= prev { + panic!("key must be ordered above the previous element"); } } if let Some((next, _)) = self.peek_next() { @@ -3170,10 +3158,11 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { } } - /// Inserts a new element into the `BTreeMap` before the current one. + /// Inserts a new element into the `BTreeMap` in the gap that the + /// `CursorMutKey` is currently pointing to. /// - /// If the cursor is pointing at the "ghost" non-element then the new element is - /// inserted at the end of the `BTreeMap`. + /// After the insertion the cursor will be pointing at the gap after the + /// newly inserted element. /// /// # Panics /// @@ -3184,35 +3173,34 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { /// 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"); } } + if let Some((next, _)) = self.peek_next() { + if &key >= next { + panic!("key must be ordered below the next 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`. + /// Removes the next element from 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. + /// The element that was removed is returned. The cursor position is + /// unchanged (before the removed element). #[unstable(feature = "btree_cursors", issue = "107540")] - pub fn remove_current(&mut self) -> Option<(K, V)> { + pub fn remove_next(&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(); + let (kv, pos) = current + .next_kv() + .ok()? + .remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone()); + self.current = Some(pos); *self.length -= 1; if emptied_internal_root { // SAFETY: This is safe since current does not point within the now @@ -3223,20 +3211,19 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { 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`. + /// Removes the precending element from 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. + /// The element that was removed is returned. The cursor position is + /// unchanged (after the removed element). #[unstable(feature = "btree_cursors", issue = "107540")] - pub fn remove_current_and_move_back(&mut self) -> Option<(K, V)> { + pub fn remove_prev(&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(); + let (kv, pos) = current + .next_back_kv() + .ok()? + .remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone()); + self.current = Some(pos); *self.length -= 1; if emptied_internal_root { // SAFETY: This is safe since current does not point within the now @@ -3248,5 +3235,97 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { } } +impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { + /// Inserts a new element into the `BTreeMap` in the gap that the + /// `CursorMut` is currently pointing to. + /// + /// After the insertion the cursor will be pointing at the gap before the + /// newly inserted element. + /// + /// # 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) { + unsafe { self.inner.insert_after_unchecked(key, value) } + } + + /// Inserts a new element into the `BTreeMap` in the gap that the + /// `CursorMut` is currently pointing to. + /// + /// After the insertion the cursor will be pointing at the gap after the + /// newly inserted element. + /// + /// # 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) { + unsafe { self.inner.insert_before_unchecked(key, value) } + } + + /// Inserts a new element into the `BTreeMap` in the gap that the + /// `CursorMut` is currently pointing to. + /// + /// After the insertion the cursor will be pointing at the gap before the + /// newly inserted element. + /// + /// # 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) { + self.inner.insert_after(key, value) + } + + /// Inserts a new element into the `BTreeMap` in the gap that the + /// `CursorMut` is currently pointing to. + /// + /// After the insertion the cursor will be pointing at the gap after the + /// newly inserted element. + /// + /// # 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) { + self.inner.insert_before(key, value) + } + + /// Removes the next element from the `BTreeMap`. + /// + /// The element that was removed is returned. The cursor position is + /// unchanged (before the removed element). + #[unstable(feature = "btree_cursors", issue = "107540")] + pub fn remove_next(&mut self) -> Option<(K, V)> { + self.inner.remove_next() + } + + /// Removes the precending element from the `BTreeMap`. + /// + /// The element that was removed is returned. The cursor position is + /// unchanged (after the removed element). + #[unstable(feature = "btree_cursors", issue = "107540")] + pub fn remove_prev(&mut self) -> Option<(K, V)> { + self.inner.remove_prev() + } +} + #[cfg(test)] mod tests; diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 8681cfcd61757..eea00cc5d6726 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -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, _))); diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index 3233a575ecf25..78ccb3af66dbb 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -648,17 +648,36 @@ impl NodeRef { impl<'a, K: 'a, V: 'a> NodeRef, 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, 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, K, V, marker::Internal> { @@ -1100,10 +1119,10 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle, 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) } From 166e3485649f1c4a7e6c4766756a1b0db01f95e1 Mon Sep 17 00:00:00 2001 From: Amanieu d'Antras Date: Thu, 23 Nov 2023 12:37:20 +0000 Subject: [PATCH 2/3] Update library/alloc/src/collections/btree/map.rs Co-authored-by: Joe ST --- library/alloc/src/collections/btree/map.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index b36298cde3675..f0a1561f6508e 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -2717,7 +2717,7 @@ impl BTreeMap { /// /// A `Cursor` is like an iterator, except that it can freely seek back-and-forth. /// -/// Cursors always point to a gao between two elements in the map, and can +/// Cursors always point to a gap between two elements in the map, and can /// operate on the two immediately adjacent elements. /// /// A `Cursor` is created with the [`BTreeMap::lower_bound`] and [`BTreeMap::upper_bound`] methods. From d085f34a2d90a6c32609a1dc8a7ba97d77ce1184 Mon Sep 17 00:00:00 2001 From: Amanieu d'Antras Date: Thu, 23 Nov 2023 16:05:03 +0000 Subject: [PATCH 3/3] Add UnorderedKeyError --- library/alloc/src/collections/btree/map.rs | 40 ++++++++++++++----- .../alloc/src/collections/btree/map/tests.rs | 34 ++++++---------- 2 files changed, 43 insertions(+), 31 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index f0a1561f6508e..b585e2082f137 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1,6 +1,7 @@ use crate::vec::Vec; use core::borrow::Borrow; use core::cmp::Ordering; +use core::error::Error; use core::fmt::{self, Debug}; use core::hash::{Hash, Hasher}; use core::iter::FusedIterator; @@ -2750,7 +2751,7 @@ impl Debug for Cursor<'_, K, V> { /// references is tied to its own lifetime, instead of just the underlying map. This means /// cursors cannot yield multiple elements at once. /// -/// Cursors always point to a gao between two elements in the map, and can +/// Cursors always point to a gap between two elements in the map, and can /// operate on the two immediately adjacent elements. /// /// A `CursorMut` is created with the [`BTreeMap::lower_bound_mut`] and [`BTreeMap::upper_bound_mut`] @@ -2780,7 +2781,7 @@ impl Debug for CursorMut<'_, K, V, A> { /// references is tied to its own lifetime, instead of just the underlying map. This means /// cursors cannot yield multiple elements at once. /// -/// Cursors always point to a gao between two elements in the map, and can +/// Cursors always point to a gap between two elements in the map, and can /// operate on the two immediately adjacent elements. /// /// A `CursorMutKey` is created from a [`CursorMut`] with the @@ -3142,20 +3143,21 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> { /// - 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) { + pub fn insert_after(&mut self, key: K, value: V) -> Result<(), UnorderedKeyError> { if let Some((prev, _)) = self.peek_prev() { if &key <= prev { - panic!("key must be ordered above the previous element"); + return Err(UnorderedKeyError {}); } } if let Some((next, _)) = self.peek_next() { if &key >= next { - panic!("key must be ordered below the next element"); + return Err(UnorderedKeyError {}); } } unsafe { self.insert_after_unchecked(key, value); } + Ok(()) } /// Inserts a new element into the `BTreeMap` in the gap that the @@ -3172,20 +3174,21 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> { /// - 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) { + pub fn insert_before(&mut self, key: K, value: V) -> Result<(), UnorderedKeyError> { if let Some((prev, _)) = self.peek_prev() { if &key <= prev { - panic!("key must be ordered above the previous element"); + return Err(UnorderedKeyError {}); } } if let Some((next, _)) = self.peek_next() { if &key >= next { - panic!("key must be ordered below the next element"); + return Err(UnorderedKeyError {}); } } unsafe { self.insert_before_unchecked(key, value); } + Ok(()) } /// Removes the next element from the `BTreeMap`. @@ -3286,7 +3289,7 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { /// - 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) { + pub fn insert_after(&mut self, key: K, value: V) -> Result<(), UnorderedKeyError> { self.inner.insert_after(key, value) } @@ -3304,7 +3307,7 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { /// - 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) { + pub fn insert_before(&mut self, key: K, value: V) -> Result<(), UnorderedKeyError> { self.inner.insert_before(key, value) } @@ -3327,5 +3330,22 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { } } +/// Error type returned by [`CursorMut::insert_before`] and +/// [`CursorMut::insert_after`] if the key being inserted is not properly +/// ordered with regards to adjacent keys. +#[derive(Clone, PartialEq, Eq, Debug)] +#[unstable(feature = "btree_cursors", issue = "107540")] +pub struct UnorderedKeyError {} + +#[unstable(feature = "btree_cursors", issue = "107540")] +impl fmt::Display for UnorderedKeyError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "key is not properly ordered relative to neighbors") + } +} + +#[unstable(feature = "btree_cursors", issue = "107540")] +impl Error for UnorderedKeyError {} + #[cfg(test)] mod tests; diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index eea00cc5d6726..c1e5df5120478 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2378,14 +2378,14 @@ fn test_cursor_mut() { assert_eq!(cur.peek_next(), Some((&5, &mut 'e'))); assert_eq!(cur.peek_prev(), Some((&3, &mut 'c'))); - cur.insert_before(4, 'd'); + cur.insert_before(4, 'd').unwrap(); assert_eq!(cur.peek_next(), Some((&5, &mut 'e'))); assert_eq!(cur.peek_prev(), Some((&4, &mut 'd'))); 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'); + cur.insert_before(6, 'f').unwrap(); assert_eq!(cur.peek_next(), None); assert_eq!(cur.peek_prev(), Some((&6, &mut 'f'))); assert_eq!(cur.remove_prev(), Some((6, 'f'))); @@ -2409,14 +2409,14 @@ fn test_cursor_mut_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'); + cur.insert_before(4, 'd').unwrap(); 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'); + cur.insert_before(6, 'f').unwrap(); assert_eq!(cur.peek_next(), None); assert_eq!(cur.peek_prev(), Some((&mut 6, &mut 'f'))); assert_eq!(cur.remove_prev(), Some((6, 'f'))); @@ -2439,74 +2439,66 @@ fn test_cursor_empty() { 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); + cur.insert_after(0, 0).unwrap(); 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")] #[test] fn test_cursor_mut_insert_before_1() { let mut map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]); let mut cur = map.upper_bound_mut(Bound::Included(&2)); - cur.insert_before(0, 'd'); + cur.insert_before(0, 'd').unwrap_err(); } -#[should_panic(expected = "key must be ordered above the previous element")] #[test] fn test_cursor_mut_insert_before_2() { let mut map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]); let mut cur = map.upper_bound_mut(Bound::Included(&2)); - cur.insert_before(1, 'd'); + cur.insert_before(1, 'd').unwrap_err(); } -#[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')]); let mut cur = map.upper_bound_mut(Bound::Included(&2)); - cur.insert_before(2, 'd'); + cur.insert_before(2, 'd').unwrap_err(); } -#[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')]); let mut cur = map.upper_bound_mut(Bound::Included(&2)); - cur.insert_before(3, 'd'); + cur.insert_before(3, 'd').unwrap_err(); } -#[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')]); let mut cur = map.upper_bound_mut(Bound::Included(&2)); - cur.insert_after(1, 'd'); + cur.insert_after(1, 'd').unwrap_err(); } -#[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')]); let mut cur = map.upper_bound_mut(Bound::Included(&2)); - cur.insert_after(2, 'd'); + cur.insert_after(2, 'd').unwrap_err(); } -#[should_panic(expected = "key must be ordered below the next element")] #[test] fn test_cursor_mut_insert_after_3() { let mut map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]); let mut cur = map.upper_bound_mut(Bound::Included(&2)); - cur.insert_after(3, 'd'); + cur.insert_after(3, 'd').unwrap_err(); } -#[should_panic(expected = "key must be ordered below the next element")] #[test] fn test_cursor_mut_insert_after_4() { let mut map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]); let mut cur = map.upper_bound_mut(Bound::Included(&2)); - cur.insert_after(4, 'd'); + cur.insert_after(4, 'd').unwrap_err(); } #[test]