HashMap get or insert and return a reference from a function

Viewed 41

I have the following code where I would like to return a reference from a HashMap (either get or insert a new value) and only do a single lookup. Is this possible in Rust?

Here is a playground link to play around with.

fn get(key: u32, map: &mut HashMap<u32, String>) -> &mut String {
    match map.entry(key) {
        Entry::Occupied(mut entry) => entry.get_mut(),
        Entry::Vacant(entry) => {
            entry.insert("Hello".into())
        }
    }
}
2 Answers

Do you want this?

fn get(key: u32, map: &mut HashMap<u32, String>) -> &mut String {
    map.entry(key).or_insert_with(|| "Hello".into())
}

The code using Entry::get_mut() doesn't compile because get_mut() returns a lifetime associated with the entry. You need to use Entry::into_mut() instead, which consumes the entry and returns a reference connected to the HashMap:

fn get(key: u32, map: &mut HashMap<u32, String>) -> &mut String {
    match map.entry(key) {
        Entry::Occupied(entry) => entry.into_mut(),
        Entry::Vacant(entry) => entry.insert("Hello".into()),
    }
}

Playground

Or, as the other answer says, you can use or_insert_with().

Related