Skip to content

Commit

Permalink
Unrolled build for rust-lang#133548
Browse files Browse the repository at this point in the history
Rollup merge of rust-lang#133548 - cuviper:btreeset-entry-api, r=Mark-Simulacrum

Add `BTreeSet` entry APIs to match `HashSet`

The following methods are added, along with the corresponding `Entry` implementation.

```rust
impl<T, A: Allocator + Clone> BTreeSet<T, A> {
    pub fn get_or_insert(&mut self, value: T) -> &T
    where
        T: Ord,
    {...}
    pub fn get_or_insert_with<Q: ?Sized, F>(&mut self, value: &Q, f: F) -> &T
    where
        T: Borrow<Q> + Ord,
        Q: Ord,
        F: FnOnce(&Q) -> T,
    {...}

    pub fn entry(&mut self, value: T) -> Entry<'_, T, A>
    where
        T: Ord,
    {...}
}
```

Tracking issue rust-lang#133549
Closes rust-lang/rfcs#1490
  • Loading branch information
rust-timer authored Nov 30, 2024
2 parents e93e096 + 21b1ab1 commit 20f5435
Show file tree
Hide file tree
Showing 4 changed files with 530 additions and 2 deletions.
29 changes: 28 additions & 1 deletion library/alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -308,11 +308,38 @@ impl<K, A: Allocator + Clone> BTreeMap<K, SetValZST, A> {
alloc: (*map.alloc).clone(),
_marker: PhantomData,
}
.insert(SetValZST::default());
.insert(SetValZST);
None
}
}
}

pub(super) fn get_or_insert_with<Q: ?Sized, F>(&mut self, q: &Q, f: F) -> &K
where
K: Borrow<Q> + Ord,
Q: Ord,
F: FnOnce(&Q) -> K,
{
let (map, dormant_map) = DormantMutRef::new(self);
let root_node =
map.root.get_or_insert_with(|| Root::new((*map.alloc).clone())).borrow_mut();
match root_node.search_tree(q) {
Found(handle) => handle.into_kv_mut().0,
GoDown(handle) => {
let key = f(q);
assert!(*key.borrow() == *q, "new value is not equal");
VacantEntry {
key,
handle: Some(handle),
dormant_map,
alloc: (*map.alloc).clone(),
_marker: PhantomData,
}
.insert_entry(SetValZST)
.into_key()
}
}
}
}

/// An iterator over the entries of a `BTreeMap`.
Expand Down
5 changes: 5 additions & 0 deletions library/alloc/src/collections/btree/map/entry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -449,6 +449,11 @@ impl<'a, K: Ord, V, A: Allocator + Clone> OccupiedEntry<'a, K, V, A> {
self.handle.reborrow().into_kv().0
}

/// Converts the entry into a reference to its key.
pub(crate) fn into_key(self) -> &'a K {
self.handle.into_kv_mut().0
}

/// Take ownership of the key and value from the map.
///
/// # Examples
Expand Down
110 changes: 109 additions & 1 deletion library/alloc/src/collections/btree/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,17 @@ use core::iter::{FusedIterator, Peekable};
use core::mem::ManuallyDrop;
use core::ops::{BitAnd, BitOr, BitXor, Bound, RangeBounds, Sub};

use super::map::{BTreeMap, Keys};
use super::map::{self, BTreeMap, Keys};
use super::merge_iter::MergeIterInner;
use super::set_val::SetValZST;
use crate::alloc::{Allocator, Global};
use crate::vec::Vec;

mod entry;

#[unstable(feature = "btree_set_entry", issue = "133549")]
pub use self::entry::{Entry, OccupiedEntry, VacantEntry};

/// An ordered set based on a B-Tree.
///
/// See [`BTreeMap`]'s documentation for a detailed discussion of this collection's performance
Expand Down Expand Up @@ -928,6 +933,109 @@ impl<T, A: Allocator + Clone> BTreeSet<T, A> {
self.map.replace(value)
}

/// Inserts the given `value` into the set if it is not present, then
/// returns a reference to the value in the set.
///
/// # Examples
///
/// ```
/// #![feature(btree_set_entry)]
///
/// use std::collections::BTreeSet;
///
/// let mut set = BTreeSet::from([1, 2, 3]);
/// assert_eq!(set.len(), 3);
/// assert_eq!(set.get_or_insert(2), &2);
/// assert_eq!(set.get_or_insert(100), &100);
/// assert_eq!(set.len(), 4); // 100 was inserted
/// ```
#[inline]
#[unstable(feature = "btree_set_entry", issue = "133549")]
pub fn get_or_insert(&mut self, value: T) -> &T
where
T: Ord,
{
self.map.entry(value).insert_entry(SetValZST).into_key()
}

/// Inserts a value computed from `f` into the set if the given `value` is
/// not present, then returns a reference to the value in the set.
///
/// # Examples
///
/// ```
/// #![feature(btree_set_entry)]
///
/// use std::collections::BTreeSet;
///
/// let mut set: BTreeSet<String> = ["cat", "dog", "horse"]
/// .iter().map(|&pet| pet.to_owned()).collect();
///
/// assert_eq!(set.len(), 3);
/// for &pet in &["cat", "dog", "fish"] {
/// let value = set.get_or_insert_with(pet, str::to_owned);
/// assert_eq!(value, pet);
/// }
/// assert_eq!(set.len(), 4); // a new "fish" was inserted
/// ```
#[inline]
#[unstable(feature = "btree_set_entry", issue = "133549")]
pub fn get_or_insert_with<Q: ?Sized, F>(&mut self, value: &Q, f: F) -> &T
where
T: Borrow<Q> + Ord,
Q: Ord,
F: FnOnce(&Q) -> T,
{
self.map.get_or_insert_with(value, f)
}

/// Gets the given value's corresponding entry in the set for in-place manipulation.
///
/// # Examples
///
/// ```
/// #![feature(btree_set_entry)]
///
/// use std::collections::BTreeSet;
/// use std::collections::btree_set::Entry::*;
///
/// let mut singles = BTreeSet::new();
/// let mut dupes = BTreeSet::new();
///
/// for ch in "a short treatise on fungi".chars() {
/// if let Vacant(dupe_entry) = dupes.entry(ch) {
/// // We haven't already seen a duplicate, so
/// // check if we've at least seen it once.
/// match singles.entry(ch) {
/// Vacant(single_entry) => {
/// // We found a new character for the first time.
/// single_entry.insert()
/// }
/// Occupied(single_entry) => {
/// // We've already seen this once, "move" it to dupes.
/// single_entry.remove();
/// dupe_entry.insert();
/// }
/// }
/// }
/// }
///
/// assert!(!singles.contains(&'t') && dupes.contains(&'t'));
/// assert!(singles.contains(&'u') && !dupes.contains(&'u'));
/// assert!(!singles.contains(&'v') && !dupes.contains(&'v'));
/// ```
#[inline]
#[unstable(feature = "btree_set_entry", issue = "133549")]
pub fn entry(&mut self, value: T) -> Entry<'_, T, A>
where
T: Ord,
{
match self.map.entry(value) {
map::Entry::Occupied(entry) => Entry::Occupied(OccupiedEntry { inner: entry }),
map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry { inner: entry }),
}
}

/// If the set contains an element equal to the value, removes it from the
/// set and drops it. Returns whether such an element was present.
///
Expand Down
Loading

0 comments on commit 20f5435

Please sign in to comment.