Skip to content

Commit

Permalink
Rollup merge of rust-lang#70996 - ChaiTRex:master, r=Amanieu
Browse files Browse the repository at this point in the history
Add or_insert_with_key to Entry of HashMap/BTreeMap

Going along with `or_insert_with`, `or_insert_with_key` provides the `Entry`'s key to the lambda, avoiding the need to either clone the key or the need to reimplement this body of this method from scratch each time.

This is useful when the initial value for a map entry is derived from the key. For example, the introductory Rust book has an example Cacher struct that takes an expensive-to-compute lambda and then can, given an argument to the lambda, produce either the cached result or execute the lambda.

---

I'm fairly new to Rust, so any optimizations, corrections to types, better names, better documentation, or whatever else would be appreciated. I'd like to thank Arnavion on freenode for helping me to implement a very similar method when I found that `or_insert_with_key` was unavailable.

As a somewhat-related note, this implements rust-lang/rfcs#1202 from 2015, so if this pull request is accepted, that should be closed.
  • Loading branch information
Centril authored Apr 11, 2020
2 parents 6bad3ee + db0c39f commit b90d0b3
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 0 deletions.
28 changes: 28 additions & 0 deletions src/liballoc/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2361,6 +2361,34 @@ impl<'a, K: Ord, V> Entry<'a, K, V> {
}
}

#[unstable(feature = "or_insert_with_key", issue = "71024")]
/// Ensures a value is in the entry by inserting, if empty, the result of the default function,
/// which takes the key as its argument, and returns a mutable reference to the value in the
/// entry.
///
/// # Examples
///
/// ```
/// #![feature(or_insert_with_key)]
/// use std::collections::BTreeMap;
///
/// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
///
/// map.entry("poneyland").or_insert_with_key(|key| key.chars().count());
///
/// assert_eq!(map["poneyland"], 9);
/// ```
#[inline]
pub fn or_insert_with_key<F: FnOnce(&K) -> V>(self, default: F) -> &'a mut V {
match self {
Occupied(entry) => entry.into_mut(),
Vacant(entry) => {
let value = default(entry.key());
entry.insert(value)
}
}
}

/// Returns a reference to this entry's key.
///
/// # Examples
Expand Down
28 changes: 28 additions & 0 deletions src/libstd/collections/hash/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1943,6 +1943,34 @@ impl<'a, K, V> Entry<'a, K, V> {
}
}

#[unstable(feature = "or_insert_with_key", issue = "71024")]
/// Ensures a value is in the entry by inserting, if empty, the result of the default function,
/// which takes the key as its argument, and returns a mutable reference to the value in the
/// entry.
///
/// # Examples
///
/// ```
/// #![feature(or_insert_with_key)]
/// use std::collections::HashMap;
///
/// let mut map: HashMap<&str, usize> = HashMap::new();
///
/// map.entry("poneyland").or_insert_with_key(|key| key.chars().count());
///
/// assert_eq!(map["poneyland"], 9);
/// ```
#[inline]
pub fn or_insert_with_key<F: FnOnce(&K) -> V>(self, default: F) -> &'a mut V {
match self {
Occupied(entry) => entry.into_mut(),
Vacant(entry) => {
let value = default(entry.key());
entry.insert(value)
}
}
}

/// Returns a reference to this entry's key.
///
/// # Examples
Expand Down

0 comments on commit b90d0b3

Please sign in to comment.