187

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?

Shepmaster
  • 388,571
  • 95
  • 1,107
  • 1,366
Yusuke Shinyama
  • 1,873
  • 2
  • 11
  • 5

2 Answers2

224

The entry API is designed for this. In manual form, it might look like

let values = match map.entry(key) {
    Entry::Occupied(o) => o.into_mut(),
    Entry::Vacant(v) => v.insert(default),
};

One can use the briefer form via Entry::or_insert_with:

let values = map.entry(key).or_insert_with(|| default);

If default is already computed, or if it's OK/cheap to compute even when it isn't inserted, you can use Entry::or_insert:

let values = map.entry(key).or_insert(default);

If the HashMap's value implements Default, you can use Entry::or_default, although you may need to provide some type hints:

let values = map.entry(key).or_default();
Shepmaster
  • 388,571
  • 95
  • 1,107
  • 1,366
huon
  • 94,605
  • 21
  • 231
  • 225
  • Thanks for a swift anser! Now I've learned I should look into a little deep down to the documents. – Yusuke Shinyama Feb 14 '15 at 05:30
  • 46
    the issue with entry() is, that you always have to clone the key, is there a way to avoid this? – Pascalius Dec 16 '17 at 20:04
  • @Pascalius you could make your key type `&T` (if the keys outlive the map, e.g. static strings) or `Rc` instead of `T` -- but it's not pretty in either case – kbolino Jul 03 '18 at 19:42
  • 1
    @Pascalius: you can use `v.key()` in the expression for `default`, and then it will get a reference to the key as it exists in the hashmap, so you may avoid a clone this way – Chris Beck Feb 07 '19 at 23:28
  • @Pascalius Note that since 1.50 there's [or_insert_with_key](https://doc.rust-lang.org/stable/std/collections/hash_map/enum.Entry.html#method.or_insert_with_key) – Bogdan Mircea Aug 16 '22 at 13:56
3

I've used huon's answer and implemented it as a trait:

use std::collections::HashMap;
use std::hash::Hash;

pub trait InsertOrGet<K: Eq + Hash, V: Default> {
    fn insert_or_get(&mut self, item: K) -> &mut V;
}

impl<K: Eq + Hash, V: Default> InsertOrGet<K, V> for HashMap<K, V> {
    fn insert_or_get(&mut self, item: K) -> &mut V {
        return match self.entry(item) {
            std::collections::hash_map::Entry::Occupied(o) => o.into_mut(),
            std::collections::hash_map::Entry::Vacant(v) => v.insert(V::default()),
        };
    }
}

Then I can do:

use crate::utils::hashmap::InsertOrGet;

let new_or_existing_value: &mut ValueType = my_map.insert_or_get(my_key.clone());
Shepmaster
  • 388,571
  • 95
  • 1,107
  • 1,366
Daniel
  • 678
  • 1
  • 5
  • 20
  • 6
    Why don't you use `my_map.entry(my_key).or_default();` instead? You dont need this trait. – Nawaz Jul 24 '22 at 18:34