diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 213924d1d0203..d1ce4e215ed9c 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -308,11 +308,38 @@ impl BTreeMap { alloc: (*map.alloc).clone(), _marker: PhantomData, } - .insert(SetValZST::default()); + .insert(SetValZST); None } } } + + pub(super) fn get_or_insert_with(&mut self, q: &Q, f: F) -> &K + where + K: Borrow + 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`. diff --git a/library/alloc/src/collections/btree/map/entry.rs b/library/alloc/src/collections/btree/map/entry.rs index 0da6af54bc22b..ea8fa363c3805 100644 --- a/library/alloc/src/collections/btree/map/entry.rs +++ b/library/alloc/src/collections/btree/map/entry.rs @@ -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 diff --git a/library/alloc/src/collections/btree/set.rs b/library/alloc/src/collections/btree/set.rs index 8daee6030c270..6f8c3b2d152b8 100644 --- a/library/alloc/src/collections/btree/set.rs +++ b/library/alloc/src/collections/btree/set.rs @@ -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 @@ -928,6 +933,109 @@ impl BTreeSet { 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 = ["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(&mut self, value: &Q, f: F) -> &T + where + T: Borrow + 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. /// diff --git a/library/alloc/src/collections/btree/set/entry.rs b/library/alloc/src/collections/btree/set/entry.rs new file mode 100644 index 0000000000000..a60d22f9ece71 --- /dev/null +++ b/library/alloc/src/collections/btree/set/entry.rs @@ -0,0 +1,388 @@ +use core::fmt::{self, Debug}; + +use Entry::*; + +use super::{SetValZST, map}; +use crate::alloc::{Allocator, Global}; + +/// A view into a single entry in a set, which may either be vacant or occupied. +/// +/// This `enum` is constructed from the [`entry`] method on [`BTreeSet`]. +/// +/// [`BTreeSet`]: super::BTreeSet +/// [`entry`]: super::BTreeSet::entry +/// +/// # Examples +/// +/// ``` +/// #![feature(btree_set_entry)] +/// +/// use std::collections::btree_set::BTreeSet; +/// +/// let mut set = BTreeSet::new(); +/// set.extend(["a", "b", "c"]); +/// assert_eq!(set.len(), 3); +/// +/// // Existing value (insert) +/// let entry = set.entry("a"); +/// let _raw_o = entry.insert(); +/// assert_eq!(set.len(), 3); +/// // Nonexistent value (insert) +/// set.entry("d").insert(); +/// +/// // Existing value (or_insert) +/// set.entry("b").or_insert(); +/// // Nonexistent value (or_insert) +/// set.entry("e").or_insert(); +/// +/// println!("Our BTreeSet: {:?}", set); +/// assert!(set.iter().eq(&["a", "b", "c", "d", "e"])); +/// ``` +#[unstable(feature = "btree_set_entry", issue = "133549")] +pub enum Entry< + 'a, + T, + #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + Clone = Global, +> { + /// An occupied entry. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::btree_set::{Entry, BTreeSet}; + /// + /// let mut set = BTreeSet::from(["a", "b"]); + /// + /// match set.entry("a") { + /// Entry::Vacant(_) => unreachable!(), + /// Entry::Occupied(_) => { } + /// } + /// ``` + #[unstable(feature = "btree_set_entry", issue = "133549")] + Occupied(OccupiedEntry<'a, T, A>), + + /// A vacant entry. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::btree_set::{Entry, BTreeSet}; + /// + /// let mut set = BTreeSet::new(); + /// + /// match set.entry("a") { + /// Entry::Occupied(_) => unreachable!(), + /// Entry::Vacant(_) => { } + /// } + /// ``` + #[unstable(feature = "btree_set_entry", issue = "133549")] + Vacant(VacantEntry<'a, T, A>), +} + +#[unstable(feature = "btree_set_entry", issue = "133549")] +impl Debug for Entry<'_, T, A> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + Vacant(ref v) => f.debug_tuple("Entry").field(v).finish(), + Occupied(ref o) => f.debug_tuple("Entry").field(o).finish(), + } + } +} + +/// A view into an occupied entry in a `BTreeSet`. +/// It is part of the [`Entry`] enum. +/// +/// # Examples +/// +/// ``` +/// #![feature(btree_set_entry)] +/// +/// use std::collections::btree_set::{Entry, BTreeSet}; +/// +/// let mut set = BTreeSet::new(); +/// set.extend(["a", "b", "c"]); +/// +/// let _entry_o = set.entry("a").insert(); +/// assert_eq!(set.len(), 3); +/// +/// // Existing key +/// match set.entry("a") { +/// Entry::Vacant(_) => unreachable!(), +/// Entry::Occupied(view) => { +/// assert_eq!(view.get(), &"a"); +/// } +/// } +/// +/// assert_eq!(set.len(), 3); +/// +/// // Existing key (take) +/// match set.entry("c") { +/// Entry::Vacant(_) => unreachable!(), +/// Entry::Occupied(view) => { +/// assert_eq!(view.remove(), "c"); +/// } +/// } +/// assert_eq!(set.get(&"c"), None); +/// assert_eq!(set.len(), 2); +/// ``` +#[unstable(feature = "btree_set_entry", issue = "133549")] +pub struct OccupiedEntry< + 'a, + T, + #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + Clone = Global, +> { + pub(super) inner: map::OccupiedEntry<'a, T, SetValZST, A>, +} + +#[unstable(feature = "btree_set_entry", issue = "133549")] +impl Debug for OccupiedEntry<'_, T, A> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("OccupiedEntry").field("value", self.get()).finish() + } +} + +/// A view into a vacant entry in a `BTreeSet`. +/// It is part of the [`Entry`] enum. +/// +/// # Examples +/// +/// ``` +/// #![feature(btree_set_entry)] +/// +/// use std::collections::btree_set::{Entry, BTreeSet}; +/// +/// let mut set = BTreeSet::<&str>::new(); +/// +/// let entry_v = match set.entry("a") { +/// Entry::Vacant(view) => view, +/// Entry::Occupied(_) => unreachable!(), +/// }; +/// entry_v.insert(); +/// assert!(set.contains("a") && set.len() == 1); +/// +/// // Nonexistent key (insert) +/// match set.entry("b") { +/// Entry::Vacant(view) => view.insert(), +/// Entry::Occupied(_) => unreachable!(), +/// } +/// assert!(set.contains("b") && set.len() == 2); +/// ``` +#[unstable(feature = "btree_set_entry", issue = "133549")] +pub struct VacantEntry< + 'a, + T, + #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + Clone = Global, +> { + pub(super) inner: map::VacantEntry<'a, T, SetValZST, A>, +} + +#[unstable(feature = "btree_set_entry", issue = "133549")] +impl Debug for VacantEntry<'_, T, A> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("VacantEntry").field(self.get()).finish() + } +} + +impl<'a, T: Ord, A: Allocator + Clone> Entry<'a, T, A> { + /// Sets the value of the entry, and returns an `OccupiedEntry`. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::BTreeSet; + /// + /// let mut set = BTreeSet::new(); + /// let entry = set.entry("horseyland").insert(); + /// + /// assert_eq!(entry.get(), &"horseyland"); + /// ``` + #[inline] + #[unstable(feature = "btree_set_entry", issue = "133549")] + pub fn insert(self) -> OccupiedEntry<'a, T, A> { + match self { + Occupied(entry) => entry, + Vacant(entry) => entry.insert_entry(), + } + } + + /// Ensures a value is in the entry by inserting if it was vacant. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::BTreeSet; + /// + /// let mut set = BTreeSet::new(); + /// + /// // nonexistent key + /// set.entry("poneyland").or_insert(); + /// assert!(set.contains("poneyland")); + /// + /// // existing key + /// set.entry("poneyland").or_insert(); + /// assert!(set.contains("poneyland")); + /// assert_eq!(set.len(), 1); + /// ``` + #[inline] + #[unstable(feature = "btree_set_entry", issue = "133549")] + pub fn or_insert(self) { + if let Vacant(entry) = self { + entry.insert(); + } + } + + /// Returns a reference to this entry's value. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::BTreeSet; + /// + /// let mut set = BTreeSet::new(); + /// set.entry("poneyland").or_insert(); + /// + /// // existing key + /// assert_eq!(set.entry("poneyland").get(), &"poneyland"); + /// // nonexistent key + /// assert_eq!(set.entry("horseland").get(), &"horseland"); + /// ``` + #[inline] + #[unstable(feature = "btree_set_entry", issue = "133549")] + pub fn get(&self) -> &T { + match *self { + Occupied(ref entry) => entry.get(), + Vacant(ref entry) => entry.get(), + } + } +} + +impl<'a, T: Ord, A: Allocator + Clone> OccupiedEntry<'a, T, A> { + /// Gets a reference to the value in the entry. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::btree_set::{Entry, BTreeSet}; + /// + /// let mut set = BTreeSet::new(); + /// set.entry("poneyland").or_insert(); + /// + /// match set.entry("poneyland") { + /// Entry::Vacant(_) => panic!(), + /// Entry::Occupied(entry) => assert_eq!(entry.get(), &"poneyland"), + /// } + /// ``` + #[inline] + #[unstable(feature = "btree_set_entry", issue = "133549")] + pub fn get(&self) -> &T { + self.inner.key() + } + + /// Takes the value out of the entry, and returns it. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::BTreeSet; + /// use std::collections::btree_set::Entry; + /// + /// let mut set = BTreeSet::new(); + /// set.entry("poneyland").or_insert(); + /// + /// if let Entry::Occupied(o) = set.entry("poneyland") { + /// assert_eq!(o.remove(), "poneyland"); + /// } + /// + /// assert_eq!(set.contains("poneyland"), false); + /// ``` + #[inline] + #[unstable(feature = "btree_set_entry", issue = "133549")] + pub fn remove(self) -> T { + self.inner.remove_entry().0 + } +} + +impl<'a, T: Ord, A: Allocator + Clone> VacantEntry<'a, T, A> { + /// Gets a reference to the value that would be used when inserting + /// through the `VacantEntry`. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::BTreeSet; + /// + /// let mut set = BTreeSet::new(); + /// assert_eq!(set.entry("poneyland").get(), &"poneyland"); + /// ``` + #[inline] + #[unstable(feature = "btree_set_entry", issue = "133549")] + pub fn get(&self) -> &T { + self.inner.key() + } + + /// Take ownership of the value. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::btree_set::{Entry, BTreeSet}; + /// + /// let mut set = BTreeSet::new(); + /// + /// match set.entry("poneyland") { + /// Entry::Occupied(_) => panic!(), + /// Entry::Vacant(v) => assert_eq!(v.into_value(), "poneyland"), + /// } + /// ``` + #[inline] + #[unstable(feature = "btree_set_entry", issue = "133549")] + pub fn into_value(self) -> T { + self.inner.into_key() + } + + /// Sets the value of the entry with the VacantEntry's value. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_set_entry)] + /// + /// use std::collections::BTreeSet; + /// use std::collections::btree_set::Entry; + /// + /// let mut set = BTreeSet::new(); + /// + /// if let Entry::Vacant(o) = set.entry("poneyland") { + /// o.insert(); + /// } + /// assert!(set.contains("poneyland")); + /// ``` + #[inline] + #[unstable(feature = "btree_set_entry", issue = "133549")] + pub fn insert(self) { + self.inner.insert(SetValZST); + } + + #[inline] + fn insert_entry(self) -> OccupiedEntry<'a, T, A> { + OccupiedEntry { inner: self.inner.insert_entry(SetValZST) } + } +}