Skip to content

Commit 9866fb3

Browse files
borsgitbot
authored and
gitbot
committed
Auto merge of rust-lang#133659 - jieyouxu:rollup-576gh4p, r=jieyouxu
Rollup of 6 pull requests Successful merges: - rust-lang#131551 (Support input/output in vector registers of PowerPC inline assembly) - rust-lang#132515 (Fix and undeprecate home_dir()) - rust-lang#132721 (CI: split x86_64-mingw job) - rust-lang#133106 (changes old intrinsic declaration to new declaration) - rust-lang#133496 (thread::available_parallelism for wasm32-wasip1-threads) - rust-lang#133548 (Add `BTreeSet` entry APIs to match `HashSet`) r? `@ghost` `@rustbot` modify labels: rollup
2 parents ebfc2b0 + f4252c5 commit 9866fb3

File tree

9 files changed

+1506
-618
lines changed

9 files changed

+1506
-618
lines changed

alloc/src/collections/btree/map.rs

+28-1
Original file line numberDiff line numberDiff line change
@@ -308,11 +308,38 @@ impl<K, A: Allocator + Clone> BTreeMap<K, SetValZST, A> {
308308
alloc: (*map.alloc).clone(),
309309
_marker: PhantomData,
310310
}
311-
.insert(SetValZST::default());
311+
.insert(SetValZST);
312312
None
313313
}
314314
}
315315
}
316+
317+
pub(super) fn get_or_insert_with<Q: ?Sized, F>(&mut self, q: &Q, f: F) -> &K
318+
where
319+
K: Borrow<Q> + Ord,
320+
Q: Ord,
321+
F: FnOnce(&Q) -> K,
322+
{
323+
let (map, dormant_map) = DormantMutRef::new(self);
324+
let root_node =
325+
map.root.get_or_insert_with(|| Root::new((*map.alloc).clone())).borrow_mut();
326+
match root_node.search_tree(q) {
327+
Found(handle) => handle.into_kv_mut().0,
328+
GoDown(handle) => {
329+
let key = f(q);
330+
assert!(*key.borrow() == *q, "new value is not equal");
331+
VacantEntry {
332+
key,
333+
handle: Some(handle),
334+
dormant_map,
335+
alloc: (*map.alloc).clone(),
336+
_marker: PhantomData,
337+
}
338+
.insert_entry(SetValZST)
339+
.into_key()
340+
}
341+
}
342+
}
316343
}
317344

318345
/// An iterator over the entries of a `BTreeMap`.

alloc/src/collections/btree/map/entry.rs

+5
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,11 @@ impl<'a, K: Ord, V, A: Allocator + Clone> OccupiedEntry<'a, K, V, A> {
449449
self.handle.reborrow().into_kv().0
450450
}
451451

452+
/// Converts the entry into a reference to its key.
453+
pub(crate) fn into_key(self) -> &'a K {
454+
self.handle.into_kv_mut().0
455+
}
456+
452457
/// Take ownership of the key and value from the map.
453458
///
454459
/// # Examples

alloc/src/collections/btree/set.rs

+109-1
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,17 @@ use core::iter::{FusedIterator, Peekable};
77
use core::mem::ManuallyDrop;
88
use core::ops::{BitAnd, BitOr, BitXor, Bound, RangeBounds, Sub};
99

10-
use super::map::{BTreeMap, Keys};
10+
use super::map::{self, BTreeMap, Keys};
1111
use super::merge_iter::MergeIterInner;
1212
use super::set_val::SetValZST;
1313
use crate::alloc::{Allocator, Global};
1414
use crate::vec::Vec;
1515

16+
mod entry;
17+
18+
#[unstable(feature = "btree_set_entry", issue = "133549")]
19+
pub use self::entry::{Entry, OccupiedEntry, VacantEntry};
20+
1621
/// An ordered set based on a B-Tree.
1722
///
1823
/// See [`BTreeMap`]'s documentation for a detailed discussion of this collection's performance
@@ -928,6 +933,109 @@ impl<T, A: Allocator + Clone> BTreeSet<T, A> {
928933
self.map.replace(value)
929934
}
930935

