diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 4bdd9639557..b585e2082f1 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; @@ -2521,14 +2522,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 +2539,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 +2558,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 +2579,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 +2597,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. + /// Returns a [`Cursor`] pointing at the last gap 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. + /// Passing [`Bound::Unbounded`] will return a cursor pointing to the end + /// of the map. /// /// # Examples /// @@ -2637,14 +2635,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 +2654,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. + /// Returns a [`CursorMut`] pointing at the last gap 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. + /// Passing [`Bound::Unbounded`] will return a cursor pointing to the end + /// of the map. /// /// # Examples /// @@ -2678,14 +2674,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 +2692,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 +2718,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 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. #[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 +2740,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 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_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,288 +2763,272 @@ 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 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 +/// [`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. + /// 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(&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`. + /// 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. - #[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 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: - /// - /// * 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`. + /// 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)) + let (k, v) = self.inner.peek_next()?; + Some((&*k, v)) } - /// 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)) + let (k, v) = self.inner.peek_prev()?; + Some((&*k, v)) } - /// Returns a read-only cursor pointing to the current element. + /// Returns a read-only cursor pointing to the same location as the + /// `CursorMut`. /// /// 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> { + self.inner.as_cursor() + } + + /// Converts the cursor into a [`CursorMutKey`], which allows mutating + /// the key of elements in the tree. + /// + /// # 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 unsafe fn with_mutable_key(self) -> CursorMutKey<'a, K, V, A> { + self.inner + } +} + +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. + /// + /// 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 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 + } + } + } + + /// Advances the cursor to the previous gap, returning the key and value of + /// the element that it moved over. + /// + /// 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 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 the key and value of the next element without + /// moving the cursor. + /// + /// 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<(&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 the key and value of the previous element + /// without moving the cursor. + /// + /// 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<(&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 same location as the + /// `CursorMutKey`. + /// + /// The lifetime of the returned `Cursor` is bound to that of the + /// `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 { // SAFETY: The tree is immutable while the cursor exists. @@ -3054,11 +3039,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 +3057,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 +3079,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 +3105,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 +3125,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 /// @@ -3154,26 +3143,28 @@ 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) { - if let Some(current) = self.key() { - if &key <= current { - panic!("key must be ordered above the current element"); + pub fn insert_after(&mut self, key: K, value: V) -> Result<(), UnorderedKeyError> { + if let Some((prev, _)) = self.peek_prev() { + if &key <= prev { + 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` 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 /// @@ -3183,36 +3174,36 @@ 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) { - if let Some(current) = self.key() { - if &key >= current { - panic!("key must be ordered below the current element"); - } - } + 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 { + return Err(UnorderedKeyError {}); } } unsafe { self.insert_before_unchecked(key, value); } + Ok(()) } - /// Removes the current element from the `BTreeMap`. + /// Removes the next 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. + /// 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 +3214,19 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { Some(kv) } - /// Removes the current element from the `BTreeMap`. + /// Removes the precending 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. + /// 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 +3238,114 @@ 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) -> Result<(), UnorderedKeyError> { + 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) -> Result<(), UnorderedKeyError> { + 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() + } +} + +/// 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 a1b7cfe6b17..56620cf890d 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2349,112 +2349,151 @@ 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)); - 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((&3, &mut 'c'))); + + cur.insert_before(4, 'd').unwrap(); + 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); - 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.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').unwrap(); + 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').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').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'))); + 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).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 below the current 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 current 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 current 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 current 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] @@ -2462,14 +2501,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 3233a575ecf..78ccb3af66d 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) }