diff --git a/src/map.rs b/src/map.rs index 5a5a92c29c..1629f0eddd 100644 --- a/src/map.rs +++ b/src/map.rs @@ -1257,11 +1257,13 @@ pub struct RawEntryBuilderMut<'a, K, V, S> { /// /// This is a lower-level version of [`Entry`]. /// -/// This `enum` is constructed from the [`raw_entry`] method on [`HashMap`]. +/// This `enum` is constructed through the [`raw_entry_mut`] method on [`HashMap`], +/// then calling one of the methods of that [`RawEntryBuilderMut`]. /// /// [`HashMap`]: struct.HashMap.html /// [`Entry`]: enum.Entry.html -/// [`raw_entry`]: struct.HashMap.html#method.raw_entry +/// [`raw_entry_mut`]: struct.HashMap.html#method.raw_entry_mut +/// [`RawEntryBuilderMut`]: struct.RawEntryBuilderMut.html pub enum RawEntryMut<'a, K, V, S> { /// An occupied entry. Occupied(RawOccupiedEntryMut<'a, K, V>), diff --git a/src/set.rs b/src/set.rs index ade77267ac..001bfca55c 100644 --- a/src/set.rs +++ b/src/set.rs @@ -601,6 +601,65 @@ where self.map.get_key_value(value).map(|(k, _)| k) } + /// Inserts the given `value` into the set if it is not present, then + /// returns a reference to the value in the set. + /// + /// # Examples + /// + /// ``` + /// use hashbrown::HashSet; + /// + /// let mut set: HashSet<_> = [1, 2, 3].iter().cloned().collect(); + /// 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] + pub fn get_or_insert(&mut self, value: T) -> &T { + // Although the raw entry gives us `&mut T`, we only return `&T` to be consistent with + // `get`. Key mutation is "raw" because you're not supposed to affect `Eq` or `Hash`. + self.map + .raw_entry_mut() + .from_key(&value) + .or_insert(value, ()) + .0 + } + + /// 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 + /// + /// ``` + /// use hashbrown::HashSet; + /// + /// let mut set: HashSet = ["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] + pub fn get_or_insert_with(&mut self, value: &Q, f: F) -> &T + where + T: Borrow, + Q: Hash + Eq, + F: FnOnce(&Q) -> T, + { + // Although the raw entry gives us `&mut T`, we only return `&T` to be consistent with + // `get`. Key mutation is "raw" because you're not supposed to affect `Eq` or `Hash`. + self.map + .raw_entry_mut() + .from_key(value) + .or_insert_with(|| (f(value), ())) + .0 + } + /// Returns `true` if `self` has no elements in common with `other`. /// This is equivalent to checking for an empty intersection. ///