diff --git a/src/set.rs b/src/set.rs index ce794ce50a..88e285edf6 100644 --- a/src/set.rs +++ b/src/set.rs @@ -951,7 +951,8 @@ where } /// 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. + /// not present, then returns a reference to the value in the set. The value + /// computed by `f` should have the same hash and compare equivalent to `value`. /// /// # Examples /// @@ -983,6 +984,46 @@ where .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. The value + /// computed by `f` should have the same hash and compare equivalent to `value`. + /// + /// # Examples + /// + /// ``` + /// use hashbrown::HashSet; + /// + /// let mut set: HashSet = ["cat", "dog", "horse"] + /// .iter().map(|&pet| pet.to_owned()).collect(); + /// + /// assert_eq!(set.len(), 3); + /// { + /// let mut pet = String::from("cat"); + /// let value = set.get_or_insert_with_mut(&mut pet, |x| std::mem::take(x)); + /// assert!(!pet.is_empty()); // string contents were not taken because the value was present + /// } + /// { + /// let mut pet = String::from("fish"); + /// let value = set.get_or_insert_with_mut(&mut pet, |x| std::mem::take(x)); + /// assert!(pet.is_empty()); // string contents were taken + /// assert_eq!(set.len(), 4); // a new "fish" was inserted + /// } + /// ``` + #[cfg_attr(feature = "inline-more", inline)] + pub fn get_or_insert_with_mut(&mut self, value: &mut Q, f: F) -> &T + where + Q: Hash + Equivalent, + F: FnOnce(&mut 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 + } + /// Gets the given value's corresponding entry in the set for in-place manipulation. /// /// # Examples