CODEWITHSUNDEEP
rusthashmap
rozina
rozina

Reputation: 4232

HashMap get or insert and return a reference from a function

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())
        }
    }
}

Upvotes: 2

Views: 1243

Answers (3)

Tigregalis
Tigregalis

Reputation: 676

While the accepted answer using or_insert_with() ( https://stackoverflow.com/a/73801683/8535397 ) is more idiomatic in most cases, the answer using match map.entry() and into_mut() ( https://stackoverflow.com/a/73805068/8535397 ) was what helped me, and I wanted to write this as an answer because I think there is a specific case where the match is more useful, and I want this to be easily found in the future.

Specifically, using a match allows you to do an early return (or break) within the match statement, useful for example if insertion can fail, and you want to pass the error to the caller, and also allows you to use the question mark ? operator ( https://doc.rust-lang.org/rust-by-example/std/result/question_mark.html ). Inverting the example:

use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::num::ParseIntError;

fn get_parsed(
    number_text: String,
    cache: &mut HashMap<String, u32>
) -> Result<&mut u32, ParseIntError> {
    let parsed = match cache.entry(number_text) {
        Entry::Occupied(entry) => entry.into_mut(),
        Entry::Vacant(entry) => {
            // imagine this is an expensive operation that can fail
            let parsed: u32 = entry.key().parse()?;
            entry.insert(parsed)
        }
    };
    Ok(parsed)
}

Upvotes: 0

user4815162342
user4815162342

Reputation: 155376

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().

Upvotes: 4

phimuemue
phimuemue

Reputation: 36071

Do you want this?

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

Upvotes: 6

Related Questions