@@ -7,12 +7,17 @@ use core::iter::{FusedIterator, Peekable};
77use core:: mem:: ManuallyDrop ;
88use core:: ops:: { BitAnd , BitOr , BitXor , Bound , RangeBounds , Sub } ;
99
10- use super :: map:: { BTreeMap , Keys } ;
10+ use super :: map:: { self , BTreeMap , Keys } ;
1111use super :: merge_iter:: MergeIterInner ;
1212use super :: set_val:: SetValZST ;
1313use crate :: alloc:: { Allocator , Global } ;
1414use 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