936+
/// Inserts the given `value` into the set if it is not present, then
937+
/// returns a reference to the value in the set.
938+
///
939+
/// # Examples
940+
///
941+
/// ```
942+
/// #![feature(btree_set_entry)]
943+
///
944+
/// use std::collections::BTreeSet;
945+
///
946+
/// let mut set = BTreeSet::from([1, 2, 3]);
947+
/// assert_eq!(set.len(), 3);
948+
/// assert_eq!(set.get_or_insert(2), &2);
949+
/// assert_eq!(set.get_or_insert(100), &100);
950+
/// assert_eq!(set.len(), 4); // 100 was inserted
951+
/// ```
952+
#[inline]
953+
#[unstable(feature = "btree_set_entry", issue = "133549")]
954+
pub fn get_or_insert(&mut self, value: T) -> &T
955+
where
956+
T: Ord,
957+
{
958+
self.map.entry(value).insert_entry(SetValZST).into_key()
959+
}
960+
961+
/// Inserts a value computed from `f` into the set if the given `value` is
962+
/// not present, then returns a reference to the value in the set.
963+
///
964+
/// # Examples
965+
///
966+
/// ```
967+
/// #![feature(btree_set_entry)]
968+
///
969+
/// use std::collections::BTreeSet;
970+
///
971+
/// let mut set: BTreeSet<String> = ["cat", "dog", "horse"]
972+
/// .iter().map(|&pet| pet.to_owned()).collect();
973+
///
974+
/// assert_eq!(set.len(), 3);
975+
/// for &pet in &["cat", "dog", "fish"] {
976+
/// let value = set.get_or_insert_with(pet, str::to_owned);
977+
/// assert_eq!(value, pet);
978+
/// }
979+
/// assert_eq!(set.len(), 4); // a new "fish" was inserted
980+
/// ```
981+
#[inline]
982+
#[unstable(feature = "btree_set_entry", issue = "133549")]
983+
pub fn get_or_insert_with<Q: ?Sized, F>(&mut self, value: &Q, f: F) -> &T
984+
where
985+
T: Borrow<Q> + Ord,
986+
Q: Ord,
987+
F: FnOnce(&Q) -> T,
988+
{
989+
self.map.get_or_insert_with(value, f)
990+
}
991+
992+
/// Gets the given value's corresponding entry in the set for in-place manipulation.
993+
///
994+
/// # Examples
995+
///
996+
/// ```
997+
/// #![feature(btree_set_entry)]
998+
///
999+
/// use std::collections::BTreeSet;
1000+
/// use std::collections::btree_set::Entry::*;
1001+
///
1002+
/// let mut singles = BTreeSet::new();
1003+
/// let mut dupes = BTreeSet::new();
1004+
///
1005+
/// for ch in "a short treatise on fungi".chars() {
1006+
/// if let Vacant(dupe_entry) = dupes.entry(ch) {
1007+
/// // We haven't already seen a duplicate, so
1008+
/// // check if we've at least seen it once.
1009+
/// match singles.entry(ch) {
1010+
/// Vacant(single_entry) => {
1011+
/// // We found a new character for the first time.
1012+
/// single_entry.insert()
1013+
/// }
1014+
/// Occupied(single_entry) => {
1015+
/// // We've already seen this once, "move" it to dupes.
1016+
/// single_entry.remove();
1017+
/// dupe_entry.insert();
1018+
/// }
1019+
/// }
1020+
/// }
1021+
/// }
1022+
///
1023+
/// assert!(!singles.contains(&'t') && dupes.contains(&'t'));
1024+
/// assert!(singles.contains(&'u') && !dupes.contains(&'u'));
1025+
/// assert!(!singles.contains(&'v') && !dupes.contains(&'v'));
1026+
/// ```
1027+
#[inline]
1028+
#[unstable(feature = "btree_set_entry", issue = "133549")]
1029+
pub fn entry(&mut self, value: T) -> Entry<'_, T, A>
1030+
where
1031+
T: Ord,
1032+
{
1033+
match self.map.entry(value) {
1034+
map::Entry::Occupied(entry) => Entry::Occupied(OccupiedEntry { inner: entry }),
1035+
map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry { inner: entry }),
1036+
}
1037+
}
1038+
9311039
/// If the set contains an element equal to the value, removes it from the
9321040
/// set and drops it. Returns whether such an element was present.
9331041
///

0 commit comments

Comments
 (0)