@@ -7,12 +7,17 @@ use core::iter::{FusedIterator, Peekable};
7
7
use core:: mem:: ManuallyDrop ;
8
8
use core:: ops:: { BitAnd , BitOr , BitXor , Bound , RangeBounds , Sub } ;
9
9
10
- use super :: map:: { BTreeMap , Keys } ;
10
+ use super :: map:: { self , BTreeMap , Keys } ;
11
11
use super :: merge_iter:: MergeIterInner ;
12
12
use super :: set_val:: SetValZST ;
13
13
use crate :: alloc:: { Allocator , Global } ;
14
14
use crate :: vec:: Vec ;
15
15
16
+ mod entry;
17
+
18
+ #[ unstable( feature = "btree_set_entry" , issue = "133549" ) ]
19
+ pub use self :: entry:: { Entry , OccupiedEntry , VacantEntry } ;
20
+
16
21
/// An ordered set based on a B-Tree.
17
22
///
18
23
/// 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> {
928
933
self . map . replace ( value)
929
934
}
930
935
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
+
931
1039
/// If the set contains an element equal to the value, removes it from the
932
1040
/// set and drops it. Returns whether such an element was present.
933
1041
///
0 commit comments