HashMap borrow issue when trying to implement find or insert

Question

I tried to implement own analogue of find_or_insert method that looks like this:

use std::collections::HashMap;

pub struct SomeManager {
    next: i32,
    types: HashMap<i32, i32>,
}

impl SomeManager {
    pub fn get_type<'a>(&'a mut self, k: i32) -> &'a i32 {
        match self.types.get(&k) {
            Some(ref x) => return *x,
            None => {
                self.types.insert(k, self.next);
                self.next += 1;
                return self.types.get(&k).unwrap();
            }
        }
    }
}

fn main() {}

Error:

error[E0502]: cannot borrow `self.types` as mutable because it is also borrowed as immutable
  --> src/main.rs:13:17
   |
10 |         match self.types.get(&k) {
   |               ---------- immutable borrow occurs here
...
13 |                 self.types.insert(k, self.next);
   |                 ^^^^^^^^^^ mutable borrow occurs here
...
18 |     }
   |     - immutable borrow ends here

I know that there are some standard methods that implement this functionality, but I want this method to be as light as possible - it will be called very-very often and almost all of the time the values will already exist.

As I understand it, when we call self.types.get we borrow it to scope of match statement, so we can't call self.types.insert here. I have tried to move methods from None branch out of match statement, but it also fails.

The only working solution that I found requires invoking get twice:

pub fn get_type<'a>(&'a mut self, k: i32) -> &'a i32 {
    let is_none = match self.types.get(&k) {
        Some(ref x) => false,
        None => true,
    };
    if is_none {
        self.types.insert(k, self.next);
        self.next += 1;
    }
    self.types.get(&k).unwrap()
}

How can I work around such situations?

Answer 1

There are a handful of methods on HashMap to achieve these sorts of complex cases. Most notably, for your case, HashMap::entry and Entry::or_insert_with:

pub fn get_type<'a>(&'a mut self, k: i32) -> &'a i32 {
    self.types.entry(k).or_insert_with(|| {
        let value = self.next;
        self.next += 1;
        value
    })
}

In your case, however, there’s the borrow of self inside, so that won’t do.

We thus shift the borrow of self.next outside of the closure so the compiler can reason about it as disjoint from self.types. Problem solved with only one lookup, as it should be.

pub fn get_type<'a>(&'a mut self, k: i32) -> &'a i32 {
    let next = &mut self.next;

    self.types.entry(k).or_insert_with(|| {
        let value = *next;
        *next += 1;
        value
    })
}

Answer 2

Note that in your first case you're doing one lookup when the key exists in the map and three when it does not exist. Your last attempt does two lookups in either case. This is somewhat prettified version of the latter:

pub fn get_type<'a>(&'a mut self, k: i32) -> &'a i32 {
    let contains = self.types.contains_key(&k);
    if !contains {
        self.types.insert(k, self.next);
        self.next += 1;
    }
    self.types.get(&k).unwrap()
}

I don't think it is possible to avoid the second lookup without some support from the map's implementation because of borrowing restrictions.

In any case, using the solution by Chris Morgan is superior to the above one (for example, it may be more efficient and in fact require less lookups), so I suggest sticking with it.