How to lookup from and insert into a HashMap efficiently?

ghz 7months ago ⋅ 69 views

I'd like to do the following:

  • Lookup a Vec for a certain key, and store it for later use.
  • If it doesn't exist, create an empty Vec for the key, but still keep it in the variable.

How to do this efficiently? Naturally I thought I could use match:

use std::collections::HashMap;

// This code doesn't compile.
let mut map = HashMap::new();
let key = "foo";
let values: &Vec<isize> = match map.get(key) {
    Some(v) => v,
    None => {
        let default: Vec<isize> = Vec::new();
        map.insert(key, default);
        &default
    }
};

When I tried it, it gave me errors like:

error[E0502]: cannot borrow `map` as mutable because it is also borrowed as immutable
  --> src/main.rs:11:13
   |
7  |     let values: &Vec<isize> = match map.get(key) {
   |                                     --- immutable borrow occurs here
...
11 |             map.insert(key, default);
   |             ^^^ mutable borrow occurs here
...
15 | }
   | - immutable borrow ends here

I ended up with doing something like this, but I don't like the fact that it performs the lookup twice (map.contains_key and map.get):

// This code does compile.
let mut map = HashMap::new();
let key = "foo";
if !map.contains_key(key) {
    let default: Vec<isize> = Vec::new();
    map.insert(key, default);
}
let values: &Vec<isize> = match map.get(key) {
    Some(v) => v,
    None => {
        panic!("impossiburu!");
    }
};

Is there a safe way to do this with just one match?

Answers

You can achieve your goal by using the entry API provided by HashMap. The entry method allows you to insert a value if the key is not present and return a mutable reference to the value, or return a mutable reference to the existing value if the key is present. This avoids the need for multiple lookups and ensures safe mutable borrowing.

Here's how you can rewrite your code using the entry API:

use std::collections::HashMap;

let mut map = HashMap::new();
let key = "foo";

let values: &Vec<isize> = match map.entry(key) {
    std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(),
    std::collections::hash_map::Entry::Vacant(entry) => entry.insert(Vec::new()),
};

In this code:

  • map.entry(key) returns an enum variant representing whether the key is occupied or vacant.
  • If the key is occupied, entry.into_mut() returns a mutable reference to the existing value.
  • If the key is vacant, entry.insert(Vec::new()) inserts a new empty vector into the map and returns a mutable reference to it.

This way, you only perform one lookup and achieve your goal safely and efficiently.

#rust